aiecs 1.0.0__py3-none-any.whl

aiecs/domain/task/task_context.py

@@ -0,0 +1,257 @@
+import time
+import logging
+import json
+import os
+from typing import Dict, Any, Optional, AsyncGenerator, List, Tuple
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+import asyncio
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class ContextUpdate:
+    """Represents a single update to the context (e.g., message, metadata, or resource)."""
+    timestamp: float
+    update_type: str  # e.g., "message", "metadata", "resource"
+    data: Any  # Content of the update (e.g., message text, metadata dict)
+    metadata: Dict[str, Any]  # Additional metadata (e.g., file paths, model info)
+
+class TaskContext:
+    """
+    Enhanced context manager for task execution with:
+    - Context history tracking and checkpointing
+    - Resource acquisition and release
+    - Performance tracking
+    - File and model tracking
+    - Persistent storage
+    - Metadata toggles
+    - Enhanced error handling
+    """
+    def __init__(self, data: dict, task_dir: str = "./tasks"):
+        self.user_id = data.get("user_id", "anonymous")
+        self.chat_id = data.get("chat_id", "none")
+        # Ensure metadata includes aiPreference
+        self.metadata = data.get("metadata", {})
+        if "aiPreference" in data:
+            self.metadata["aiPreference"] = data["aiPreference"]
+        self.task_dir = Path(task_dir)
+        self.start_time: Optional[float] = None
+        self.resources: Dict[str, Any] = {}
+        self.context_history: List[ContextUpdate] = []
+        self.file_tracker: Dict[str, Dict[str, Any]] = {}  # Tracks file operations
+        self.model_tracker: List[Dict[str, Any]] = []  # Tracks model usage
+        self.metadata_toggles: Dict[str, bool] = data.get("metadata_toggles", {})
+        self._initialize_persistence()
+
+    def _initialize_persistence(self):
+        """Initialize persistent storage for context history."""
+        try:
+            self.task_dir.mkdir(parents=True, exist_ok=True)
+            history_file = self.task_dir / f"context_history_{self.chat_id}.json"
+            if history_file.exists():
+                with open(history_file, "r") as f:
+                    raw_history = json.load(f)
+                    self.context_history = [
+                        ContextUpdate(
+                            timestamp=entry["timestamp"],
+                            update_type=entry["update_type"],
+                            data=entry["data"],
+                            metadata=entry["metadata"]
+                        )
+                        for entry in raw_history
+                    ]
+                    logger.debug(f"Loaded context history from {history_file}")
+        except Exception as e:
+            logger.error(f"Failed to initialize context history: {e}")
+
+    async def _save_context_history(self):
+        """Save context history to disk."""
+        try:
+            history_file = self.task_dir / f"context_history_{self.chat_id}.json"
+            serialized_history = [
+                {
+                    "timestamp": update.timestamp,
+                    "update_type": update.update_type,
+                    "data": update.data,
+                    "metadata": update.metadata
+                }
+                for update in self.context_history
+            ]
+            with open(history_file, "w") as f:
+                json.dump(serialized_history, f, indent=2)
+            logger.debug(f"Saved context history to {history_file}")
+        except Exception as e:
+            logger.error(f"Failed to save context history: {e}")
+
+    def add_context_update(self, update_type: str, data: Any, metadata: Dict[str, Any] = None):
+        """Add a context update (e.g., message, metadata change)."""
+        update = ContextUpdate(
+            timestamp=time.time(),
+            update_type=update_type,
+            data=data,
+            metadata=metadata or {}
+        )
+        self.context_history.append(update)
+        logger.debug(f"Added context update: {update_type}")
+
+    def add_resource(self, name: str, resource: Any) -> None:
+        """Add a resource that needs cleanup."""
+        self.resources[name] = resource
+        self.add_context_update("resource", {"name": name}, {"type": type(resource).__name__})
+        logger.debug(f"Added resource: {name}")
+
+    def track_file_operation(self, file_path: str, operation: str, source: str = "task"):
+        """Track a file operation (e.g., read, edit)."""
+        self.file_tracker[file_path] = {
+            "operation": operation,
+            "source": source,
+            "timestamp": time.time(),
+            "state": "active"
+        }
+        self.add_context_update("file_operation", {"path": file_path, "operation": operation}, {"source": source})
+        logger.debug(f"Tracked file operation: {operation} on {file_path}")
+
+    def track_model_usage(self, model_id: str, provider_id: str, mode: str):
+        """Track AI model usage."""
+        model_entry = {
+            "model_id": model_id,
+            "provider_id": provider_id,
+            "mode": mode,
+            "timestamp": time.time()
+        }
+        # Avoid duplicates
+        if not self.model_tracker or self.model_tracker[-1] != model_entry:
+            self.model_tracker.append(model_entry)
+            self.add_context_update("model_usage", model_entry)
+            logger.debug(f"Tracked model usage: {model_id} ({provider_id}, {mode})")
+
+    def optimize_context(self, max_size: int = 1000) -> bool:
+        """Optimize context by removing duplicates and old entries."""
+        deduplicated = {}
+        optimized_history = []
+        total_size = 0
+
+        for update in reversed(self.context_history):
+            key = f"{update.update_type}:{json.dumps(update.data, sort_keys=True)}"
+            if key not in deduplicated:
+                deduplicated[key] = update
+                data_size = len(str(update.data))
+                if total_size + data_size <= max_size:
+                    optimized_history.append(update)
+                    total_size += data_size
+
+        self.context_history = list(reversed(optimized_history))
+        if len(deduplicated) < len(self.context_history):
+            logger.debug(f"Optimized context: removed {len(self.context_history) - len(deduplicated)} duplicates")
+            return True
+        return False
+
+    async def truncate_context_history(self, timestamp: float):
+        """Truncate context history after a given timestamp."""
+        original_len = len(self.context_history)
+        self.context_history = [update for update in self.context_history if update.timestamp <= timestamp]
+        if len(self.context_history) < original_len:
+            await self._save_context_history()
+            logger.debug(f"Truncated context history at timestamp {timestamp}")
+
+    def get_active_metadata(self) -> Dict[str, Any]:
+        """Return metadata filtered by toggles."""
+        return {
+            key: value
+            for key, value in self.metadata.items()
+            if key not in self.metadata_toggles or self.metadata_toggles[key] is not False
+        }
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert context to dictionary."""
+        return {
+            "user_id": self.user_id,
+            "chat_id": self.chat_id,
+            "metadata": self.get_active_metadata(),
+            "context_history": [
+                {
+                    "timestamp": update.timestamp,
+                    "update_type": update.update_type,
+                    "data": update.data,
+                    "metadata": update.metadata
+                }
+                for update in self.context_history
+            ],
+            "file_tracker": self.file_tracker,
+            "model_tracker": self.model_tracker
+        }
+
+    def __enter__(self):
+        """Synchronous context entry."""
+        self.start_time = time.time()
+        logger.debug(f"Starting task context for user {self.user_id}, chat {self.chat_id}")
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Synchronous context exit with cleanup."""
+        duration = time.time() - self.start_time
+        logger.debug(f"Completed task context in {duration:.2f}s for user {self.user_id}")
+        for resource_name, resource in self.resources.items():
+            try:
+                if hasattr(resource, 'close'):
+                    resource.close()
+                logger.debug(f"Cleaned up resource: {resource_name}")
+            except Exception as e:
+                logger.error(f"Error cleaning up resource {resource_name}: {e}")
+        if exc_type:
+            logger.error(f"Task context exited with error: {exc_val}")
+        return False
+
+    async def __aenter__(self):
+        """Asynchronous context entry."""
+        self.start_time = time.time()
+        logger.debug(f"Starting async task context for user {self.user_id}, chat {self.chat_id}")
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Asynchronous context exit with cleanup."""
+        duration = time.time() - self.start_time
+        logger.debug(f"Completed async task context in {duration:.2f}s for user {self.user_id}")
+        for resource_name, resource in self.resources.items():
+            try:
+                if hasattr(resource, 'close'):
+                    if callable(getattr(resource, 'close')):
+                        if hasattr(resource.close, '__await__'):
+                            await resource.close()
+                        else:
+                            resource.close()
+                logger.debug(f"Cleaned up async resource: {resource_name}")
+            except Exception as e:
+                logger.error(f"Error cleaning up async resource {resource_name}: {e}")
+        if exc_type:
+            logger.error(f"Async task context exited with error: {exc_val}")
+        await self._save_context_history()
+        return False
+
+
+def build_context(data: dict) -> dict:
+    """Build a simple context dictionary (for backward compatibility)."""
+    context = TaskContext(data)
+    return context.to_dict()
+
+
+@asynccontextmanager
+async def task_context(data: dict, task_dir: str = "./tasks") -> AsyncGenerator[TaskContext, None]:
+    """
+    Async context manager for task execution.
+
+    Usage:
+        async with task_context(request_data, task_dir="/path/to/tasks") as context:
+            context.add_context_update("message", "User input", {"source": "user"})
+            context.track_file_operation("example.py", "read", "tool")
+            result = await service_instance.run(data, context)
+    """
+    context = TaskContext(data, task_dir)
+    try:
+        await context.__aenter__()
+        yield context
+    finally:
+        await context.__aexit__(None, None, None)

aiecs/infrastructure/__init__.py

@@ -0,0 +1,26 @@
+"""Infrastructure layer module
+
+Contains external system integrations and technical concerns.
+"""
+
+from .messaging.celery_task_manager import CeleryTaskManager
+from .messaging.websocket_manager import WebSocketManager, UserConfirmation
+from .persistence.database_manager import DatabaseManager
+from .persistence.redis_client import RedisClient
+from .monitoring.executor_metrics import ExecutorMetrics
+from .monitoring.tracing_manager import TracingManager
+
+__all__ = [
+    # Messaging
+    "CeleryTaskManager",
+    "WebSocketManager",
+    "UserConfirmation",
+
+    # Persistence
+    "DatabaseManager",
+    "RedisClient",
+
+    # Monitoring
+    "ExecutorMetrics",
+    "TracingManager",
+]

aiecs/infrastructure/messaging/__init__.py

@@ -0,0 +1,13 @@
+"""Infrastructure messaging module
+
+Contains messaging and communication infrastructure.
+"""
+
+from .celery_task_manager import CeleryTaskManager
+from .websocket_manager import WebSocketManager, UserConfirmation
+
+__all__ = [
+    "CeleryTaskManager",
+    "WebSocketManager",
+    "UserConfirmation",
+]

aiecs/infrastructure/messaging/celery_task_manager.py

@@ -0,0 +1,341 @@
+import asyncio
+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 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)}