250 lines
10 KiB
Python
250 lines
10 KiB
Python
# Copyright (c) 2025 Patrick Motsch
|
|
# All rights reserved.
|
|
import logging
|
|
import asyncio
|
|
import uuid
|
|
import base64
|
|
from typing import Dict, Any, List, Union, Tuple, Optional
|
|
from dataclasses import dataclass
|
|
import time
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
from modules.aichat.aicore.aicoreModelRegistry import modelRegistry
|
|
from modules.aichat.aicore.aicoreModelSelector import modelSelector
|
|
from modules.datamodels.datamodelAi import (
|
|
AiModel,
|
|
AiCallOptions,
|
|
AiCallRequest,
|
|
AiCallResponse,
|
|
OperationTypeEnum,
|
|
AiModelCall,
|
|
AiModelResponse,
|
|
)
|
|
from modules.datamodels.datamodelExtraction import ContentPart, MergeStrategy
|
|
|
|
|
|
# Dynamic model registry - models are now loaded from connectors via aicore system
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class AiObjects:
|
|
"""Centralized AI interface: dynamically discovers and uses AI models. Includes web functionality."""
|
|
|
|
def __post_init__(self) -> None:
|
|
# Auto-discover and register all available connectors
|
|
self._discoverAndRegisterConnectors()
|
|
|
|
def _discoverAndRegisterConnectors(self):
|
|
"""Auto-discover and register all available AI connectors."""
|
|
logger.info("Auto-discovering AI connectors...")
|
|
|
|
# Use the model registry's built-in discovery mechanism
|
|
discoveredConnectors = modelRegistry.discoverConnectors()
|
|
|
|
# Register each discovered connector
|
|
for connector in discoveredConnectors:
|
|
modelRegistry.registerConnector(connector)
|
|
logger.info(f"Registered connector: {connector.getConnectorType()}")
|
|
|
|
logger.info(f"Total connectors registered: {len(discoveredConnectors)}")
|
|
logger.info("All AI connectors registered with dynamic model registry")
|
|
|
|
@classmethod
|
|
async def create(cls) -> "AiObjects":
|
|
"""Create AiObjects instance with auto-discovered connectors."""
|
|
# No need to manually create connectors - they're auto-discovered
|
|
return cls()
|
|
|
|
def _selectModel(self, prompt: str, context: str, options: AiCallOptions) -> str:
|
|
"""Select the best model using dynamic model selection system. Returns displayName (unique identifier)."""
|
|
# Get available models from the dynamic registry
|
|
availableModels = modelRegistry.getAvailableModels()
|
|
|
|
if not availableModels:
|
|
logger.error("No models available in the registry")
|
|
raise ValueError("No AI models available")
|
|
|
|
# Use the dynamic model selector
|
|
selectedModel = modelSelector.selectModel(prompt, context, options, availableModels)
|
|
|
|
if not selectedModel:
|
|
logger.error("No suitable model found for the given criteria")
|
|
raise ValueError("No suitable AI model found")
|
|
|
|
logger.info(f"Selected model: {selectedModel.name} ({selectedModel.displayName})")
|
|
return selectedModel.displayName
|
|
|
|
|
|
# AI for Extraction, Processing, Generation
|
|
async def callWithTextContext(self, request: AiCallRequest) -> AiCallResponse:
|
|
"""Call AI model for traditional text/context calls with fallback mechanism."""
|
|
prompt = request.prompt
|
|
context = request.context or ""
|
|
options = request.options
|
|
|
|
# Input bytes will be calculated inside _callWithModel
|
|
|
|
# Generation parameters are handled inside _callWithModel
|
|
|
|
# Get failover models for this operation type
|
|
availableModels = modelRegistry.getAvailableModels()
|
|
failoverModelList = modelSelector.getFailoverModelList(prompt, context, options, availableModels)
|
|
|
|
if not failoverModelList:
|
|
errorMsg = f"No suitable models found for operation {options.operationType}"
|
|
logger.error(errorMsg)
|
|
return AiCallResponse(
|
|
content=errorMsg,
|
|
modelName="error",
|
|
priceUsd=0.0,
|
|
processingTime=0.0,
|
|
bytesSent=0,
|
|
bytesReceived=0,
|
|
errorCount=1
|
|
)
|
|
|
|
# Try each model in failover sequence
|
|
lastError = None
|
|
for attempt, model in enumerate(failoverModelList):
|
|
try:
|
|
logger.info(f"Attempting AI call with model: {model.name} (attempt {attempt + 1}/{len(failoverModelList)})")
|
|
|
|
# Call the model directly - no truncation or compression here
|
|
response = await self._callWithModel(model, prompt, context, options)
|
|
|
|
logger.info(f"✅ AI call successful with model: {model.name}")
|
|
return response
|
|
|
|
except Exception as e:
|
|
lastError = e
|
|
logger.warning(f"❌ AI call failed with model {model.name}: {str(e)}")
|
|
|
|
# If this is not the last model, try the next one
|
|
if attempt < len(failoverModelList) - 1:
|
|
logger.info(f"🔄 Trying next failover model...")
|
|
continue
|
|
else:
|
|
# All models failed
|
|
logger.error(f"💥 All {len(failoverModelList)} models failed for operation {options.operationType}")
|
|
break
|
|
|
|
# All failover attempts failed - return error response
|
|
errorMsg = f"All AI models failed for operation {options.operationType}. Last error: {str(lastError)}"
|
|
logger.error(errorMsg)
|
|
return AiCallResponse(
|
|
content=errorMsg,
|
|
modelName="error",
|
|
priceUsd=0.0,
|
|
processingTime=0.0,
|
|
bytesSent=0,
|
|
bytesReceived=0,
|
|
errorCount=1
|
|
)
|
|
|
|
def _createErrorResponse(self, errorMsg: str, inputBytes: int, outputBytes: int) -> AiCallResponse:
|
|
"""Create an error response."""
|
|
return AiCallResponse(
|
|
content=errorMsg,
|
|
modelName="error",
|
|
priceUsd=0.0,
|
|
processingTime=0.0,
|
|
bytesSent=inputBytes,
|
|
bytesReceived=outputBytes,
|
|
errorCount=1
|
|
)
|
|
|
|
async def _callWithModel(self, model: AiModel, prompt: str, context: str, options: AiCallOptions = None) -> AiCallResponse:
|
|
"""Call a specific model and return the response."""
|
|
# Calculate input bytes from prompt and context
|
|
inputBytes = len((prompt + context).encode('utf-8'))
|
|
|
|
# Replace <TOKEN_LIMIT> placeholder with model's maxTokens value
|
|
if "<TOKEN_LIMIT>" in prompt:
|
|
if model.maxTokens > 0:
|
|
tokenLimit = str(model.maxTokens)
|
|
modelPrompt = prompt.replace("<TOKEN_LIMIT>", tokenLimit)
|
|
logger.debug(f"Replaced <TOKEN_LIMIT> with {tokenLimit} for model {model.name}")
|
|
else:
|
|
raise ValueError(f"Model {model.name} has invalid maxTokens ({model.maxTokens}). Cannot set token limit.")
|
|
else:
|
|
modelPrompt = prompt
|
|
|
|
# Update messages array with replaced content
|
|
messages = []
|
|
if context:
|
|
messages.append({"role": "system", "content": f"Context from documents:\n{context}"})
|
|
messages.append({"role": "user", "content": modelPrompt})
|
|
|
|
# Start timing
|
|
startTime = time.time()
|
|
|
|
# Call the model's function directly - completely generic
|
|
if model.functionCall:
|
|
# Create standardized call object
|
|
modelCall = AiModelCall(
|
|
messages=messages,
|
|
model=model,
|
|
options=options or {}
|
|
)
|
|
|
|
# Log before calling model
|
|
contextSize = len(context.encode('utf-8')) if context else 0
|
|
promptSize = len(modelPrompt.encode('utf-8')) if modelPrompt else 0
|
|
totalInputSize = contextSize + promptSize
|
|
logger.debug(f"Calling model {model.name} with {len(messages)} messages, context size: {contextSize} bytes, prompt size: {promptSize} bytes, total input: {totalInputSize} bytes")
|
|
|
|
# Call the model with standardized interface
|
|
modelResponse = await model.functionCall(modelCall)
|
|
|
|
# Log after successful call
|
|
logger.debug(f"Model {model.name} returned successfully")
|
|
|
|
# Extract content from standardized response
|
|
if not modelResponse.success:
|
|
raise ValueError(f"Model call failed: {modelResponse.error}")
|
|
content = modelResponse.content
|
|
else:
|
|
raise ValueError(f"Model {model.name} has no function call defined")
|
|
|
|
# Calculate timing and output bytes
|
|
endTime = time.time()
|
|
processingTime = endTime - startTime
|
|
outputBytes = len(content.encode("utf-8"))
|
|
|
|
# Calculate price using model's own price calculation method
|
|
priceUsd = model.calculatePriceUsd(processingTime, inputBytes, outputBytes)
|
|
|
|
return AiCallResponse(
|
|
content=content,
|
|
modelName=model.name,
|
|
priceUsd=priceUsd,
|
|
processingTime=processingTime,
|
|
bytesSent=inputBytes,
|
|
bytesReceived=outputBytes,
|
|
errorCount=0
|
|
)
|
|
|
|
|
|
# Utility methods
|
|
async def listAvailableModels(self, connectorType: str = None) -> List[Dict[str, Any]]:
|
|
"""List available models, optionally filtered by connector type."""
|
|
models = modelRegistry.getAvailableModels()
|
|
if connectorType:
|
|
return [model.model_dump() for model in models if model.connectorType == connectorType]
|
|
return [model.model_dump() for model in models]
|
|
|
|
async def getModelInfo(self, displayName: str) -> Dict[str, Any]:
|
|
"""Get information about a specific model by displayName."""
|
|
model = modelRegistry.getModel(displayName)
|
|
if not model:
|
|
raise ValueError(f"Model with displayName '{displayName}' not found")
|
|
return model.model_dump()
|
|
|
|
async def getModelsByTag(self, tag: str) -> List[str]:
|
|
"""Get model displayNames that have a specific tag. Returns displayNames (unique identifiers)."""
|
|
models = modelRegistry.getModelsByTag(tag)
|
|
return [model.displayName for model in models]
|
|
|
|
|
|
|