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

aiecs/domain/task/task_context.py

@@ -1,23 +1,24 @@
 import time
 import logging
 import json
-import os
-from typing import Dict, Any, Optional, AsyncGenerator, List, Tuple
+from typing import Dict, Any, Optional, AsyncGenerator, List
 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)
+    # Additional metadata (e.g., file paths, model info)
+    metadata: Dict[str, Any]
+
 
 class TaskContext:
     """
@@ -30,6 +31,7 @@ class TaskContext:
     - 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")
@@ -41,7 +43,8 @@ class TaskContext:
         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
+        # Tracks file operations
+        self.file_tracker: Dict[str, Dict[str, Any]] = {}
         self.model_tracker: List[Dict[str, Any]] = []  # Tracks model usage
         self.metadata_toggles: Dict[str, bool] = data.get("metadata_toggles", {})
         self._initialize_persistence()
@@ -59,7 +62,7 @@ class TaskContext:
                             timestamp=entry["timestamp"],
                             update_type=entry["update_type"],
                             data=entry["data"],
-                            metadata=entry["metadata"]
+                            metadata=entry["metadata"],
                         )
                         for entry in raw_history
                     ]
@@ -76,7 +79,7 @@ class TaskContext:
                     "timestamp": update.timestamp,
                     "update_type": update.update_type,
                     "data": update.data,
-                    "metadata": update.metadata
+                    "metadata": update.metadata,
                 }
                 for update in self.context_history
             ]
@@ -86,13 +89,13 @@ class TaskContext:
         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):
+    def add_context_update(self, update_type: str, data: Any, metadata: Optional[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 {}
+            metadata=metadata or {},
         )
         self.context_history.append(update)
         logger.debug(f"Added context update: {update_type}")
@@ -109,9 +112,13 @@ class TaskContext:
             "operation": operation,
             "source": source,
             "timestamp": time.time(),
-            "state": "active"
+            "state": "active",
         }
-        self.add_context_update("file_operation", {"path": file_path, "operation": operation}, {"source": source})
+        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):
@@ -120,7 +127,7 @@ class TaskContext:
             "model_id": model_id,
             "provider_id": provider_id,
             "mode": mode,
-            "timestamp": time.time()
+            "timestamp": time.time(),
         }
         # Avoid duplicates
         if not self.model_tracker or self.model_tracker[-1] != model_entry:
@@ -159,11 +166,7 @@ class TaskContext:
 
     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
-        }
+        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."""
@@ -176,12 +179,12 @@ class TaskContext:
                     "timestamp": update.timestamp,
                     "update_type": update.update_type,
                     "data": update.data,
-                    "metadata": update.metadata
+                    "metadata": update.metadata,
                 }
                 for update in self.context_history
             ],
             "file_tracker": self.file_tracker,
-            "model_tracker": self.model_tracker
+            "model_tracker": self.model_tracker,
         }
 
     def __enter__(self):
@@ -196,7 +199,7 @@ class TaskContext:
         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'):
+                if hasattr(resource, "close"):
                     resource.close()
                 logger.debug(f"Cleaned up resource: {resource_name}")
             except Exception as e:
@@ -217,9 +220,9 @@ class TaskContext:
         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__'):
+                if hasattr(resource, "close"):
+                    if callable(getattr(resource, "close")):
+                        if hasattr(resource.close, "__await__"):
                             await resource.close()
                         else:
                             resource.close()

aiecs/infrastructure/__init__.py

@@ -15,11 +15,9 @@ __all__ = [
     "CeleryTaskManager",
     "WebSocketManager",
     "UserConfirmation",
-
     # Persistence
     "DatabaseManager",
     "RedisClient",
-
     # Monitoring
     "ExecutorMetrics",
     "TracingManager",

aiecs/infrastructure/graph_storage/__init__.py

@@ -0,0 +1,11 @@
+from aiecs.infrastructure.graph_storage.base import GraphStore
+from aiecs.infrastructure.graph_storage.in_memory import InMemoryGraphStore
+from aiecs.infrastructure.graph_storage.sqlite import SQLiteGraphStore
+from aiecs.infrastructure.graph_storage.postgres import PostgresGraphStore
+
+__all__ = [
+    "GraphStore",
+    "InMemoryGraphStore",
+    "SQLiteGraphStore",
+    "PostgresGraphStore",
+]