355 lines
No EOL
14 KiB
Python
355 lines
No EOL
14 KiB
Python
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 AgentTask, AgentAction, AgentResult, Action, TaskStatus, ActionStatus
|
|
from modules.interfaces.serviceManagementClass import ServiceManagement
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class ServiceContainer:
|
|
"""Service container for dependency injection and service management."""
|
|
|
|
def __init__(self):
|
|
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, Any] = {} # Will be populated with AgentTask 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()
|
|
} |