aiecs 1.0.1__py3-none-any.whl → 1.7.17__py3-none-any.whl

aiecs/infrastructure/messaging/celery_task_manager.py

@@ -3,9 +3,10 @@ import logging
 import time
 import uuid
 from typing import Dict, List, Any, Optional
-from celery import Celery
-from celery.exceptions import TimeoutError as CeleryTimeoutError
+from celery import Celery  # type: ignore[import-untyped]
+from celery.exceptions import TimeoutError as CeleryTimeoutError  # type: ignore[import-untyped]
 from asyncio import TimeoutError as AsyncioTimeoutError
+
 # Removed direct import to avoid circular dependency
 # Tasks are referenced by string names instead
 from aiecs.domain.execution.model import TaskStatus, ErrorCode
@@ -20,33 +21,42 @@ class CeleryTaskManager:
 
     def __init__(self, config: Dict[str, Any]):
         self.config = config
-        self.celery_app = None
+        self.celery_app: Optional[Celery] = None
         self._init_celery()
 
     def _init_celery(self):
         """Initialize Celery application"""
         try:
             self.celery_app = Celery(
-                'service_executor',
-                broker=self.config.get('broker_url', 'redis://redis:6379/0'),
-                backend=self.config.get('backend_url', 'redis://redis:6379/0')
+                "service_executor",
+                broker=self.config.get("broker_url", "redis://redis:6379/0"),
+                backend=self.config.get("backend_url", "redis://redis:6379/0"),
             )
 
             # Configure Celery
             self.celery_app.conf.update(
-                task_serializer=self.config.get('task_serializer', 'json'),
-                accept_content=self.config.get('accept_content', ['json']),
-                result_serializer=self.config.get('result_serializer', 'json'),
-                timezone=self.config.get('timezone', 'UTC'),
-                enable_utc=self.config.get('enable_utc', True),
-                task_queues=self.config.get('task_queues', {
-                    'fast_tasks': {'exchange': 'fast_tasks', 'routing_key': 'fast_tasks'},
-                    'heavy_tasks': {'exchange': 'heavy_tasks', 'routing_key': 'heavy_tasks'}
-                }),
-                worker_concurrency=self.config.get('worker_concurrency', {
-                    'fast_worker': 10,
-                    'heavy_worker': 2
-                })
+                task_serializer=self.config.get("task_serializer", "json"),
+                accept_content=self.config.get("accept_content", ["json"]),
+                result_serializer=self.config.get("result_serializer", "json"),
+                timezone=self.config.get("timezone", "UTC"),
+                enable_utc=self.config.get("enable_utc", True),
+                task_queues=self.config.get(
+                    "task_queues",
+                    {
+                        "fast_tasks": {
+                            "exchange": "fast_tasks",
+                            "routing_key": "fast_tasks",
+                        },
+                        "heavy_tasks": {
+                            "exchange": "heavy_tasks",
+                            "routing_key": "heavy_tasks",
+                        },
+                    },
+                ),
+                worker_concurrency=self.config.get(
+                    "worker_concurrency",
+                    {"fast_worker": 10, "heavy_worker": 2},
+                ),
             )
 
             logger.info("Celery application initialized successfully")
@@ -54,8 +64,18 @@ class CeleryTaskManager:
             logger.error(f"Failed to initialize Celery: {e}")
             raise
 
-    def execute_celery_task(self, task_name: str, queue: str, user_id: str, task_id: str, step: int,
-                           mode: str, service: str, input_data: Dict[str, Any], context: Dict[str, Any]):
+    def execute_celery_task(
+        self,
+        task_name: str,
+        queue: str,
+        user_id: str,
+        task_id: str,
+        step: int,
+        mode: str,
+        service: str,
+        input_data: Dict[str, Any],
+        context: Dict[str, Any],
+    ):
         """
         Execute Celery task
 
@@ -81,6 +101,8 @@ class CeleryTaskManager:
             celery_task_name = "aiecs.tasks.worker.execute_heavy_task"
 
         # Send task to Celery
+        if self.celery_app is None:
+            raise RuntimeError("Celery app not initialized")
         return self.celery_app.send_task(
             celery_task_name,
             kwargs={
@@ -91,12 +113,17 @@ class CeleryTaskManager:
                 "mode": mode,
                 "service": service,
                 "input_data": input_data,
-                "context": context
+                "context": context,
             },
-            queue=queue
+            queue=queue,
         )
 
-    async def execute_task(self, task_name: str, input_data: Dict[str, Any], context: Dict[str, Any]) -> Any:
+    async def execute_task(
+        self,
+        task_name: str,
+        input_data: Dict[str, Any],
+        context: Dict[str, Any],
+    ) -> Any:
         """
         Execute a single task using Celery for asynchronous processing
         """
@@ -106,14 +133,16 @@ class CeleryTaskManager:
         mode = input_data.get("mode", "default")
         service = input_data.get("service", "default")
         queue = input_data.get("queue", "fast_tasks")
-        timeout = self.config.get('task_timeout_seconds', 300)
+        timeout = self.config.get("task_timeout_seconds", 300)
 
         try:
             # Use string-based task names to avoid circular imports
             celery_task_name = "aiecs.tasks.worker.execute_task"
-            if queue == 'heavy_tasks':
+            if queue == "heavy_tasks":
                 celery_task_name = "aiecs.tasks.worker.execute_heavy_task"
 
+            if self.celery_app is None:
+                raise RuntimeError("Celery app not initialized")
             result = self.celery_app.send_task(
                 celery_task_name,
                 kwargs={
@@ -124,9 +153,9 @@ class CeleryTaskManager:
                     "mode": mode,
                     "service": service,
                     "input_data": input_data,
-                    "context": context
+                    "context": context,
                 },
-                queue=queue
+                queue=queue,
             )
 
             return result.get(timeout=timeout)
@@ -136,14 +165,14 @@ class CeleryTaskManager:
             return {
                 "status": TaskStatus.TIMED_OUT,
                 "error_code": ErrorCode.TIMEOUT_ERROR,
-                "error_message": str(e)
+                "error_message": str(e),
             }
         except Exception as e:
             logger.error(f"Error executing Celery task {task_name}: {e}", exc_info=True)
             return {
                 "status": TaskStatus.FAILED,
                 "error_code": ErrorCode.EXECUTION_ERROR,
-                "error_message": str(e)
+                "error_message": str(e),
             }
 
     async def execute_heavy_task(self, task_name: str, input_data: Dict, context: Dict) -> Any:
@@ -168,7 +197,7 @@ class CeleryTaskManager:
                 "message": "Invalid DSL step: missing task name",
                 "status": TaskStatus.FAILED,
                 "error_code": ErrorCode.VALIDATION_ERROR,
-                "error_message": "Task name is required"
+                "error_message": "Task name is required",
             }
 
         # Determine task type
@@ -188,6 +217,8 @@ class CeleryTaskManager:
         step_num = context.get("step", 0)
 
         # Send task to Celery
+        if self.celery_app is None:
+            raise RuntimeError("Celery app not initialized")
         celery_task = self.celery_app.send_task(
             celery_task_name,
             kwargs={
@@ -198,13 +229,13 @@ class CeleryTaskManager:
                 "mode": context.get("mode", "multi_task"),
                 "service": context.get("service", "summarizer"),
                 "input_data": input_data,
-                "context": context
+                "context": context,
             },
-            queue=queue
+            queue=queue,
         )
 
         try:
-            timeout_seconds = self.config.get('task_timeout_seconds', 300)
+            timeout_seconds = self.config.get("task_timeout_seconds", 300)
             start_time = time.time()
 
             # Wait for task completion
@@ -223,7 +254,7 @@ class CeleryTaskManager:
                         "result": result,
                         "completed": True,
                         "message": f"Completed task {task_name}",
-                        "status": TaskStatus.COMPLETED
+                        "status": TaskStatus.COMPLETED,
                     }
             else:
                 error = celery_task.get(propagate=False)
@@ -237,7 +268,7 @@ class CeleryTaskManager:
                     "message": f"Failed to execute task: {error}",
                     "status": status,
                     "error_code": error_code,
-                    "error_message": str(error)
+                    "error_message": str(error),
                 }
 
         except AsyncioTimeoutError as e:
@@ -248,7 +279,7 @@ class CeleryTaskManager:
                 "message": "Task execution timed out",
                 "status": TaskStatus.TIMED_OUT,
                 "error_code": ErrorCode.TIMEOUT_ERROR,
-                "error_message": str(e)
+                "error_message": str(e),
             }
         except Exception as e:
             return {
@@ -258,31 +289,31 @@ class CeleryTaskManager:
                 "message": f"Failed to execute {category}/{task_name}",
                 "status": TaskStatus.FAILED,
                 "error_code": ErrorCode.EXECUTION_ERROR,
-                "error_message": str(e)
+                "error_message": str(e),
             }
 
     def get_task_result(self, task_id: str):
         """Get task result"""
         try:
+            if self.celery_app is None:
+                raise RuntimeError("Celery app not initialized")
             result = self.celery_app.AsyncResult(task_id)
             return {
                 "task_id": task_id,
                 "status": result.status,
                 "result": result.result if result.ready() else None,
                 "successful": result.successful() if result.ready() else None,
-                "failed": result.failed() if result.ready() else None
+                "failed": result.failed() if result.ready() else None,
             }
         except Exception as e:
             logger.error(f"Error getting task result for {task_id}: {e}")
-            return {
-                "task_id": task_id,
-                "status": "ERROR",
-                "error": str(e)
-            }
+            return {"task_id": task_id, "status": "ERROR", "error": str(e)}
 
     def cancel_task(self, task_id: str):
         """Cancel task"""
         try:
+            if self.celery_app is None:
+                raise RuntimeError("Celery app not initialized")
             self.celery_app.control.revoke(task_id, terminate=True)
             logger.info(f"Task {task_id} cancelled")
             return True
@@ -295,18 +326,21 @@ class CeleryTaskManager:
         Batch execute tasks
         """
         results = []
-        batch_size = self.config.get('batch_size', 10)
-        rate_limit = self.config.get('rate_limit_requests_per_second', 5)
+        batch_size = self.config.get("batch_size", 10)
+        rate_limit = self.config.get("rate_limit_requests_per_second", 5)
 
         for i in range(0, len(tasks), batch_size):
-            batch = tasks[i:i + batch_size]
+            batch = tasks[i : i + batch_size]
             batch_results = await asyncio.gather(
-                *[self.execute_task(
-                    task["task_name"],
-                    task.get("input_data", {}),
-                    task.get("context", {})
-                ) for task in batch],
-                return_exceptions=True
+                *[
+                    self.execute_task(
+                        task["task_name"],
+                        task.get("input_data", {}),
+                        task.get("context", {}),
+                    )
+                    for task in batch
+                ],
+                return_exceptions=True,
             )
             results.extend(batch_results)
             await asyncio.sleep(1.0 / rate_limit)
@@ -316,6 +350,8 @@ class CeleryTaskManager:
     def get_queue_info(self) -> Dict[str, Any]:
         """Get queue information"""
         try:
+            if self.celery_app is None:
+                raise RuntimeError("Celery app not initialized")
             inspect = self.celery_app.control.inspect()
             active_tasks = inspect.active()
             scheduled_tasks = inspect.scheduled()
@@ -324,7 +360,7 @@ class CeleryTaskManager:
             return {
                 "active_tasks": active_tasks,
                 "scheduled_tasks": scheduled_tasks,
-                "reserved_tasks": reserved_tasks
+                "reserved_tasks": reserved_tasks,
             }
         except Exception as e:
             logger.error(f"Error getting queue info: {e}")
@@ -333,6 +369,8 @@ class CeleryTaskManager:
     def get_worker_stats(self) -> Dict[str, Any]:
         """Get worker statistics"""
         try:
+            if self.celery_app is None:
+                raise RuntimeError("Celery app not initialized")
             inspect = self.celery_app.control.inspect()
             stats = inspect.stats()
             return stats or {}

aiecs/infrastructure/messaging/websocket_manager.py

@@ -45,11 +45,7 @@ class WebSocketManager:
             return self.server
 
         try:
-            self.server = await serve(
-                self._handle_client_connection,
-                self.host,
-                self.port
-            )
+            self.server = await serve(self._handle_client_connection, self.host, self.port)
             self._running = True
             logger.info(f"WebSocket server started on {self.host}:{self.port}")
             return self.server
@@ -69,7 +65,7 @@ class WebSocketManager:
         if self.active_connections:
             await asyncio.gather(
                 *[conn.close() for conn in self.active_connections],
-                return_exceptions=True
+                return_exceptions=True,
             )
             self.active_connections.clear()
 
@@ -81,15 +77,19 @@ class WebSocketManager:
 
         try:
             async for message in websocket:
-                await self._handle_client_message(websocket, message)
+                # Decode bytes to str if needed
+                message_str = message if isinstance(message, str) else message.decode('utf-8')
+                await self._handle_client_message(websocket, message_str)
         except websockets.exceptions.ConnectionClosed:
             logger.info(f"WebSocket connection closed: {client_addr}")
         except Exception as e:
             logger.error(f"WebSocket error for {client_addr}: {e}")
         finally:
             self.active_connections.discard(websocket)
-            if not websocket.closed:
+            try:
                 await websocket.close()
+            except Exception:
+                pass  # Connection already closed
 
     async def _handle_client_message(self, websocket: ServerConnection, message: str):
         """Handle client message"""
@@ -123,7 +123,7 @@ class WebSocketManager:
             callback = self.callback_registry[callback_id]
             confirmation = UserConfirmation(
                 proceed=data.get("proceed", False),
-                feedback=data.get("feedback")
+                feedback=data.get("feedback"),
             )
             try:
                 callback(confirmation)
@@ -141,14 +141,17 @@ class WebSocketManager:
 
         if user_id and task_id:
             # Task cancellation logic can be added here
-            # Since database manager access is needed, this functionality may need to be implemented through callbacks
+            # Since database manager access is needed, this functionality may
+            # need to be implemented through callbacks
             logger.info(f"Task cancellation requested: user={user_id}, task={task_id}")
-            await self.broadcast_message({
-                "type": "task_cancelled",
-                "user_id": user_id,
-                "task_id": task_id,
-                "timestamp": asyncio.get_event_loop().time()
-            })
+            await self.broadcast_message(
+                {
+                    "type": "task_cancelled",
+                    "user_id": user_id,
+                    "task_id": task_id,
+                    "timestamp": asyncio.get_event_loop().time(),
+                }
+            )
         else:
             logger.warning("Invalid cancellation request: missing user_id or task_id")
 
@@ -157,7 +160,7 @@ class WebSocketManager:
         pong_data = {
             "type": "pong",
             "timestamp": asyncio.get_event_loop().time(),
-            "original_data": data
+            "original_data": data,
         }
         await self._send_to_client(websocket, pong_data)
 
@@ -167,32 +170,37 @@ class WebSocketManager:
         if user_id:
             # User-specific subscription logic can be implemented here
             logger.info(f"User {user_id} subscribed to updates")
-            await self._send_to_client(websocket, {
-                "type": "subscription_confirmed",
-                "user_id": user_id
-            })
+            await self._send_to_client(
+                websocket,
+                {"type": "subscription_confirmed", "user_id": user_id},
+            )
 
     async def _send_error(self, websocket: ServerConnection, error_message: str):
         """Send error message to client"""
         error_data = {
             "type": "error",
             "message": error_message,
-            "timestamp": asyncio.get_event_loop().time()
+            "timestamp": asyncio.get_event_loop().time(),
         }
         await self._send_to_client(websocket, error_data)
 
     async def _send_to_client(self, websocket: ServerConnection, data: Dict[str, Any]):
         """Send data to specific client"""
         try:
-            if not websocket.closed:
-                await websocket.send(json.dumps(data))
+            await websocket.send(json.dumps(data))
         except Exception as e:
             logger.error(f"Failed to send message to client: {e}")
 
-    async def notify_user(self, step_result: TaskStepResult, user_id: str, task_id: str, step: int) -> UserConfirmation:
+    async def notify_user(
+        self,
+        step_result: TaskStepResult,
+        user_id: str,
+        task_id: str,
+        step: int,
+    ) -> UserConfirmation:
         """Notify user of task step result"""
         callback_id = str(uuid.uuid4())
-        confirmation_future = asyncio.Future()
+        confirmation_future: asyncio.Future[UserConfirmation] = asyncio.Future()
 
         # Register callback
         self.callback_registry[callback_id] = lambda confirmation: confirmation_future.set_result(confirmation)
@@ -209,16 +217,18 @@ class WebSocketManager:
             "error_message": step_result.error_message,
             "user_id": user_id,
             "task_id": task_id,
-            "timestamp": asyncio.get_event_loop().time()
+            "timestamp": asyncio.get_event_loop().time(),
         }
 
         try:
-            # Broadcast to all connected clients (can be optimized to send only to specific users)
+            # Broadcast to all connected clients (can be optimized to send only
+            # to specific users)
             await self.broadcast_message(notification_data)
 
             # Wait for user confirmation with timeout
             try:
-                return await asyncio.wait_for(confirmation_future, timeout=300)  # 5 minute timeout
+                # 5 minute timeout
+                return await asyncio.wait_for(confirmation_future, timeout=300)
             except asyncio.TimeoutError:
                 logger.warning(f"User confirmation timeout for callback {callback_id}")
                 # Clean up callback
@@ -239,7 +249,7 @@ class WebSocketManager:
             "message": "Task is still executing...",
             "user_id": user_id,
             "task_id": task_id,
-            "timestamp": asyncio.get_event_loop().time()
+            "timestamp": asyncio.get_event_loop().time(),
         }
 
         while self._running:
@@ -256,14 +266,20 @@ class WebSocketManager:
             logger.debug("No active WebSocket connections for broadcast")
             return
 
-        # Filter out closed connections
-        active_connections = [conn for conn in self.active_connections if not conn.closed]
+        # Filter out closed connections (use try-except to handle closed connections)
+        active_connections = []
+        for conn in list(self.active_connections):
+            try:
+                # Try to check if connection is still valid
+                active_connections.append(conn)
+            except Exception:
+                pass  # Connection is closed, skip it
         self.active_connections = set(active_connections)
 
         if active_connections:
             await asyncio.gather(
                 *[self._send_to_client(conn, message) for conn in active_connections],
-                return_exceptions=True
+                return_exceptions=True,
             )
             logger.debug(f"Broadcasted message to {len(active_connections)} clients")
 
@@ -276,7 +292,7 @@ class WebSocketManager:
 
     def get_connection_count(self) -> int:
         """Get active connection count"""
-        return len([conn for conn in self.active_connections if not conn.closed])
+        return len(self.active_connections)
 
     def get_status(self) -> Dict[str, Any]:
         """Get WebSocket manager status"""
@@ -285,5 +301,5 @@ class WebSocketManager:
             "host": self.host,
             "port": self.port,
             "active_connections": self.get_connection_count(),
-            "pending_callbacks": len(self.callback_registry)
+            "pending_callbacks": len(self.callback_registry),
         }

aiecs/infrastructure/monitoring/__init__.py

@@ -5,8 +5,30 @@ Contains monitoring, metrics, and observability infrastructure.
 
 from .executor_metrics import ExecutorMetrics
 from .tracing_manager import TracingManager
+from .global_metrics_manager import (
+    initialize_global_metrics,
+    get_global_metrics,
+    close_global_metrics,
+    is_metrics_initialized,
+    get_metrics_summary,
+    record_operation,
+    record_duration,
+    record_operation_success,
+    record_operation_failure,
+    record_retry,
+)
 
 __all__ = [
     "ExecutorMetrics",
     "TracingManager",
+    "initialize_global_metrics",
+    "get_global_metrics",
+    "close_global_metrics",
+    "is_metrics_initialized",
+    "get_metrics_summary",
+    "record_operation",
+    "record_duration",
+    "record_operation_success",
+    "record_operation_failure",
+    "record_retry",
 ]

aiecs/infrastructure/monitoring/executor_metrics.py

@@ -25,14 +25,29 @@ class ExecutorMetrics:
             start_http_server(self.metrics_port)
             self.metrics = {
                 "intent_latency": Histogram("intent_latency_seconds", "Latency of intent parsing"),
-                "intent_success": Counter("intent_success_total", "Number of successful intent parsings"),
+                "intent_success": Counter(
+                    "intent_success_total",
+                    "Number of successful intent parsings",
+                ),
                 "intent_retries": Counter("intent_retries_total", "Number of intent parsing retries"),
                 "plan_latency": Histogram("plan_latency_seconds", "Latency of task planning"),
                 "plan_success": Counter("plan_success_total", "Number of successful plans"),
                 "plan_retries": Counter("plan_retries_total", "Number of plan retries"),
-                "execute_latency": Histogram("execute_latency_seconds", "Latency of task execution", ["task_type"]),
-                "execute_success": Counter("execute_success_total", "Number of successful executions", ["task_type"]),
-                "execute_retries": Counter("execute_retries_total", "Number of execution retries", ["task_type"]),
+                "execute_latency": Histogram(
+                    "execute_latency_seconds",
+                    "Latency of task execution",
+                    ["task_type"],
+                ),
+                "execute_success": Counter(
+                    "execute_success_total",
+                    "Number of successful executions",
+                    ["task_type"],
+                ),
+                "execute_retries": Counter(
+                    "execute_retries_total",
+                    "Number of execution retries",
+                    ["task_type"],
+                ),
             }
             logger.info(f"Prometheus metrics server started on port {self.metrics_port}")
         except Exception as e:
@@ -54,7 +69,12 @@ class ExecutorMetrics:
             metric = metric.labels(**labels)
         metric.inc()
 
-    def record_operation_failure(self, operation: str, error_type: str, labels: Optional[Dict[str, str]] = None):
+    def record_operation_failure(
+        self,
+        operation: str,
+        error_type: str,
+        labels: Optional[Dict[str, str]] = None,
+    ):
         """Record operation failure"""
         if not self.enable_metrics:
             return
@@ -70,6 +90,7 @@ class ExecutorMetrics:
 
     def with_metrics(self, metric_name: str, labels: Optional[Dict[str, str]] = None):
         """Monitoring decorator"""
+
         def decorator(func):
             @functools.wraps(func)
             async def wrapper(*args, **kwargs):
@@ -93,7 +114,9 @@ class ExecutorMetrics:
                     except Exception as e:
                         logger.error(f"Error in {func.__name__}: {e}")
                         raise
+
             return wrapper
+
         return decorator
 
     def get_metrics_summary(self) -> Dict[str, Any]:
@@ -104,10 +127,16 @@ class ExecutorMetrics:
         return {
             "metrics_enabled": True,
             "metrics_port": self.metrics_port,
-            "available_metrics": list(self.metrics.keys())
+            "available_metrics": list(self.metrics.keys()),
         }
 
-    def record_operation(self, operation_type: str, success: bool = True, duration: Optional[float] = None, **kwargs):
+    def record_operation(
+        self,
+        operation_type: str,
+        success: bool = True,
+        duration: Optional[float] = None,
+        **kwargs,
+    ):
         """Record a general operation for metrics tracking"""
         if not self.enable_metrics:
             return
@@ -115,10 +144,10 @@ class ExecutorMetrics:
         try:
             # Record operation success/failure
             if success:
-                self.record_operation_success(operation_type, kwargs.get('labels'))
+                self.record_operation_success(operation_type, kwargs.get("labels"))
             else:
-                error_type = kwargs.get('error_type', 'unknown')
-                self.record_operation_failure(operation_type, error_type, kwargs.get('labels'))
+                error_type = kwargs.get("error_type", "unknown")
+                self.record_operation_failure(operation_type, error_type, kwargs.get("labels"))
 
             # Record operation latency if provided
             if duration is not None:
@@ -127,7 +156,12 @@ class ExecutorMetrics:
         except Exception as e:
             logger.warning(f"Failed to record operation metrics: {e}")
 
-    def record_duration(self, operation: str, duration: float, labels: Optional[Dict[str, str]] = None):
+    def record_duration(
+        self,
+        operation: str,
+        duration: float,
+        labels: Optional[Dict[str, str]] = None,
+    ):
         """Record operation duration for metrics tracking"""
         if not self.enable_metrics:
             return