import logging from typing import Dict, Any, List, Optional from datetime import datetime, UTC import json import asyncio from modules.shared.configuration import APP_CONFIG from modules.methods import MethodBase, MethodResult from modules.interfaces.serviceChatModel import ( TaskStatus, UserInputRequest, ContentMetadata, ContentItem, ChatDocument, TaskDocument, ExtractedContent, TaskItem, TaskResult, ChatStat, ChatLog, ChatMessage, ChatWorkflow ) from modules.interfaces.serviceManagementClass import ServiceManagement from modules.interfaces.serviceChatClass import ChatInterface logger = logging.getLogger(__name__) class ServiceContainer: """Service container for dependency injection and service management.""" def __init__(self, chatInterface: ChatInterface): self.methods = {} self.context = {} self.workflow = None self.model = {} self.functions = {} self.documentProcessor = None self.state = { 'status': TaskStatus.PENDING, 'retryCount': 0, 'retryMax': 3, 'timeout': 300, # 5 minutes 'lastError': None, 'lastErrorTime': None } self.tasks: Dict[str, TaskItem] = {} # Will be populated with TaskItem instances # Initialize service management self.serviceManagement = ServiceManagement() # Initialize file-related functions self.functions = { 'getFileData': self.serviceManagement.getFileData, 'saveFileData': self.serviceManagement.saveFileData, 'getFileMetadata': self.serviceManagement.getFileMetadata, 'saveFileMetadata': self.serviceManagement.saveFileMetadata, 'deleteFile': self.serviceManagement.deleteFile, 'getFile': self.serviceManagement.getFile, 'getMimeType': self.serviceManagement.getMimeType, 'calculateFileHash': self.serviceManagement.calculateFileHash, 'checkForDuplicateFile': self.serviceManagement.checkForDuplicateFile } def initialize(self) -> None: """Initialize service container""" pass def registerMethod(self, methodName: str, methodInstance: Any) -> None: """Register a new method""" self.methods[methodName] = methodInstance def getMethod(self, methodName: str) -> Optional[Any]: """Get a method by name""" return self.methods.get(methodName) def removeMethod(self, methodName: str) -> None: """Remove a method""" self.methods.pop(methodName, None) def hasMethod(self, methodName: str) -> bool: """Check if a method exists""" return methodName in self.methods def listMethods(self) -> List[str]: """List all registered methods""" return list(self.methods.keys()) def getMethodInfo(self, methodName: str) -> Dict[str, Any]: """Get method information""" method = self.getMethod(methodName) if not method: return {} return { "name": methodName, "description": self.getMethodDescription(methodName), "version": self.getMethodVersion(methodName), "author": self.getMethodAuthor(methodName), "license": self.getMethodLicense(methodName), "dependencies": self.getMethodDependencies(methodName), "tags": self.getMethodTags(methodName), "examples": self.getMethodExamples(methodName), "documentation": self.getMethodDocumentation(methodName), "source": self.getMethodSource(methodName), "tests": self.getMethodTests(methodName), "benchmarks": self.getMethodBenchmarks(methodName), "metrics": self.getMethodMetrics(methodName), "logs": self.getMethodLogs(methodName), "history": self.getMethodHistory(methodName), "usage": self.getMethodUsage(methodName), "errors": self.getMethodErrors(methodName), "warnings": self.getMethodWarnings(methodName) } def getMethodSchema(self, methodName: str) -> Optional[Dict[str, Any]]: """Get method schema""" method = self.getMethod(methodName) return method.schema if method else None def getMethodParameters(self, methodName: str) -> Optional[Dict[str, Any]]: """Get method parameters""" method = self.getMethod(methodName) return method.parameters if method else None def getMethodReturnType(self, methodName: str) -> Optional[str]: """Get method return type""" method = self.getMethod(methodName) return method.returnType if method else None def getMethodDescription(self, methodName: str) -> Optional[str]: """Get method description""" method = self.getMethod(methodName) return method.description if method else None def getMethodVersion(self, methodName: str) -> Optional[str]: """Get method version""" method = self.getMethod(methodName) return method.version if method else None def getMethodAuthor(self, methodName: str) -> Optional[str]: """Get method author""" method = self.getMethod(methodName) return method.author if method else None def getMethodLicense(self, methodName: str) -> Optional[str]: """Get method license""" method = self.getMethod(methodName) return method.license if method else None def getMethodDependencies(self, methodName: str) -> Optional[List[str]]: """Get method dependencies""" method = self.getMethod(methodName) return method.dependencies if method else None def getMethodTags(self, methodName: str) -> Optional[List[str]]: """Get method tags""" method = self.getMethod(methodName) return method.tags if method else None def getMethodExamples(self, methodName: str) -> Optional[List[Dict[str, Any]]]: """Get method examples""" method = self.getMethod(methodName) return method.examples if method else None def getMethodDocumentation(self, methodName: str) -> Optional[str]: """Get method documentation""" method = self.getMethod(methodName) return method.documentation if method else None def getMethodSource(self, methodName: str) -> Optional[str]: """Get method source""" method = self.getMethod(methodName) return method.source if method else None def getMethodTests(self, methodName: str) -> Optional[List[Dict[str, Any]]]: """Get method tests""" method = self.getMethod(methodName) return method.tests if method else None def getMethodBenchmarks(self, methodName: str) -> Optional[List[Dict[str, Any]]]: """Get method benchmarks""" method = self.getMethod(methodName) return method.benchmarks if method else None def getMethodMetrics(self, methodName: str) -> Optional[Dict[str, Any]]: """Get method metrics""" method = self.getMethod(methodName) return method.metrics if method else None def getMethodLogs(self, methodName: str) -> Optional[List[Dict[str, Any]]]: """Get method logs""" method = self.getMethod(methodName) return method.logs if method else None def getMethodHistory(self, methodName: str) -> Optional[List[Dict[str, Any]]]: """Get method history""" method = self.getMethod(methodName) return method.history if method else None def getMethodUsage(self, methodName: str) -> Optional[Dict[str, Any]]: """Get method usage""" method = self.getMethod(methodName) return method.usage if method else None def getMethodErrors(self, methodName: str) -> Optional[List[Dict[str, Any]]]: """Get method errors""" method = self.getMethod(methodName) return method.errors if method else None def getMethodWarnings(self, methodName: str) -> Optional[List[Dict[str, Any]]]: """Get method warnings""" method = self.getMethod(methodName) return method.warnings if method else None def executeTask(self, task: Any) -> None: """Execute a task""" try: # Execute each action for action in task.actionList: method = self.getMethod(action.method) if method: method.executeAction(action.action, action.parameters) except Exception as e: logger.error(f"Error executing task: {str(e)}") raise def getFileData(self, fileId: str) -> bytes: """Get file data by ID""" try: # Get file data from storage if hasattr(self.functions, 'getFileData'): return self.functions.getFileData(fileId) return b"" except Exception as e: logger.error(f"Error getting file data: {str(e)}") return b"" def saveFileData(self, fileId: str, data: bytes) -> bool: """Save file data by ID""" try: # Save file data to storage if hasattr(self.functions, 'saveFileData'): return self.functions.saveFileData(fileId, data) return False except Exception as e: logger.error(f"Error saving file data: {str(e)}") return False def getFileMetadata(self, fileId: str) -> Dict[str, Any]: """Get file metadata by ID""" try: # Get file metadata from storage if hasattr(self.functions, 'getFileMetadata'): return self.functions.getFileMetadata(fileId) return {} except Exception as e: logger.error(f"Error getting file metadata: {str(e)}") return {} def saveFileMetadata(self, fileId: str, metadata: Dict[str, Any]) -> bool: """Save file metadata by ID""" try: # Save file metadata to storage if hasattr(self.functions, 'saveFileMetadata'): return self.functions.saveFileMetadata(fileId, metadata) return False except Exception as e: logger.error(f"Error saving file metadata: {str(e)}") return False async def executeTaskImproved(self, task: Any) -> None: # task: AgentTask """Execute task with improved error handling and timeout""" try: # Check for timeout if (datetime.now(UTC) - datetime.fromisoformat(task.startedAt)).seconds > self.state['timeout']: task.status = TaskStatus.TIMEOUT return # Execute actions for action in task.actionList: if not task.canExecuteAction(action): if not task.getAuthData(action.authSource): action.status = ActionStatus.FAILED task.error = f"Missing authentication for {action.authSource}" else: action.status = ActionStatus.DEPENDENCY_FAILED continue try: # Get method method = self.getMethod(action.method) if not method: raise ValueError(f"Unknown method: {action.method}") # Validate parameters if not await method.validateParameters(action.action, action.parameters): raise ValueError(f"Invalid parameters for {action.method}:{action.action}") # Get auth data if needed authData = None if action.authSource and action.authSource != "local": authData = task.getAuthData(action.authSource) if not authData: raise ValueError(f"Missing authentication data for {action.authSource}") # Execute with timeout result = await asyncio.wait_for( method.execute(action.action, action.parameters, authData), timeout=action.timeout or 60 ) if result.success: action.status = ActionStatus.SUCCESS else: if self._shouldRetry(result.data.get('error')): action.retryCount += 1 if action.retryCount > action.retryMax: action.status = ActionStatus.FAILED if action.rollbackOnFailure: await method.rollback(action.action, action.parameters, authData) else: action.status = ActionStatus.RETRY else: action.status = ActionStatus.FAILED if action.rollbackOnFailure: await method.rollback(action.action, action.parameters, authData) except asyncio.TimeoutError: action.status = ActionStatus.TIMEOUT except Exception as e: action.status = ActionStatus.FAILED if action.rollbackOnFailure: await method.rollback(action.action, action.parameters, authData) # Update task status if task.hasFailed(): task.status = TaskStatus.FAILED elif task.isComplete(): task.status = TaskStatus.SUCCESS task.finishedAt = datetime.now(UTC).isoformat() except Exception as e: task.status = TaskStatus.FAILED task.error = str(e) def _shouldRetry(self, error: str) -> bool: """Determine if error is retryable""" retryableErrors = [ "AI down", "Document not found", "Content extraction failed", "Network error", "Temporary failure" ] return any(err in error for err in retryableErrors) def getAvailableMethodsCatalog(self) -> Dict[str, Dict[str, Any]]: """Get catalog of available methods and their actions""" return { name: { "description": method.description, "actions": method.actions } for name, method in self.methods.items() }