aiecs 1.0.0__py3-none-any.whl

aiecs/infrastructure/messaging/websocket_manager.py

@@ -0,0 +1,289 @@
+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:
+                return await asyncio.wait_for(confirmation_future, timeout=300)  # 5 minute timeout
+            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)
+        }

aiecs/infrastructure/monitoring/__init__.py

@@ -0,0 +1,12 @@
+"""Infrastructure monitoring module
+
+Contains monitoring, metrics, and observability infrastructure.
+"""
+
+from .executor_metrics import ExecutorMetrics
+from .tracing_manager import TracingManager
+
+__all__ = [
+    "ExecutorMetrics",
+    "TracingManager",
+]

aiecs/infrastructure/monitoring/executor_metrics.py

@@ -0,0 +1,138 @@
+import logging
+import functools
+from typing import Dict, Optional, Any
+from prometheus_client import Counter, Histogram, start_http_server
+
+logger = logging.getLogger(__name__)
+
+
+class ExecutorMetrics:
+    """
+    Specialized handler for executor performance monitoring and metrics collection
+    """
+
+    def __init__(self, enable_metrics: bool = True, metrics_port: int = 8001):
+        self.enable_metrics = enable_metrics
+        self.metrics_port = metrics_port
+        self.metrics: Dict[str, Any] = {}
+
+        if self.enable_metrics:
+            self._init_prometheus_metrics()
+
+    def _init_prometheus_metrics(self):
+        """Initialize Prometheus metrics"""
+        try:
+            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_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"]),
+            }
+            logger.info(f"Prometheus metrics server started on port {self.metrics_port}")
+        except Exception as e:
+            logger.warning(f"Failed to start metrics server: {e}")
+            self.metrics = {}
+
+    def record_operation_latency(self, operation: str, duration: float):
+        """Record operation latency"""
+        if not self.enable_metrics or f"{operation}_latency" not in self.metrics:
+            return
+        self.metrics[f"{operation}_latency"].observe(duration)
+
+    def record_operation_success(self, operation: str, labels: Optional[Dict[str, str]] = None):
+        """Record operation success"""
+        if not self.enable_metrics or f"{operation}_success" not in self.metrics:
+            return
+        metric = self.metrics[f"{operation}_success"]
+        if labels:
+            metric = metric.labels(**labels)
+        metric.inc()
+
+    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
+        # Failure metrics can be added
+        logger.error(f"Operation {operation} failed with error type: {error_type}")
+
+    def record_retry(self, operation: str, attempt_number: int):
+        """Record retry"""
+        if not self.enable_metrics or f"{operation}_retries" not in self.metrics:
+            return
+        if attempt_number > 1:
+            self.metrics[f"{operation}_retries"].inc()
+
+    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):
+                if not self.metrics or f"{metric_name}_latency" not in self.metrics:
+                    return await func(*args, **kwargs)
+
+                labels_dict = labels or {}
+                metric = self.metrics[f"{metric_name}_latency"]
+                if labels:
+                    metric = metric.labels(**labels_dict)
+
+                with metric.time():
+                    try:
+                        result = await func(*args, **kwargs)
+                        if f"{metric_name}_success" in self.metrics:
+                            success_metric = self.metrics[f"{metric_name}_success"]
+                            if labels:
+                                success_metric = success_metric.labels(**labels_dict)
+                            success_metric.inc()
+                        return result
+                    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]:
+        """Get metrics summary"""
+        if not self.enable_metrics:
+            return {"metrics_enabled": False}
+
+        return {
+            "metrics_enabled": True,
+            "metrics_port": self.metrics_port,
+            "available_metrics": list(self.metrics.keys())
+        }
+
+    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
+
+        try:
+            # Record operation success/failure
+            if success:
+                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'))
+
+            # Record operation latency if provided
+            if duration is not None:
+                self.record_operation_latency(operation_type, duration)
+
+        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):
+        """Record operation duration for metrics tracking"""
+        if not self.enable_metrics:
+            return
+
+        try:
+            self.record_operation_latency(operation, duration)
+        except Exception as e:
+            logger.warning(f"Failed to record duration metrics: {e}")

aiecs/infrastructure/monitoring/structured_logger.py

@@ -0,0 +1,50 @@
+"""Structured logging setup for aiecs."""
+import logging
+import sys
+from typing import Optional
+
+
+def setup_structured_logging(level: str = "INFO", format_type: str = "json") -> None:
+    """
+    Setup structured logging for the application.
+
+    Args:
+        level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
+        format_type: Format type (json, text)
+    """
+    # Convert string level to logging level
+    numeric_level = getattr(logging, level.upper(), logging.INFO)
+
+    # Create formatter
+    if format_type.lower() == "json":
+        # Simple JSON-like format for now
+        formatter = logging.Formatter(
+            '{"timestamp": "%(asctime)s", "level": "%(levelname)s", "module": "%(name)s", "message": "%(message)s"}'
+        )
+    else:
+        # Standard text format
+        formatter = logging.Formatter(
+            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+        )
+
+    # Setup root logger
+    root_logger = logging.getLogger()
+    root_logger.setLevel(numeric_level)
+
+    # Remove existing handlers
+    for handler in root_logger.handlers[:]:
+        root_logger.removeHandler(handler)
+
+    # Add console handler
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setLevel(numeric_level)
+    console_handler.setFormatter(formatter)
+    root_logger.addHandler(console_handler)
+
+    # Set specific logger levels
+    logging.getLogger("aiecs").setLevel(numeric_level)
+
+    # Suppress noisy third-party loggers
+    logging.getLogger("urllib3").setLevel(logging.WARNING)
+    logging.getLogger("requests").setLevel(logging.WARNING)
+    logging.getLogger("httpx").setLevel(logging.WARNING)