import logging import asyncio import uuid 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 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 for attempt, model in enumerate(failoverModelList): try: logger.info(f"Processing content part with model: {model.name} (attempt {attempt + 1}/{len(failoverModelList)})") # 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 merge_strategy = { "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 merge_strategy = { "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 placeholder in prompt for this specific model # Use maxTokens for output limit, not contextLength if model.maxTokens > 0: tokenLimit = str(model.maxTokens) else: tokenLimit = "16000" # Default for text generation # Create a copy of the prompt for this model call modelPrompt = prompt if "" in modelPrompt: modelPrompt = modelPrompt.replace("", tokenLimit) logger.debug(f"Replaced with {tokenLimit} for model {model.name}") # 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() # Get the connector for this model connector = modelRegistry.getConnectorForModel(model.name) if not connector: raise ValueError(f"No connector found for model {model.name}") # 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 ) async def callImage(self, prompt: str, imageData: Union[str, bytes], mimeType: str = None, options: AiCallOptions = None) -> AiCallResponse: """Call AI model for image analysis with fallback mechanism.""" if options is None: options = AiCallOptions(operationType=OperationTypeEnum.IMAGE_ANALYSE) # Get fallback models for image analysis availableModels = modelRegistry.getAvailableModels() failoverModelList = modelSelector.getFailoverModelList(prompt, "", options, availableModels) if not failoverModelList: errorMsg = f"No suitable models found for image analysis" 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 fallback sequence lastError = None for attempt, model in enumerate(failoverModelList): try: logger.info(f"Attempting image analysis with model: {model.name} (attempt {attempt + 1}/{len(failoverModelList)})") # Call the model response = await self._callImageWithModel(model, prompt, imageData, mimeType) logger.info(f"✅ Image analysis successful with model: {model.name}") return response except Exception as e: lastError = e logger.warning(f"❌ Image analysis 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 analysis...") continue else: # All models failed logger.error(f"💥 All {len(failoverModelList)} models failed for image analysis") break # All fallback attempts failed - return error response errorMsg = f"All AI models failed for image analysis. 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 _callImageWithModel(self, model: AiModel, prompt: str, imageData: Union[str, bytes], mimeType: str) -> AiCallResponse: """Call a specific model for image analysis and return the response.""" # Calculate input bytes from prompt and image data promptBytes = len(prompt.encode('utf-8')) if isinstance(imageData, str): # Base64 encoded string imageBytes = len(imageData.encode('utf-8')) else: # Raw bytes imageBytes = len(imageData) inputBytes = promptBytes + imageBytes # Start timing startTime = time.time() # Create standardized call object for image analysis modelCall = AiModelCall( messages=[{"role": "user", "content": prompt}], model=model, options=AiCallOptions(imageData=imageData, mimeType=mimeType) ) # Call the model with standardized interface if model.functionCall: 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 ) 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]