aiecs 1.5.1__py3-none-any.whl

aiecs/infrastructure/messaging/celery_task_manager.py

@@ -0,0 +1,383 @@
+import asyncio
+import logging
+import time
+import uuid
+from typing import Dict, List, Any
+from celery import Celery
+from celery.exceptions import TimeoutError as CeleryTimeoutError
+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
+
+logger = logging.getLogger(__name__)
+
+
+class CeleryTaskManager:
+    """
+    Specialized handler for Celery distributed task scheduling and execution
+    """
+
+    def __init__(self, config: Dict[str, Any]):
+        self.config = config
+        self.celery_app = 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"),
+            )
+
+            # 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},
+                ),
+            )
+
+            logger.info("Celery application initialized successfully")
+        except Exception as e:
+            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],
+    ):
+        """
+        Execute Celery task
+
+        Args:
+            task_name: Task name
+            queue: Queue name ('fast_tasks' or 'heavy_tasks')
+            user_id: User ID
+            task_id: Task ID
+            step: Step number
+            mode: Service mode
+            service: Service name
+            input_data: Input data
+            context: Context information
+
+        Returns:
+            Celery AsyncResult object
+        """
+        logger.info(
+            f"Queueing task {task_name} to {queue} for user {user_id}, task {task_id}, step {step}"
+        )
+
+        # Determine Celery task to use based on queue
+        celery_task_name = "aiecs.tasks.worker.execute_task"
+        if queue == "heavy_tasks":
+            celery_task_name = "aiecs.tasks.worker.execute_heavy_task"
+
+        # Send task to Celery
+        return self.celery_app.send_task(
+            celery_task_name,
+            kwargs={
+                "task_name": task_name,
+                "user_id": user_id,
+                "task_id": task_id,
+                "step": step,
+                "mode": mode,
+                "service": service,
+                "input_data": input_data,
+                "context": context,
+            },
+            queue=queue,
+        )
+
+    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
+        """
+        user_id = context.get("user_id", "anonymous")
+        task_id = input_data.get("task_id", str(uuid.uuid4()))
+        step = input_data.get("step", 0)
+        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)
+
+        try:
+            # Use string-based task names to avoid circular imports
+            celery_task_name = "aiecs.tasks.worker.execute_task"
+            if queue == "heavy_tasks":
+                celery_task_name = "aiecs.tasks.worker.execute_heavy_task"
+
+            result = self.celery_app.send_task(
+                celery_task_name,
+                kwargs={
+                    "task_name": task_name,
+                    "user_id": user_id,
+                    "task_id": task_id,
+                    "step": step,
+                    "mode": mode,
+                    "service": service,
+                    "input_data": input_data,
+                    "context": context,
+                },
+                queue=queue,
+            )
+
+            return result.get(timeout=timeout)
+
+        except CeleryTimeoutError as e:
+            logger.error(f"Timeout executing Celery task {task_name}: {e}")
+            return {
+                "status": TaskStatus.TIMED_OUT,
+                "error_code": ErrorCode.TIMEOUT_ERROR,
+                "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),
+            }
+
+    async def execute_heavy_task(self, task_name: str, input_data: Dict, context: Dict) -> Any:
+        """
+        Execute heavy task
+        """
+        input_data["queue"] = "heavy_tasks"
+        return await self.execute_task(task_name, input_data, context)
+
+    async def execute_dsl_task_step(
+        self, step: Dict, input_data: Dict, context: Dict
+    ) -> Dict[str, Any]:
+        """
+        Execute DSL task step
+        """
+        task_name = step.get("task")
+        category = "process"
+
+        if not task_name:
+            return {
+                "step": "unknown",
+                "result": None,
+                "completed": False,
+                "message": "Invalid DSL step: missing task name",
+                "status": TaskStatus.FAILED,
+                "error_code": ErrorCode.VALIDATION_ERROR,
+                "error_message": "Task name is required",
+            }
+
+        # Determine task type
+        task_type = "fast"
+        try:
+            task_type_result = await self.execute_task(task_name, {"get_task_type": True}, context)
+            if isinstance(task_type_result, dict) and "task_type" in task_type_result:
+                task_type = task_type_result["task_type"]
+        except Exception:
+            logger.warning(f"Could not determine task type for {task_name}, defaulting to 'fast'")
+
+        queue = "heavy_tasks" if task_type == "heavy" else "fast_tasks"
+        celery_task_name = (
+            "aiecs.tasks.worker.execute_heavy_task"
+            if task_type == "heavy"
+            else "aiecs.tasks.worker.execute_task"
+        )
+
+        user_id = context.get("user_id", str(uuid.uuid4()))
+        task_id = context.get("task_id", str(uuid.uuid4()))
+        step_num = context.get("step", 0)
+
+        # Send task to Celery
+        celery_task = self.celery_app.send_task(
+            celery_task_name,
+            kwargs={
+                "task_name": task_name,
+                "user_id": user_id,
+                "task_id": task_id,
+                "step": step_num,
+                "mode": context.get("mode", "multi_task"),
+                "service": context.get("service", "summarizer"),
+                "input_data": input_data,
+                "context": context,
+            },
+            queue=queue,
+        )
+
+        try:
+            timeout_seconds = self.config.get("task_timeout_seconds", 300)
+            start_time = time.time()
+
+            # Wait for task completion
+            while not celery_task.ready():
+                if time.time() - start_time > timeout_seconds:
+                    raise AsyncioTimeoutError(
+                        f"Task {task_name} timed out after {timeout_seconds} seconds"
+                    )
+                await asyncio.sleep(0.5)
+
+            if celery_task.successful():
+                result = celery_task.get()
+                if isinstance(result, dict) and "step" in result:
+                    return result
+                else:
+                    return {
+                        "step": f"{category}/{task_name}",
+                        "result": result,
+                        "completed": True,
+                        "message": f"Completed task {task_name}",
+                        "status": TaskStatus.COMPLETED,
+                    }
+            else:
+                error = celery_task.get(propagate=False)
+                status = (
+                    TaskStatus.TIMED_OUT
+                    if isinstance(error, CeleryTimeoutError)
+                    else TaskStatus.FAILED
+                )
+                error_code = (
+                    ErrorCode.TIMEOUT_ERROR
+                    if isinstance(error, CeleryTimeoutError)
+                    else ErrorCode.EXECUTION_ERROR
+                )
+
+                return {
+                    "step": f"{category}/{task_name}",
+                    "result": None,
+                    "completed": False,
+                    "message": f"Failed to execute task: {error}",
+                    "status": status,
+                    "error_code": error_code,
+                    "error_message": str(error),
+                }
+
+        except AsyncioTimeoutError as e:
+            return {
+                "step": f"{category}/{task_name}",
+                "result": None,
+                "completed": False,
+                "message": "Task execution timed out",
+                "status": TaskStatus.TIMED_OUT,
+                "error_code": ErrorCode.TIMEOUT_ERROR,
+                "error_message": str(e),
+            }
+        except Exception as e:
+            return {
+                "step": f"{category}/{task_name}",
+                "result": None,
+                "completed": False,
+                "message": f"Failed to execute {category}/{task_name}",
+                "status": TaskStatus.FAILED,
+                "error_code": ErrorCode.EXECUTION_ERROR,
+                "error_message": str(e),
+            }
+
+    def get_task_result(self, task_id: str):
+        """Get task result"""
+        try:
+            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,
+            }
+        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)}
+
+    def cancel_task(self, task_id: str):
+        """Cancel task"""
+        try:
+            self.celery_app.control.revoke(task_id, terminate=True)
+            logger.info(f"Task {task_id} cancelled")
+            return True
+        except Exception as e:
+            logger.error(f"Error cancelling task {task_id}: {e}")
+            return False
+
+    async def batch_execute_tasks(self, tasks: List[Dict[str, Any]]) -> List[Any]:
+        """
+        Batch execute tasks
+        """
+        results = []
+        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_results = await asyncio.gather(
+                *[
+                    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)
+
+        return results
+
+    def get_queue_info(self) -> Dict[str, Any]:
+        """Get queue information"""
+        try:
+            inspect = self.celery_app.control.inspect()
+            active_tasks = inspect.active()
+            scheduled_tasks = inspect.scheduled()
+            reserved_tasks = inspect.reserved()
+
+            return {
+                "active_tasks": active_tasks,
+                "scheduled_tasks": scheduled_tasks,
+                "reserved_tasks": reserved_tasks,
+            }
+        except Exception as e:
+            logger.error(f"Error getting queue info: {e}")
+            return {"error": str(e)}
+
+    def get_worker_stats(self) -> Dict[str, Any]:
+        """Get worker statistics"""
+        try:
+            inspect = self.celery_app.control.inspect()
+            stats = inspect.stats()
+            return stats or {}
+        except Exception as e:
+            logger.error(f"Error getting worker stats: {e}")
+            return {"error": str(e)}

aiecs/infrastructure/messaging/websocket_manager.py

@@ -0,0 +1,298 @@
+import asyncio
+import json
+import logging
+import uuid
+import websockets
+from typing import Dict, Any, Set, Optional, Callable
+from websockets import serve, ServerConnection
+from pydantic import BaseModel
+
+logger = logging.getLogger(__name__)
+
+
+class UserConfirmation(BaseModel):
+    proceed: bool
+    feedback: Optional[str] = None
+
+
+class TaskStepResult(BaseModel):
+    step: str
+    result: Any = None
+    completed: bool = False
+    message: str
+    status: str
+    error_code: Optional[str] = None
+    error_message: Optional[str] = None
+
+
+class WebSocketManager:
+    """
+    Specialized handler for WebSocket server and client communication
+    """
+
+    def __init__(self, host: str = "python-middleware-api", port: int = 8765):
+        self.host = host
+        self.port = port
+        self.server = None
+        self.callback_registry: Dict[str, Callable] = {}
+        self.active_connections: Set[ServerConnection] = set()
+        self._running = False
+
+    async def start_server(self):
+        """Start WebSocket server"""
+        if self.server:
+            logger.warning("WebSocket server is already running")
+            return self.server
+
+        try:
+            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
+        except Exception as e:
+            logger.error(f"Failed to start WebSocket server: {e}")
+            raise
+
+    async def stop_server(self):
+        """Stop WebSocket server"""
+        if self.server:
+            self.server.close()
+            await self.server.wait_closed()
+            self._running = False
+            logger.info("WebSocket server stopped")
+
+        # Close all active connections
+        if self.active_connections:
+            await asyncio.gather(
+                *[conn.close() for conn in self.active_connections],
+                return_exceptions=True,
+            )
+            self.active_connections.clear()
+
+    async def _handle_client_connection(self, websocket: ServerConnection, path: str):
+        """Handle client connection"""
+        self.active_connections.add(websocket)
+        client_addr = websocket.remote_address
+        logger.info(f"New WebSocket connection from {client_addr}")
+
+        try:
+            async for message in websocket:
+                await self._handle_client_message(websocket, message)
+        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:
+                await websocket.close()
+
+    async def _handle_client_message(self, websocket: ServerConnection, message: str):
+        """Handle client message"""
+        try:
+            data = json.loads(message)
+            action = data.get("action")
+
+            if action == "confirm":
+                await self._handle_confirmation(data)
+            elif action == "cancel":
+                await self._handle_cancellation(data)
+            elif action == "ping":
+                await self._handle_ping(websocket, data)
+            elif action == "subscribe":
+                await self._handle_subscription(websocket, data)
+            else:
+                logger.warning(f"Unknown action received: {action}")
+                await self._send_error(websocket, f"Unknown action: {action}")
+
+        except json.JSONDecodeError as e:
+            logger.error(f"Invalid JSON received: {e}")
+            await self._send_error(websocket, "Invalid JSON format")
+        except Exception as e:
+            logger.error(f"Error handling client message: {e}")
+            await self._send_error(websocket, f"Internal error: {str(e)}")
+
+    async def _handle_confirmation(self, data: Dict[str, Any]):
+        """Handle user confirmation"""
+        callback_id = data.get("callback_id")
+        if callback_id and callback_id in self.callback_registry:
+            callback = self.callback_registry[callback_id]
+            confirmation = UserConfirmation(
+                proceed=data.get("proceed", False),
+                feedback=data.get("feedback"),
+            )
+            try:
+                callback(confirmation)
+                del self.callback_registry[callback_id]
+                logger.debug(f"Processed confirmation for callback {callback_id}")
+            except Exception as e:
+                logger.error(f"Error processing confirmation callback: {e}")
+        else:
+            logger.warning(f"No callback found for confirmation ID: {callback_id}")
+
+    async def _handle_cancellation(self, data: Dict[str, Any]):
+        """Handle task cancellation"""
+        user_id = data.get("user_id")
+        task_id = data.get("task_id")
+
+        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
+            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(),
+                }
+            )
+        else:
+            logger.warning("Invalid cancellation request: missing user_id or task_id")
+
+    async def _handle_ping(self, websocket: ServerConnection, data: Dict[str, Any]):
+        """Handle heartbeat detection"""
+        pong_data = {
+            "type": "pong",
+            "timestamp": asyncio.get_event_loop().time(),
+            "original_data": data,
+        }
+        await self._send_to_client(websocket, pong_data)
+
+    async def _handle_subscription(self, websocket: ServerConnection, data: Dict[str, Any]):
+        """Handle subscription request"""
+        user_id = data.get("user_id")
+        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},
+            )
+
+    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(),
+        }
+        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))
+        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:
+        """Notify user of task step result"""
+        callback_id = str(uuid.uuid4())
+        confirmation_future = asyncio.Future()
+
+        # Register callback
+        self.callback_registry[callback_id] = lambda confirmation: confirmation_future.set_result(
+            confirmation
+        )
+
+        # Prepare notification data
+        notification_data = {
+            "type": "task_step_result",
+            "callback_id": callback_id,
+            "step": step,
+            "message": step_result.message,
+            "result": step_result.result,
+            "status": step_result.status,
+            "error_code": step_result.error_code,
+            "error_message": step_result.error_message,
+            "user_id": user_id,
+            "task_id": task_id,
+            "timestamp": asyncio.get_event_loop().time(),
+        }
+
+        try:
+            # 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:
+                # 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
+                self.callback_registry.pop(callback_id, None)
+                return UserConfirmation(proceed=True)  # Default to proceed
+
+        except Exception as e:
+            logger.error(f"WebSocket notification error: {e}")
+            # Clean up callback
+            self.callback_registry.pop(callback_id, None)
+            return UserConfirmation(proceed=True)  # Default to proceed
+
+    async def send_heartbeat(self, user_id: str, task_id: str, interval: int = 30):
+        """Send heartbeat message"""
+        heartbeat_data = {
+            "type": "heartbeat",
+            "status": "heartbeat",
+            "message": "Task is still executing...",
+            "user_id": user_id,
+            "task_id": task_id,
+            "timestamp": asyncio.get_event_loop().time(),
+        }
+
+        while self._running:
+            try:
+                await self.broadcast_message(heartbeat_data)
+                await asyncio.sleep(interval)
+            except Exception as e:
+                logger.error(f"WebSocket heartbeat error: {e}")
+                break
+
+    async def broadcast_message(self, message: Dict[str, Any]):
+        """Broadcast message to all connected clients"""
+        if not self.active_connections:
+            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]
+        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,
+            )
+            logger.debug(f"Broadcasted message to {len(active_connections)} clients")
+
+    async def send_to_user(self, user_id: str, message: Dict[str, Any]):
+        """Send message to specific user (requires user connection mapping implementation)"""
+        # User ID to WebSocket connection mapping can be implemented here
+        # Currently simplified to broadcast
+        message["target_user_id"] = user_id
+        await self.broadcast_message(message)
+
+    def get_connection_count(self) -> int:
+        """Get active connection count"""
+        return len([conn for conn in self.active_connections if not conn.closed])
+
+    def get_status(self) -> Dict[str, Any]:
+        """Get WebSocket manager status"""
+        return {
+            "running": self._running,
+            "host": self.host,
+            "port": self.port,
+            "active_connections": self.get_connection_count(),
+            "pending_callbacks": len(self.callback_registry),
+        }