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

aiecs/infrastructure/persistence/context_engine_client.py

@@ -0,0 +1,184 @@
+"""
+Global ContextEngine Manager
+
+This module provides a singleton ContextEngine instance that can be shared
+across all components in the application. It follows the same pattern as
+the Redis client initialization in aiecs.infrastructure.persistence.redis_client.
+
+Usage:
+    # In main.py startup:
+    await initialize_context_engine()
+
+    # In any component:
+    from aiecs.infrastructure.persistence.context_engine_client import get_context_engine
+    context_engine = get_context_engine()
+"""
+
+import logging
+from typing import Optional, TYPE_CHECKING
+import asyncio
+
+if TYPE_CHECKING:
+    from aiecs.domain.context.context_engine import ContextEngine
+
+logger = logging.getLogger(__name__)
+
+# Global singleton instance
+_global_context_engine: Optional["ContextEngine"] = None
+_initialization_lock = asyncio.Lock()
+_initialized = False
+
+
+def _get_context_engine_class():
+    """Lazy import of ContextEngine to avoid circular dependencies."""
+    try:
+        from aiecs.domain.context.context_engine import ContextEngine
+
+        return ContextEngine
+    except ImportError as e:
+        logger.warning(f"ContextEngine not available - {e}")
+        return None
+
+
+async def initialize_context_engine(
+    use_existing_redis: bool = True,
+) -> Optional["ContextEngine"]:
+    """
+    Initialize the global ContextEngine instance.
+
+    This should be called once during application startup (in main.py lifespan).
+
+    Args:
+        use_existing_redis: Whether to use existing Redis client (default: True)
+
+    Returns:
+        The initialized ContextEngine instance or None if initialization fails
+
+    Example:
+        @asynccontextmanager
+        async def lifespan(app: FastAPI):
+            # Startup
+            await initialize_redis_client()
+            await initialize_context_engine()  # Initialize after Redis
+            yield
+            # Shutdown
+            await close_context_engine()
+            await close_redis_client()
+    """
+    global _global_context_engine, _initialized
+
+    if _initialized and _global_context_engine:
+        logger.info("ContextEngine already initialized")
+        return _global_context_engine
+
+    async with _initialization_lock:
+        # Double-check after acquiring lock
+        if _initialized and _global_context_engine:
+            return _global_context_engine
+
+        ContextEngine = _get_context_engine_class()
+        if not ContextEngine:
+            logger.error("ContextEngine class not available - cannot initialize")
+            return None
+
+        try:
+            logger.info("Initializing global ContextEngine...")
+            _global_context_engine = ContextEngine(use_existing_redis=use_existing_redis)
+            await _global_context_engine.initialize()
+            _initialized = True
+            logger.info("✅ Global ContextEngine initialized successfully")
+            return _global_context_engine
+
+        except Exception as e:
+            logger.error(f"❌ Failed to initialize global ContextEngine: {e}")
+            logger.warning("Application will continue without ContextEngine (degraded mode)")
+            _global_context_engine = None
+            _initialized = False
+            return None
+
+
+def get_context_engine() -> Optional["ContextEngine"]:
+    """
+    Get the global ContextEngine instance.
+
+    Returns:
+        The global ContextEngine instance or None if not initialized
+
+    Example:
+        from aiecs.infrastructure.persistence.context_engine_client import get_context_engine
+
+        context_engine = get_context_engine()
+        if context_engine:
+            await context_engine.add_conversation_message(...)
+        else:
+            # Fallback to local storage
+            logger.warning("ContextEngine not available")
+    """
+    if not _initialized:
+        logger.warning("ContextEngine not initialized. Call initialize_context_engine() " "during application startup.")
+    return _global_context_engine
+
+
+async def close_context_engine() -> None:
+    """
+    Close and cleanup the global ContextEngine instance.
+
+    This should be called during application shutdown (in main.py lifespan).
+
+    Example:
+        @asynccontextmanager
+        async def lifespan(app: FastAPI):
+            # Startup
+            await initialize_context_engine()
+            yield
+            # Shutdown
+            await close_context_engine()
+    """
+    global _global_context_engine, _initialized
+
+    async with _initialization_lock:
+        if _global_context_engine:
+            try:
+                logger.info("Closing global ContextEngine...")
+                # ContextEngine cleanup if needed
+                if hasattr(_global_context_engine, "close"):
+                    await _global_context_engine.close()
+                logger.info("✅ Global ContextEngine closed successfully")
+            except Exception as e:
+                logger.error(f"Error closing ContextEngine: {e}")
+            finally:
+                _global_context_engine = None
+                _initialized = False
+
+
+def is_context_engine_initialized() -> bool:
+    """
+    Check if the global ContextEngine is initialized.
+
+    Returns:
+        True if ContextEngine is initialized and available, False otherwise
+    """
+    return _initialized and _global_context_engine is not None
+
+
+# Convenience function for testing
+async def reset_context_engine() -> None:
+    """
+    Reset the global ContextEngine instance.
+
+    This is primarily for testing purposes to allow re-initialization.
+    Should NOT be used in production code.
+    """
+    global _global_context_engine, _initialized
+
+    async with _initialization_lock:
+        if _global_context_engine:
+            try:
+                if hasattr(_global_context_engine, "close"):
+                    await _global_context_engine.close()
+            except Exception as e:
+                logger.warning(f"Error during ContextEngine reset: {e}")
+
+        _global_context_engine = None
+        _initialized = False
+        logger.info("ContextEngine reset completed")

aiecs/infrastructure/persistence/database_manager.py

@@ -1,4 +1,4 @@
-import asyncpg
+import asyncpg  # type: ignore[import-untyped]
 import json
 import logging
 from datetime import datetime
@@ -41,11 +41,7 @@ class DatabaseManager:
     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
-            )
+            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}")
@@ -80,7 +76,8 @@ class DatabaseManager:
 
     async def _create_tables(self, conn):
         """Create database tables"""
-        await conn.execute('''
+        await conn.execute(
+            """
             CREATE TABLE IF NOT EXISTS task_history (
                 id SERIAL PRIMARY KEY,
                 user_id TEXT NOT NULL,
@@ -94,9 +91,16 @@ class DatabaseManager:
             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):
+    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()
@@ -105,15 +109,25 @@ class DatabaseManager:
             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
+                        "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
+                        "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()
@@ -133,25 +147,27 @@ class DatabaseManager:
             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
+                        "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
+                        "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']
+                    "step": r["step"],
+                    "result": json.loads(r["result"]),
+                    "timestamp": r["timestamp"].isoformat(),
+                    "status": r["status"],
                 }
                 for r in records
             ]
@@ -168,15 +184,19 @@ class DatabaseManager:
             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
+                        "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
+                        "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()
@@ -196,20 +216,22 @@ class DatabaseManager:
             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
+                        "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
+                        "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
+            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}")
@@ -223,7 +245,7 @@ class DatabaseManager:
             if self.connection_pool:
                 async with self.connection_pool.acquire() as conn:
                     records = await conn.fetch(
-                        '''SELECT DISTINCT task_id,
+                        """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
@@ -232,14 +254,15 @@ class DatabaseManager:
                            WHERE user_id = $1
                            GROUP BY task_id
                            ORDER BY last_updated DESC
-                           LIMIT $2''',
-                        user_id, limit
+                           LIMIT $2""",
+                        user_id,
+                        limit,
                     )
             else:
                 conn = await asyncpg.connect(**self.db_config)
                 try:
                     records = await conn.fetch(
-                        '''SELECT DISTINCT task_id,
+                        """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
@@ -248,17 +271,18 @@ class DatabaseManager:
                            WHERE user_id = $1
                            GROUP BY task_id
                            ORDER BY last_updated DESC
-                           LIMIT $2''',
-                        user_id, limit
+                           LIMIT $2""",
+                        user_id,
+                        limit,
                     )
                 finally:
                     await conn.close()
 
             return [
                 {
-                    "task_id": r['task_id'],
-                    "last_updated": r['last_updated'].isoformat(),
-                    "status": r['status']
+                    "task_id": r["task_id"],
+                    "last_updated": r["last_updated"].isoformat(),
+                    "status": r["status"],
                 }
                 for r in records
             ]
@@ -275,15 +299,15 @@ class DatabaseManager:
             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
+                        "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
+                        "DELETE FROM task_history WHERE timestamp < NOW() - INTERVAL %s DAY",
+                        days_old,
                     )
                 finally:
                     await conn.close()