gateway/modules/interfaces/interfaceAiObjects.py
2025-10-29 00:38:57 +01:00

617 lines
27 KiB
Python

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.aicore.aicoreModelRegistry import modelRegistry
from modules.aicore.aicoreModelSelector import modelSelector
from modules.datamodels.datamodelAi import (
AiModel,
AiCallOptions,
AiCallRequest,
AiCallResponse,
OperationTypeEnum,
AiModelCall,
AiModelResponse,
)
from modules.datamodels.datamodelExtraction import ContentPart
# 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."""
# 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.name
# AI for Extraction and Text Generation
async def call(self, request: AiCallRequest) -> AiCallResponse:
"""Call AI model for text generation with model-aware chunking."""
# Handle content parts (unified path)
if hasattr(request, 'contentParts') and request.contentParts:
return await self._callWithContentParts(request)
# Handle traditional text/context calls
return await self._callWithTextContext(request)
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
# Compress optionally (prompt/context) - simple truncation fallback kept here
def _maybeTruncate(text: str, limit: int) -> str:
data = text.encode("utf-8")
if len(data) <= limit:
return text
return data[:limit].decode("utf-8", errors="ignore") + "... [truncated]"
if options.compressPrompt and len(prompt.encode("utf-8")) > 2000:
prompt = _maybeTruncate(prompt, 2000)
if options.compressContext and len(context.encode("utf-8")) > 70000:
context = _maybeTruncate(context, 70000)
# 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
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
)
async def _callWithContentParts(self, request: AiCallRequest) -> AiCallResponse:
"""Process content parts with model-aware chunking (unified for single and multiple parts)."""
prompt = request.prompt
options = request.options
contentParts = request.contentParts
# Get failover models
availableModels = modelRegistry.getAvailableModels()
failoverModelList = modelSelector.getFailoverModelList(prompt, "", options, availableModels)
if not failoverModelList:
return self._createErrorResponse("No suitable models found", 0, 0)
# Process each content part
allResults = []
for contentPart in contentParts:
partResult = await self._processContentPartWithFallback(contentPart, prompt, options, failoverModelList)
allResults.append(partResult)
# Merge all results
mergedContent = self._mergePartResults(allResults)
return AiCallResponse(
content=mergedContent,
modelName="multiple",
priceUsd=sum(r.priceUsd for r in allResults),
processingTime=sum(r.processingTime for r in allResults),
bytesSent=sum(r.bytesSent for r in allResults),
bytesReceived=sum(r.bytesReceived for r in allResults),
errorCount=sum(r.errorCount for r in allResults)
)
async def _processContentPartWithFallback(self, contentPart, prompt: str, options, failoverModelList) -> AiCallResponse:
"""Process a single content part with model-aware chunking and fallback."""
lastError = None
# Check if this is an image - Vision models need special handling
isImage = (contentPart.typeGroup == "image") or (contentPart.mimeType and contentPart.mimeType.startswith("image/"))
for attempt, model in enumerate(failoverModelList):
try:
logger.info(f"Processing content part with model: {model.name} (attempt {attempt + 1}/{len(failoverModelList)})")
# Special handling for images with Vision models
if isImage and hasattr(model, 'functionCall'):
# Call model's functionCall directly (for Vision models this is callAiImage)
from modules.datamodels.datamodelAi import AiModelCall, AiCallOptions as AiCallOpts
try:
modelCall = AiModelCall(
messages=[
{
"role": "user",
"content": [
{"type": "text", "text": prompt},
{
"type": "image_url",
"image_url": {
"url": f"data:{contentPart.mimeType};base64,{contentPart.data}" if isinstance(contentPart.data, str) else
f"data:{contentPart.mimeType};base64,{base64.b64encode(contentPart.data).decode('utf-8')}"
}
}
]
}
],
model=model,
options=AiCallOpts(operationType=options.operationType)
)
modelResponse = await model.functionCall(modelCall)
if not modelResponse.success:
raise ValueError(f"Model call failed: {modelResponse.error}")
logger.info(f"✅ Image content part processed successfully with model: {model.name}")
# Convert to AiCallResponse format
return AiCallResponse(
content=modelResponse.content,
modelName=model.name,
priceUsd=modelResponse.priceUsd if hasattr(modelResponse, 'priceUsd') else 0.0,
processingTime=modelResponse.processingTime if hasattr(modelResponse, 'processingTime') else 0.0,
bytesSent=0, # Will be calculated elsewhere
bytesReceived=0, # Will be calculated elsewhere
errorCount=0
)
except Exception as e:
# Image processing failed with this model
lastError = e
logger.warning(f"❌ Image processing 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 fallback model for image processing...")
continue
else:
# All models failed
logger.error(f"💥 All {len(failoverModelList)} models failed for image processing")
raise
# For non-image parts, check if part fits in model context
partSize = len(contentPart.data.encode('utf-8')) if contentPart.data else 0
modelContextBytes = model.contextLength * 4 # Convert tokens to bytes
if partSize <= modelContextBytes:
# Part fits - call AI directly
response = await self._callWithModel(model, prompt, contentPart.data, options)
logger.info(f"✅ Content part processed successfully with model: {model.name}")
return response
else:
# Part too large - chunk it
chunks = await self._chunkContentPart(contentPart, model, options)
if not chunks:
raise ValueError(f"Failed to chunk content part for model {model.name}")
# Process each chunk
chunkResults = []
for chunk in chunks:
chunkResponse = await self._callWithModel(model, prompt, chunk['data'], options)
chunkResults.append(chunkResponse)
# Merge chunk results
mergedContent = self._mergeChunkResults(chunkResults)
totalPrice = sum(r.priceUsd for r in chunkResults)
totalTime = sum(r.processingTime for r in chunkResults)
totalBytesSent = sum(r.bytesSent for r in chunkResults)
totalBytesReceived = sum(r.bytesReceived for r in chunkResults)
totalErrors = sum(r.errorCount for r in chunkResults)
logger.info(f"✅ Content part chunked and processed with model: {model.name} ({len(chunks)} chunks)")
return AiCallResponse(
content=mergedContent,
modelName=model.name,
priceUsd=totalPrice,
processingTime=totalTime,
bytesSent=totalBytesSent,
bytesReceived=totalBytesReceived,
errorCount=totalErrors
)
except Exception as e:
lastError = e
logger.warning(f"❌ Model {model.name} failed for content part: {str(e)}")
if attempt < len(failoverModelList) - 1:
logger.info(f"🔄 Trying next failover model...")
continue
else:
logger.error(f"💥 All {len(failoverModelList)} models failed for content part")
break
# All models failed
return self._createErrorResponse(f"All models failed: {str(lastError)}", 0, 0)
async def _chunkContentPart(self, contentPart, model, options) -> List[Dict[str, Any]]:
"""Chunk a content part based on model capabilities."""
# Calculate model-specific chunk sizes
modelContextBytes = model.contextLength * 4 # Convert tokens to bytes
maxContextBytes = int(modelContextBytes * 0.9) # 90% of context length
textChunkSize = int(maxContextBytes * 0.7) # 70% of max context for text chunks
imageChunkSize = int(maxContextBytes * 0.8) # 80% of max context for image chunks
# Build chunking options
chunkingOptions = {
"textChunkSize": textChunkSize,
"imageChunkSize": imageChunkSize,
"maxSize": maxContextBytes,
"chunkAllowed": True
}
# Get appropriate chunker
from modules.services.serviceExtraction.subRegistry import ChunkerRegistry
chunkerRegistry = ChunkerRegistry()
chunker = chunkerRegistry.resolve(contentPart.typeGroup)
if not chunker:
logger.warning(f"No chunker found for typeGroup: {contentPart.typeGroup}")
return []
# Chunk the content part
try:
chunks = chunker.chunk(contentPart, chunkingOptions)
logger.debug(f"Created {len(chunks)} chunks for {contentPart.typeGroup} part")
return chunks
except Exception as e:
logger.error(f"Chunking failed for {contentPart.typeGroup}: {str(e)}")
return []
def _mergePartResults(self, partResults: List[AiCallResponse]) -> str:
"""Merge part results using the existing sophisticated merging system."""
if not partResults:
return ""
# Convert AiCallResponse results to ContentParts for merging
from modules.datamodels.datamodelExtraction import ContentPart
from modules.services.serviceExtraction.subUtils import makeId
content_parts = []
for i, result in enumerate(partResults):
if result.content:
content_part = ContentPart(
id=str(uuid.uuid4()),
parentId=None,
label=f"ai_result_{i}",
typeGroup="text", # Default to text for AI results
mimeType="text/plain",
data=result.content,
metadata={
"aiResult": True,
"modelName": result.modelName,
"priceUsd": result.priceUsd,
"processingTime": result.processingTime,
"bytesSent": result.bytesSent,
"bytesReceived": result.bytesReceived
}
)
content_parts.append(content_part)
# Use existing merging system
from modules.datamodels.datamodelExtraction import MergeStrategy
merge_strategy = MergeStrategy(
useIntelligentMerging=True,
groupBy="typeGroup",
orderBy="id",
mergeType="concatenate"
)
from modules.services.serviceExtraction.subPipeline import _applyMerging
merged_parts = _applyMerging(content_parts, merge_strategy)
# Convert merged parts back to final string
final_content = "\n\n".join([part.data for part in merged_parts])
logger.info(f"Merged {len(partResults)} AI results using existing merging system")
return final_content.strip()
def _mergeChunkResults(self, chunkResults: List[AiCallResponse]) -> str:
"""Merge chunk results using the existing sophisticated merging system."""
if not chunkResults:
return ""
# Convert AiCallResponse results to ContentParts for merging
content_parts = []
for i, result in enumerate(chunkResults):
if result.content:
content_part = ContentPart(
id=str(uuid.uuid4()),
parentId=None,
label=f"chunk_result_{i}",
typeGroup="text", # Default to text for AI results
mimeType="text/plain",
data=result.content,
metadata={
"aiResult": True,
"chunk": True,
"modelName": result.modelName,
"priceUsd": result.priceUsd,
"processingTime": result.processingTime,
"bytesSent": result.bytesSent,
"bytesReceived": result.bytesReceived
}
)
content_parts.append(content_part)
# Use existing merging system
from modules.datamodels.datamodelExtraction import MergeStrategy
merge_strategy = MergeStrategy(
useIntelligentMerging=True,
groupBy="typeGroup",
orderBy="id",
mergeType="concatenate"
)
from modules.services.serviceExtraction.subPipeline import _applyMerging
merged_parts = _applyMerging(content_parts, merge_strategy)
# Convert merged parts back to final string
final_content = "\n\n".join([part.data for part in merged_parts])
logger.info(f"Merged {len(chunkResults)} chunk results using existing merging system")
return final_content.strip()
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 {}
)
# Call the model with standardized interface
modelResponse = await model.functionCall(modelCall)
# 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
)
# AI for Image Generation
async def generateImage(self, prompt: str, size: str = "1024x1024", quality: str = "standard", style: str = "vivid", options: AiCallOptions = None) -> AiCallResponse:
"""Generate an image using AI."""
if options is None:
options = AiCallOptions(operationType=OperationTypeEnum.IMAGE_GENERATE)
# Calculate input bytes
inputBytes = len(prompt.encode("utf-8"))
try:
# Select the best model for image generation
modelName = self._selectModel(prompt, "", options)
selectedModel = modelRegistry.getModel(modelName)
if not selectedModel:
raise ValueError(f"Selected model {modelName} not found in registry")
# Get the connector for this model
connector = modelRegistry.getConnectorForModel(modelName)
if not connector:
raise ValueError(f"No connector found for model {modelName}")
# Start timing
startTime = time.time()
# Create standardized call object for image generation
modelCall = AiModelCall(
messages=[{"role": "user", "content": prompt}],
model=selectedModel,
options=AiCallOptions(size=size, quality=quality, style=style)
)
# Call the model with standardized interface
if selectedModel.functionCall:
modelResponse = await selectedModel.functionCall(modelCall)
# 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 {modelName} 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 = selectedModel.calculatePriceUsd(processingTime, inputBytes, outputBytes)
logger.info(f"✅ Image generation successful with model: {modelName}")
return AiCallResponse(
success=True,
content=content,
modelName=modelName,
processingTime=processingTime,
priceUsd=priceUsd,
bytesSent=inputBytes,
bytesReceived=outputBytes,
errorCount=0
)
except Exception as e:
logger.error(f"❌ Image generation failed with model {modelName}: {str(e)}")
return AiCallResponse(
content=f"Image generation failed: {str(e)}",
modelName=modelName,
priceUsd=0.0,
processingTime=0.0,
bytesSent=inputBytes,
bytesReceived=0,
errorCount=1
)
# 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, modelName: str) -> Dict[str, Any]:
"""Get information about a specific model."""
model = modelRegistry.getModel(modelName)
if not model:
raise ValueError(f"Model {modelName} not found")
return model.model_dump()
async def getModelsByTag(self, tag: str) -> List[str]:
"""Get model names that have a specific tag."""
models = modelRegistry.getModelsByTag(tag)
return [model.name for model in models]