129 lines
No EOL
5.5 KiB
Python
129 lines
No EOL
5.5 KiB
Python
import logging
|
|
from typing import Dict, Any, Optional
|
|
from datetime import datetime, UTC
|
|
import asyncio
|
|
|
|
from modules.methods import MethodBase, MethodResult
|
|
from modules.interfaces.serviceChatModel import AgentTask, AgentAction, AgentResult, Action, TaskStatus, ActionStatus
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class ServiceContainer:
|
|
"""Service container with improved state management"""
|
|
|
|
def __init__(self):
|
|
self.state = {
|
|
'status': TaskStatus.PENDING,
|
|
'retryCount': 0,
|
|
'retryMax': 3,
|
|
'timeout': 300, # 5 minutes
|
|
'lastError': None,
|
|
'lastErrorTime': None
|
|
}
|
|
self.methods: Dict[str, MethodBase] = {}
|
|
self.tasks: Dict[str, Any] = {} # Will be populated with AgentTask instances
|
|
|
|
def register_method(self, method: MethodBase) -> None:
|
|
"""Register a method in the container"""
|
|
self.methods[method.name] = method
|
|
logger.info(f"Registered method: {method.name}")
|
|
|
|
async def execute_task(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.can_execute_action(action):
|
|
if not task.get_auth_data(action.auth_source):
|
|
action.status = ActionStatus.FAILED
|
|
task.error = f"Missing authentication for {action.auth_source}"
|
|
else:
|
|
action.status = ActionStatus.DEPENDENCY_FAILED
|
|
continue
|
|
|
|
try:
|
|
# Get method
|
|
method = self.methods.get(action.method)
|
|
if not method:
|
|
raise ValueError(f"Unknown method: {action.method}")
|
|
|
|
# Validate parameters
|
|
if not await method.validate_parameters(action.action, action.parameters):
|
|
raise ValueError(f"Invalid parameters for {action.method}:{action.action}")
|
|
|
|
# Get auth data if needed
|
|
auth_data = None
|
|
if action.auth_source and action.auth_source != "local":
|
|
auth_data = task.get_auth_data(action.auth_source)
|
|
if not auth_data:
|
|
raise ValueError(f"Missing authentication data for {action.auth_source}")
|
|
|
|
# Execute with timeout
|
|
result = await asyncio.wait_for(
|
|
method.execute(action.action, action.parameters, auth_data),
|
|
timeout=action.timeout or 60
|
|
)
|
|
|
|
if result.success:
|
|
action.status = ActionStatus.SUCCESS
|
|
else:
|
|
if self._should_retry(result.data.get('error')):
|
|
action.retryCount += 1
|
|
if action.retryCount > action.retryMax:
|
|
action.status = ActionStatus.FAILED
|
|
if action.rollback_on_failure:
|
|
await method.rollback(action.action, action.parameters, auth_data)
|
|
else:
|
|
action.status = ActionStatus.RETRY
|
|
else:
|
|
action.status = ActionStatus.FAILED
|
|
if action.rollback_on_failure:
|
|
await method.rollback(action.action, action.parameters, auth_data)
|
|
|
|
except asyncio.TimeoutError:
|
|
action.status = ActionStatus.TIMEOUT
|
|
except Exception as e:
|
|
action.status = ActionStatus.FAILED
|
|
if action.rollback_on_failure:
|
|
await method.rollback(action.action, action.parameters, auth_data)
|
|
|
|
# Update task status
|
|
if task.has_failed():
|
|
task.status = TaskStatus.FAILED
|
|
elif task.is_complete():
|
|
task.status = TaskStatus.SUCCESS
|
|
task.finishedAt = datetime.now(UTC).isoformat()
|
|
|
|
except Exception as e:
|
|
task.status = TaskStatus.FAILED
|
|
task.error = str(e)
|
|
|
|
def _should_retry(self, error: str) -> bool:
|
|
"""Determine if error is retryable"""
|
|
retryable_errors = [
|
|
"AI down",
|
|
"Document not found",
|
|
"Content extraction failed",
|
|
"Network error",
|
|
"Temporary failure"
|
|
]
|
|
return any(err in error for err in retryable_errors)
|
|
|
|
def get_method(self, name: str) -> Optional[MethodBase]:
|
|
"""Get a method by name"""
|
|
return self.methods.get(name)
|
|
|
|
def get_available_methods(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()
|
|
} |