aiecs 1.0.0__py3-none-any.whl

aiecs/infrastructure/monitoring/tracing_manager.py

@@ -0,0 +1,376 @@
+import functools
+import logging
+import os
+from typing import Dict, Any, Optional
+import jaeger_client
+import jaeger_client.config
+from opentracing import tracer, Span
+
+logger = logging.getLogger(__name__)
+
+
+class TracingManager:
+    """
+    Specialized handler for distributed tracing and link tracking
+    """
+
+    def __init__(self, service_name: str = "service_executor",
+                 jaeger_host: Optional[str] = None,
+                 jaeger_port: Optional[int] = None,
+                 enable_tracing: Optional[bool] = None):
+        self.service_name = service_name
+        # Get configuration from environment variables, use defaults if not available
+        self.jaeger_host = jaeger_host or os.getenv("JAEGER_AGENT_HOST", "jaeger")
+        self.jaeger_port = jaeger_port or int(os.getenv("JAEGER_AGENT_PORT", "6831"))
+        self.enable_tracing = enable_tracing if enable_tracing is not None else os.getenv("JAEGER_ENABLE_TRACING", "true").lower() == "true"
+        self.tracer = None
+
+        if self.enable_tracing:
+            self._init_tracer()
+
+    def _init_tracer(self):
+        """Initialize Jaeger tracer"""
+        try:
+            config = jaeger_client.config.Config(
+                config={
+                    'sampler': {
+                        'type': 'const',
+                        'param': 1,
+                    },
+                    'local_agent': {
+                        'reporting_host': self.jaeger_host,
+                        'reporting_port': self.jaeger_port,
+                    },
+                    'logging': True,
+                },
+                service_name=self.service_name,
+                validate=True
+            )
+            self.tracer = config.initialize_tracer()
+            logger.info(f"Jaeger tracer initialized for service '{self.service_name}' at {self.jaeger_host}:{self.jaeger_port}")
+        except Exception as e:
+            logger.warning(f"Failed to initialize Jaeger tracer: {e}")
+            self.tracer = None
+            self.enable_tracing = False
+
+    def start_span(self, operation_name: str, parent_span: Optional[Span] = None,
+                   tags: Optional[Dict[str, Any]] = None) -> Optional[Span]:
+        """
+        Start a tracing span
+
+        Args:
+            operation_name: Operation name
+            parent_span: Parent span
+            tags: Initial tags
+
+        Returns:
+            Span object or None (if tracing is not enabled)
+        """
+        if not self.enable_tracing or not self.tracer:
+            return None
+
+        try:
+            span = self.tracer.start_span(
+                operation_name=operation_name,
+                child_of=parent_span
+            )
+
+            # Set initial tags
+            if tags:
+                for key, value in tags.items():
+                    span.set_tag(key, value)
+
+            # Set service information
+            span.set_tag("service.name", self.service_name)
+            span.set_tag("span.kind", "server")
+
+            return span
+        except Exception as e:
+            logger.error(f"Error starting span '{operation_name}': {e}")
+            return None
+
+    def finish_span(self, span: Optional[Span], tags: Optional[Dict[str, Any]] = None,
+                    logs: Optional[Dict[str, Any]] = None, error: Optional[Exception] = None):
+        """
+        Finish tracing span
+
+        Args:
+            span: Span to finish
+            tags: Additional tags
+            logs: Log information
+            error: Error information
+        """
+        if not span or not self.enable_tracing:
+            return
+
+        try:
+            # Add additional tags
+            if tags:
+                for key, value in tags.items():
+                    span.set_tag(key, value)
+
+            # Record error
+            if error:
+                span.set_tag("error", True)
+                span.set_tag("error.kind", type(error).__name__)
+                span.set_tag("error.message", str(error))
+                span.log_kv({"event": "error", "error.object": error})
+
+            # Add logs
+            if logs:
+                span.log_kv(logs)
+
+            span.finish()
+        except Exception as e:
+            logger.error(f"Error finishing span: {e}")
+
+    def with_tracing(self, operation_name: str, tags: Optional[Dict[str, Any]] = None):
+        """
+        Tracing decorator
+
+        Args:
+            operation_name: Operation name
+            tags: Initial tags
+        """
+        def decorator(func):
+            @functools.wraps(func)
+            async def async_wrapper(*args, **kwargs):
+                if not self.enable_tracing or not self.tracer:
+                    return await func(*args, **kwargs)
+
+                span = self.start_span(operation_name, tags=tags)
+
+                try:
+                    # Add function arguments as tags
+                    self._add_function_args_to_span(span, args, kwargs)
+
+                    result = await func(*args, **kwargs)
+
+                    # Record success
+                    if span:
+                        span.set_tag("success", True)
+
+                    return result
+                except Exception as e:
+                    self.finish_span(span, error=e)
+                    raise
+                finally:
+                    if span and not span.finished:
+                        self.finish_span(span)
+
+            @functools.wraps(func)
+            def sync_wrapper(*args, **kwargs):
+                if not self.enable_tracing or not self.tracer:
+                    return func(*args, **kwargs)
+
+                span = self.start_span(operation_name, tags=tags)
+
+                try:
+                    # Add function arguments as tags
+                    self._add_function_args_to_span(span, args, kwargs)
+
+                    result = func(*args, **kwargs)
+
+                    # Record success
+                    if span:
+                        span.set_tag("success", True)
+
+                    return result
+                except Exception as e:
+                    self.finish_span(span, error=e)
+                    raise
+                finally:
+                    if span and not span.finished:
+                        self.finish_span(span)
+
+            # Return appropriate wrapper based on function type
+            import asyncio
+            if asyncio.iscoroutinefunction(func):
+                return async_wrapper
+            else:
+                return sync_wrapper
+
+        return decorator
+
+    def _add_function_args_to_span(self, span: Optional[Span], args: tuple, kwargs: Dict[str, Any]):
+        """Add function arguments to span tags"""
+        if not span:
+            return
+
+        try:
+            # Add positional arguments
+            for i, arg in enumerate(args):
+                if isinstance(arg, (str, int, float, bool)):
+                    span.set_tag(f"arg_{i}", arg)
+                elif hasattr(arg, '__class__'):
+                    span.set_tag(f"arg_{i}_type", arg.__class__.__name__)
+
+            # Add keyword arguments
+            for key, value in kwargs.items():
+                if isinstance(value, (str, int, float, bool)):
+                    span.set_tag(key, value)
+                elif isinstance(value, dict) and len(str(value)) < 1000:  # Avoid overly large dictionaries
+                    span.set_tag(f"{key}_json", str(value))
+                elif hasattr(value, '__class__'):
+                    span.set_tag(f"{key}_type", value.__class__.__name__)
+        except Exception as e:
+            logger.debug(f"Error adding function args to span: {e}")
+
+    def trace_database_operation(self, operation: str, table: str = None, query: str = None):
+        """Database operation tracing decorator"""
+        def decorator(func):
+            @functools.wraps(func)
+            async def wrapper(*args, **kwargs):
+                tags = {
+                    "component": "database",
+                    "db.type": "postgresql",
+                    "db.statement.type": operation
+                }
+
+                if table:
+                    tags["db.table"] = table
+                if query:
+                    tags["db.statement"] = query[:500]  # Limit query length
+
+                span = self.start_span(f"db.{operation}", tags=tags)
+
+                try:
+                    result = await func(*args, **kwargs)
+                    if span:
+                        span.set_tag("db.rows_affected", len(result) if isinstance(result, list) else 1)
+                    return result
+                except Exception as e:
+                    self.finish_span(span, error=e)
+                    raise
+                finally:
+                    if span and not span.finished:
+                        self.finish_span(span)
+
+            return wrapper
+        return decorator
+
+    def trace_external_call(self, service_name: str, endpoint: str = None):
+        """External service call tracing decorator"""
+        def decorator(func):
+            @functools.wraps(func)
+            async def wrapper(*args, **kwargs):
+                tags = {
+                    "component": "http",
+                    "span.kind": "client",
+                    "peer.service": service_name
+                }
+
+                if endpoint:
+                    tags["http.url"] = endpoint
+
+                span = self.start_span(f"http.{service_name}", tags=tags)
+
+                try:
+                    result = await func(*args, **kwargs)
+                    if span:
+                        span.set_tag("http.status_code", 200)
+                    return result
+                except Exception as e:
+                    if span:
+                        span.set_tag("http.status_code", 500)
+                    self.finish_span(span, error=e)
+                    raise
+                finally:
+                    if span and not span.finished:
+                        self.finish_span(span)
+
+            return wrapper
+        return decorator
+
+    def trace_tool_execution(self, tool_name: str, operation: str):
+        """Tool execution tracing decorator"""
+        def decorator(func):
+            @functools.wraps(func)
+            async def wrapper(*args, **kwargs):
+                tags = {
+                    "component": "tool",
+                    "tool.name": tool_name,
+                    "tool.operation": operation
+                }
+
+                span = self.start_span(f"tool.{tool_name}.{operation}", tags=tags)
+
+                try:
+                    result = await func(*args, **kwargs)
+                    if span:
+                        span.set_tag("tool.success", True)
+                        if hasattr(result, '__len__'):
+                            span.set_tag("tool.result_size", len(result))
+                    return result
+                except Exception as e:
+                    if span:
+                        span.set_tag("tool.success", False)
+                    self.finish_span(span, error=e)
+                    raise
+                finally:
+                    if span and not span.finished:
+                        self.finish_span(span)
+
+            return wrapper
+        return decorator
+
+    def create_child_span(self, parent_span: Optional[Span], operation_name: str,
+                         tags: Optional[Dict[str, Any]] = None) -> Optional[Span]:
+        """Create child span"""
+        if not self.enable_tracing or not parent_span:
+            return None
+
+        return self.start_span(operation_name, parent_span=parent_span, tags=tags)
+
+    def inject_span_context(self, span: Optional[Span], carrier: Dict[str, str]):
+        """Inject span context into carrier (for cross-service propagation)"""
+        if not self.enable_tracing or not span or not self.tracer:
+            return
+
+        try:
+            from opentracing.propagation import Format
+            self.tracer.inject(span.context, Format.TEXT_MAP, carrier)
+        except Exception as e:
+            logger.error(f"Error injecting span context: {e}")
+
+    def extract_span_context(self, carrier: Dict[str, str]) -> Optional[Any]:
+        """Extract span context from carrier"""
+        if not self.enable_tracing or not self.tracer:
+            return None
+
+        try:
+            from opentracing.propagation import Format
+            return self.tracer.extract(Format.TEXT_MAP, carrier)
+        except Exception as e:
+            logger.error(f"Error extracting span context: {e}")
+            return None
+
+    def get_active_span(self) -> Optional[Span]:
+        """Get current active span"""
+        if not self.enable_tracing or not self.tracer:
+            return None
+
+        try:
+            return self.tracer.active_span
+        except Exception as e:
+            logger.error(f"Error getting active span: {e}")
+            return None
+
+    def close_tracer(self):
+        """Close tracer"""
+        if self.tracer:
+            try:
+                self.tracer.close()
+                logger.info("Tracer closed successfully")
+            except Exception as e:
+                logger.error(f"Error closing tracer: {e}")
+
+    def get_tracer_info(self) -> Dict[str, Any]:
+        """Get tracer information"""
+        return {
+            "enabled": self.enable_tracing,
+            "service_name": self.service_name,
+            "jaeger_host": self.jaeger_host,
+            "jaeger_port": self.jaeger_port,
+            "tracer_initialized": self.tracer is not None
+        }

aiecs/infrastructure/persistence/__init__.py

@@ -0,0 +1,12 @@
+"""Infrastructure persistence module
+
+Contains data persistence and storage infrastructure.
+"""
+
+from .database_manager import DatabaseManager
+from .redis_client import RedisClient
+
+__all__ = [
+    "DatabaseManager",
+    "RedisClient",
+]

aiecs/infrastructure/persistence/database_manager.py

@@ -0,0 +1,286 @@
+import asyncpg
+import json
+import logging
+from datetime import datetime
+from typing import Dict, List, Any, Optional
+from aiecs.domain.execution.model import TaskStatus, TaskStepResult
+from aiecs.config.config import get_settings
+
+logger = logging.getLogger(__name__)
+
+
+class DatabaseManager:
+    """
+    Specialized handler for database connections, operations, and task history management
+    """
+
+    def __init__(self, db_config: Optional[Dict[str, Any]] = None):
+        if db_config is None:
+            settings = get_settings()
+            self.db_config = settings.database_config
+        else:
+            self.db_config = db_config
+        self.connection_pool = None
+        self._initialized = False
+
+    async def init_connection_pool(self, min_size: int = 10, max_size: int = 20):
+        """Initialize database connection pool"""
+        try:
+            self.connection_pool = await asyncpg.create_pool(
+                **self.db_config,
+                min_size=min_size,
+                max_size=max_size
+            )
+            logger.info("Database connection pool initialized successfully")
+        except Exception as e:
+            logger.error(f"Failed to initialize database connection pool: {e}")
+            raise
+
+    async def _get_connection(self):
+        """Get database connection"""
+        if self.connection_pool:
+            return self.connection_pool.acquire()
+        else:
+            return asyncpg.connect(**self.db_config)
+
+    async def init_database_schema(self):
+        """Initialize database table structure"""
+        try:
+            if self.connection_pool:
+                async with self.connection_pool.acquire() as conn:
+                    await self._create_tables(conn)
+            else:
+                conn = await asyncpg.connect(**self.db_config)
+                try:
+                    await self._create_tables(conn)
+                finally:
+                    await conn.close()
+
+            self._initialized = True
+            logger.info("Database schema initialized successfully")
+            return True
+        except Exception as e:
+            logger.error(f"Database initialization error: {e}")
+            return False
+
+    async def _create_tables(self, conn):
+        """Create database tables"""
+        await conn.execute('''
+            CREATE TABLE IF NOT EXISTS task_history (
+                id SERIAL PRIMARY KEY,
+                user_id TEXT NOT NULL,
+                task_id TEXT NOT NULL,
+                step INTEGER NOT NULL,
+                result JSONB NOT NULL,
+                timestamp TIMESTAMP NOT NULL,
+                status TEXT NOT NULL DEFAULT 'pending'
+            );
+            CREATE INDEX IF NOT EXISTS idx_task_history_user_id ON task_history (user_id);
+            CREATE INDEX IF NOT EXISTS idx_task_history_task_id ON task_history (task_id);
+            CREATE INDEX IF NOT EXISTS idx_task_history_status ON task_history (status);
+            CREATE INDEX IF NOT EXISTS idx_task_history_timestamp ON task_history (timestamp);
+        ''')
+
+    async def save_task_history(self, user_id: str, task_id: str, step: int, step_result: TaskStepResult):
+        """Save task execution history"""
+        if not self._initialized:
+            await self.init_database_schema()
+
+        try:
+            if self.connection_pool:
+                async with self.connection_pool.acquire() as conn:
+                    await conn.execute(
+                        'INSERT INTO task_history (user_id, task_id, step, result, timestamp, status) VALUES ($1, $2, $3, $4, $5, $6)',
+                        user_id, task_id, step, json.dumps(step_result.dict()), datetime.now(), step_result.status
+                    )
+            else:
+                conn = await asyncpg.connect(**self.db_config)
+                try:
+                    await conn.execute(
+                        'INSERT INTO task_history (user_id, task_id, step, result, timestamp, status) VALUES ($1, $2, $3, $4, $5, $6)',
+                        user_id, task_id, step, json.dumps(step_result.dict()), datetime.now(), step_result.status
+                    )
+                finally:
+                    await conn.close()
+
+            logger.debug(f"Saved task history for user {user_id}, task {task_id}, step {step}")
+            return True
+        except Exception as e:
+            logger.error(f"Database error saving task history: {e}")
+            raise Exception(f"Database error: {e}")
+
+    async def load_task_history(self, user_id: str, task_id: str) -> List[Dict]:
+        """Load task execution history"""
+        if not self._initialized:
+            await self.init_database_schema()
+
+        try:
+            if self.connection_pool:
+                async with self.connection_pool.acquire() as conn:
+                    records = await conn.fetch(
+                        'SELECT step, result, timestamp, status FROM task_history WHERE user_id = $1 AND task_id = $2 ORDER BY step ASC',
+                        user_id, task_id
+                    )
+            else:
+                conn = await asyncpg.connect(**self.db_config)
+                try:
+                    records = await conn.fetch(
+                        'SELECT step, result, timestamp, status FROM task_history WHERE user_id = $1 AND task_id = $2 ORDER BY step ASC',
+                        user_id, task_id
+                    )
+                finally:
+                    await conn.close()
+
+            return [
+                {
+                    "step": r['step'],
+                    "result": json.loads(r['result']),
+                    "timestamp": r['timestamp'].isoformat(),
+                    "status": r['status']
+                }
+                for r in records
+            ]
+        except Exception as e:
+            logger.error(f"Database error loading task history: {e}")
+            raise Exception(f"Database error: {e}")
+
+    async def mark_task_as_cancelled(self, user_id: str, task_id: str):
+        """Mark task as cancelled"""
+        if not self._initialized:
+            await self.init_database_schema()
+
+        try:
+            if self.connection_pool:
+                async with self.connection_pool.acquire() as conn:
+                    await conn.execute(
+                        'UPDATE task_history SET status = $1 WHERE user_id = $2 AND task_id = $3',
+                        TaskStatus.CANCELLED.value, user_id, task_id
+                    )
+            else:
+                conn = await asyncpg.connect(**self.db_config)
+                try:
+                    await conn.execute(
+                        'UPDATE task_history SET status = $1 WHERE user_id = $2 AND task_id = $3',
+                        TaskStatus.CANCELLED.value, user_id, task_id
+                    )
+                finally:
+                    await conn.close()
+
+            logger.info(f"Marked task {task_id} as cancelled for user {user_id}")
+            return True
+        except Exception as e:
+            logger.error(f"Database error marking task as cancelled: {e}")
+            raise Exception(f"Database error: {e}")
+
+    async def check_task_status(self, user_id: str, task_id: str) -> TaskStatus:
+        """Check task status"""
+        if not self._initialized:
+            await self.init_database_schema()
+
+        try:
+            if self.connection_pool:
+                async with self.connection_pool.acquire() as conn:
+                    record = await conn.fetchrow(
+                        'SELECT status FROM task_history WHERE user_id = $1 AND task_id = $2 ORDER BY step DESC LIMIT 1',
+                        user_id, task_id
+                    )
+            else:
+                conn = await asyncpg.connect(**self.db_config)
+                try:
+                    record = await conn.fetchrow(
+                        'SELECT status FROM task_history WHERE user_id = $1 AND task_id = $2 ORDER BY step DESC LIMIT 1',
+                        user_id, task_id
+                    )
+                finally:
+                    await conn.close()
+
+            return TaskStatus(record['status']) if record else TaskStatus.PENDING
+        except Exception as e:
+            logger.error(f"Database error checking task status: {e}")
+            raise Exception(f"Database error: {e}")
+
+    async def get_user_tasks(self, user_id: str, limit: int = 100) -> List[Dict]:
+        """Get user task list"""
+        if not self._initialized:
+            await self.init_database_schema()
+
+        try:
+            if self.connection_pool:
+                async with self.connection_pool.acquire() as conn:
+                    records = await conn.fetch(
+                        '''SELECT DISTINCT task_id,
+                           MAX(timestamp) as last_updated,
+                           (SELECT status FROM task_history th2
+                            WHERE th2.user_id = $1 AND th2.task_id = th1.task_id
+                            ORDER BY step DESC LIMIT 1) as status
+                           FROM task_history th1
+                           WHERE user_id = $1
+                           GROUP BY task_id
+                           ORDER BY last_updated DESC
+                           LIMIT $2''',
+                        user_id, limit
+                    )
+            else:
+                conn = await asyncpg.connect(**self.db_config)
+                try:
+                    records = await conn.fetch(
+                        '''SELECT DISTINCT task_id,
+                           MAX(timestamp) as last_updated,
+                           (SELECT status FROM task_history th2
+                            WHERE th2.user_id = $1 AND th2.task_id = th1.task_id
+                            ORDER BY step DESC LIMIT 1) as status
+                           FROM task_history th1
+                           WHERE user_id = $1
+                           GROUP BY task_id
+                           ORDER BY last_updated DESC
+                           LIMIT $2''',
+                        user_id, limit
+                    )
+                finally:
+                    await conn.close()
+
+            return [
+                {
+                    "task_id": r['task_id'],
+                    "last_updated": r['last_updated'].isoformat(),
+                    "status": r['status']
+                }
+                for r in records
+            ]
+        except Exception as e:
+            logger.error(f"Database error getting user tasks: {e}")
+            raise Exception(f"Database error: {e}")
+
+    async def cleanup_old_tasks(self, days_old: int = 30):
+        """Clean up old task records"""
+        if not self._initialized:
+            await self.init_database_schema()
+
+        try:
+            if self.connection_pool:
+                async with self.connection_pool.acquire() as conn:
+                    result = await conn.execute(
+                        'DELETE FROM task_history WHERE timestamp < NOW() - INTERVAL %s DAY',
+                        days_old
+                    )
+            else:
+                conn = await asyncpg.connect(**self.db_config)
+                try:
+                    result = await conn.execute(
+                        'DELETE FROM task_history WHERE timestamp < NOW() - INTERVAL %s DAY',
+                        days_old
+                    )
+                finally:
+                    await conn.close()
+
+            logger.info(f"Cleaned up old task records: {result}")
+            return True
+        except Exception as e:
+            logger.error(f"Database error during cleanup: {e}")
+            return False
+
+    async def close(self):
+        """Close database connection pool"""
+        if self.connection_pool:
+            await self.connection_pool.close()
+            logger.info("Database connection pool closed")