aiecs 1.5.1__py3-none-any.whl

aiecs/utils/execution_utils.py

@@ -0,0 +1,184 @@
+import json
+import asyncio
+import threading
+import time
+from typing import Any, Callable, Dict, Optional
+from cachetools import LRUCache
+from contextlib import contextmanager
+import logging
+from tenacity import retry, stop_after_attempt, wait_exponential
+
+logger = logging.getLogger(__name__)
+
+
+class ExecutionUtils:
+    """
+    Provides common utility set for execution layer, including caching and retry logic.
+    """
+
+    def __init__(
+        self,
+        cache_size: int = 100,
+        cache_ttl: int = 3600,
+        retry_attempts: int = 3,
+        retry_backoff: float = 1.0,
+    ):
+        """
+        Initialize execution utility class.
+
+        Args:
+            cache_size (int): Maximum number of cache entries
+            cache_ttl (int): Cache time-to-live (seconds)
+            retry_attempts (int): Number of retry attempts
+            retry_backoff (float): Retry backoff factor
+        """
+        self.cache_size = cache_size
+        self.cache_ttl = cache_ttl
+        self.retry_attempts = retry_attempts
+        self.retry_backoff = retry_backoff
+        self._cache = LRUCache(maxsize=self.cache_size) if cache_size > 0 else None
+        self._cache_lock = threading.Lock()
+        self._cache_ttl_dict: Dict[str, float] = {}
+
+    def generate_cache_key(
+        self,
+        func_name: str,
+        user_id: str,
+        task_id: str,
+        args: tuple,
+        kwargs: Dict[str, Any],
+    ) -> str:
+        """
+        Generate context-based cache key including user ID, task ID, function name and parameters.
+
+        Args:
+            func_name (str): Function name
+            user_id (str): User ID
+            task_id (str): Task ID
+            args (tuple): Positional arguments
+            kwargs (Dict[str, Any]): Keyword arguments
+
+        Returns:
+            str: Cache key
+        """
+        key_dict = {
+            "func": func_name,
+            "user_id": user_id,
+            "task_id": task_id,
+            "args": args,
+            "kwargs": {k: v for k, v in kwargs.items() if k != "self"},
+        }
+        try:
+            key_str = json.dumps(key_dict, sort_keys=True)
+        except (TypeError, ValueError):
+            key_str = str(key_dict)
+        return hash(key_str).__str__()
+
+    def get_from_cache(self, cache_key: str) -> Optional[Any]:
+        """
+        Get result from cache if it exists and is not expired.
+
+        Args:
+            cache_key (str): Cache key
+
+        Returns:
+            Optional[Any]: Cached result or None
+        """
+        if self._cache is None:
+            return None
+        with self._cache_lock:
+            if cache_key in self._cache:
+                if (
+                    cache_key in self._cache_ttl_dict
+                    and time.time() > self._cache_ttl_dict[cache_key]
+                ):
+                    del self._cache[cache_key]
+                    del self._cache_ttl_dict[cache_key]
+                    return None
+                return self._cache[cache_key]
+        return None
+
+    def add_to_cache(self, cache_key: str, result: Any, ttl: Optional[int] = None) -> None:
+        """
+        Add result to cache with optional time-to-live setting.
+
+        Args:
+            cache_key (str): Cache key
+            result (Any): Cached result
+            ttl (Optional[int]): Time-to-live (seconds)
+        """
+        if self._cache is None:
+            return
+        with self._cache_lock:
+            self._cache[cache_key] = result
+            ttl = ttl if ttl is not None else self.cache_ttl
+            if ttl > 0:
+                self._cache_ttl_dict[cache_key] = time.time() + ttl
+
+    def create_retry_strategy(self, metric_name: Optional[str] = None) -> Callable:
+        """
+        Create retry strategy for execution operations.
+
+        Args:
+            metric_name (Optional[str]): Metric name for logging
+
+        Returns:
+            Callable: Retry decorator
+        """
+
+        def after_retry(retry_state):
+            logger.warning(
+                f"Retry {retry_state.attempt_number}/{self.retry_attempts} for {metric_name or 'operation'} after {retry_state.idle_for}s: {retry_state.outcome.exception()}"
+            )
+
+        return retry(
+            stop=stop_after_attempt(self.retry_attempts),
+            wait=wait_exponential(multiplier=self.retry_backoff, min=1, max=10),
+            after=after_retry,
+        )
+
+    @contextmanager
+    def timeout_context(self, seconds: int):
+        """
+        Context manager for enforcing operation timeout.
+
+        Args:
+            seconds (int): Timeout duration (seconds)
+
+        Raises:
+            TimeoutError: If operation exceeds timeout duration
+        """
+        loop = asyncio.get_event_loop()
+        future = asyncio.Future()
+        handle = loop.call_later(
+            seconds,
+            lambda: future.set_exception(TimeoutError(f"Operation timed out after {seconds}s")),
+        )
+        try:
+            yield future
+        finally:
+            handle.cancel()
+
+    async def execute_with_retry_and_timeout(
+        self, func: Callable, timeout: int, *args, **kwargs
+    ) -> Any:
+        """
+        Execute operation with retry and timeout mechanism.
+
+        Args:
+            func (Callable): Function to execute
+            timeout (int): Timeout duration (seconds)
+            *args: Positional arguments
+            **kwargs: Keyword arguments
+
+        Returns:
+            Any: Operation result
+
+        Raises:
+            OperationError: If all retry attempts fail
+        """
+        retry_strategy = self.create_retry_strategy(func.__name__)
+        try:
+            return await asyncio.wait_for(retry_strategy(func)(*args, **kwargs), timeout=timeout)
+        except asyncio.TimeoutError:
+            raise TimeoutError(f"Operation timed out after {timeout}s")

aiecs/utils/logging.py

@@ -0,0 +1 @@
+# Placeholder for logging.py

aiecs/utils/prompt_loader.py

@@ -0,0 +1,14 @@
+import yaml
+import os
+
+
+def get_prompt(mode: str, service: str) -> str:
+    """
+    Load the prompt for the specified service from services/{mode}/prompts.yaml.
+    """
+    path = f"app/services/{mode}/prompts.yaml"
+    if not os.path.exists(path):
+        return "[Default prompt]"
+    with open(path, "r", encoding="utf-8") as f:
+        data = yaml.safe_load(f)
+    return data.get(service, "[No specific prompt found]")

aiecs/utils/token_usage_repository.py

@@ -0,0 +1,323 @@
+from datetime import datetime
+from typing import Optional, Dict, Any
+import logging
+from aiecs.infrastructure.persistence.redis_client import get_redis_client
+
+logger = logging.getLogger(__name__)
+
+
+class TokenUsageRepository:
+    """Encapsulates all Redis operations related to user token usage"""
+
+    def _get_key_for_current_period(
+        self, user_id: str, cycle_start_date: Optional[str] = None
+    ) -> str:
+        """
+        Generate Redis key for current billing period
+
+        Args:
+            user_id: User ID
+            cycle_start_date: Cycle start date in YYYY-MM-DD format, defaults to current month if not provided
+
+        Returns:
+            Redis key string
+        """
+        if cycle_start_date:
+            # Use provided cycle start date
+            period = cycle_start_date
+        else:
+            # Use current month as default period
+            period = datetime.now().strftime("%Y-%m-%d")
+
+        return f"token_usage:{user_id}:{period}"
+
+    async def increment_prompt_tokens(
+        self,
+        user_id: str,
+        prompt_tokens: int,
+        cycle_start_date: Optional[str] = None,
+    ):
+        """
+        Increment prompt token usage for specified user
+
+        Args:
+            user_id: User ID
+            prompt_tokens: Number of input tokens
+            cycle_start_date: Cycle start date
+        """
+        if not user_id or prompt_tokens <= 0:
+            return
+
+        redis_key = self._get_key_for_current_period(user_id, cycle_start_date)
+
+        try:
+            # Use HINCRBY for atomic increment
+            client = await get_redis_client()
+            await client.hincrby(redis_key, "prompt_tokens", prompt_tokens)
+            logger.info(
+                f"[Repository] User '{user_id}' prompt tokens incremented by {prompt_tokens} in key '{redis_key}'."
+            )
+        except Exception as e:
+            logger.error(f"Failed to increment prompt tokens for user {user_id}: {e}")
+            raise
+
+    async def increment_completion_tokens(
+        self,
+        user_id: str,
+        completion_tokens: int,
+        cycle_start_date: Optional[str] = None,
+    ):
+        """
+        Increment completion token usage for specified user
+
+        Args:
+            user_id: User ID
+            completion_tokens: Number of output tokens
+            cycle_start_date: Cycle start date
+        """
+        if not user_id or completion_tokens <= 0:
+            return
+
+        redis_key = self._get_key_for_current_period(user_id, cycle_start_date)
+
+        try:
+            # Use HINCRBY for atomic increment
+            client = await get_redis_client()
+            await client.hincrby(redis_key, "completion_tokens", completion_tokens)
+            logger.info(
+                f"[Repository] User '{user_id}' completion tokens incremented by {completion_tokens} in key '{redis_key}'."
+            )
+        except Exception as e:
+            logger.error(f"Failed to increment completion tokens for user {user_id}: {e}")
+            raise
+
+    async def increment_total_usage(
+        self,
+        user_id: str,
+        total_tokens: int,
+        cycle_start_date: Optional[str] = None,
+    ):
+        """
+        Increment total token usage for specified user
+
+        Args:
+            user_id: User ID
+            total_tokens: Total number of tokens
+            cycle_start_date: Cycle start date
+        """
+        if not user_id or total_tokens <= 0:
+            return
+
+        redis_key = self._get_key_for_current_period(user_id, cycle_start_date)
+
+        try:
+            # Use HINCRBY for atomic increment
+            client = await get_redis_client()
+            await client.hincrby(redis_key, "total_tokens", total_tokens)
+            logger.info(
+                f"[Repository] User '{user_id}' total usage incremented by {total_tokens} tokens in key '{redis_key}'."
+            )
+        except Exception as e:
+            logger.error(f"Failed to increment total tokens for user {user_id}: {e}")
+            raise
+
+    async def increment_detailed_usage(
+        self,
+        user_id: str,
+        prompt_tokens: int,
+        completion_tokens: int,
+        cycle_start_date: Optional[str] = None,
+    ):
+        """
+        Increment both prompt and completion token usage for specified user
+
+        Args:
+            user_id: User ID
+            prompt_tokens: Number of input tokens
+            completion_tokens: Number of output tokens
+            cycle_start_date: Cycle start date
+        """
+        if not user_id or (prompt_tokens <= 0 and completion_tokens <= 0):
+            return
+
+        redis_key = self._get_key_for_current_period(user_id, cycle_start_date)
+
+        try:
+            # Batch update multiple fields
+            updates = {}
+            if prompt_tokens > 0:
+                updates["prompt_tokens"] = prompt_tokens
+            if completion_tokens > 0:
+                updates["completion_tokens"] = completion_tokens
+
+            # Calculate total token count
+            total_tokens = prompt_tokens + completion_tokens
+            if total_tokens > 0:
+                updates["total_tokens"] = total_tokens
+
+            # Use pipeline for batch operations
+            redis_client_instance = await get_redis_client()
+            client = await redis_client_instance.get_client()
+            pipe = client.pipeline()
+
+            for field, value in updates.items():
+                pipe.hincrby(redis_key, field, value)
+
+            await pipe.execute()
+
+            logger.info(
+                f"[Repository] User '{user_id}' detailed usage updated: prompt={prompt_tokens}, completion={completion_tokens}, total={total_tokens} in key '{redis_key}'."
+            )
+        except Exception as e:
+            logger.error(f"Failed to increment detailed usage for user {user_id}: {e}")
+            raise
+
+    async def get_usage_stats(
+        self, user_id: str, cycle_start_date: Optional[str] = None
+    ) -> Dict[str, int]:
+        """
+        Get token usage statistics for specified user
+
+        Args:
+            user_id: User ID
+            cycle_start_date: Cycle start date
+
+        Returns:
+            Dictionary containing token usage statistics
+        """
+        if not user_id:
+            return {}
+
+        redis_key = self._get_key_for_current_period(user_id, cycle_start_date)
+
+        try:
+            client = await get_redis_client()
+            stats = await client.hgetall(redis_key)
+
+            # Convert to integer type
+            result = {}
+            for key, value in stats.items():
+                try:
+                    result[key] = int(value) if value else 0
+                except (ValueError, TypeError):
+                    result[key] = 0
+
+            # Ensure required fields exist
+            result.setdefault("prompt_tokens", 0)
+            result.setdefault("completion_tokens", 0)
+            result.setdefault("total_tokens", 0)
+
+            logger.debug(f"[Repository] Retrieved usage stats for user '{user_id}': {result}")
+            return result
+
+        except Exception as e:
+            logger.error(f"Failed to get usage stats for user {user_id}: {e}")
+            return {
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "total_tokens": 0,
+            }
+
+    async def reset_usage(self, user_id: str, cycle_start_date: Optional[str] = None):
+        """
+        Reset token usage for specified user
+
+        Args:
+            user_id: User ID
+            cycle_start_date: Cycle start date
+        """
+        if not user_id:
+            return
+
+        redis_key = self._get_key_for_current_period(user_id, cycle_start_date)
+
+        try:
+            redis_client_instance = await get_redis_client()
+            client = await redis_client_instance.get_client()
+            await client.delete(redis_key)
+            logger.info(f"[Repository] Reset usage for user '{user_id}' in key '{redis_key}'.")
+        except Exception as e:
+            logger.error(f"Failed to reset usage for user {user_id}: {e}")
+            raise
+
+    async def set_usage_limit(
+        self, user_id: str, limit: int, cycle_start_date: Optional[str] = None
+    ):
+        """
+        Set token usage limit for user
+
+        Args:
+            user_id: User ID
+            limit: Token usage limit
+            cycle_start_date: Cycle start date
+        """
+        if not user_id or limit <= 0:
+            return
+
+        redis_key = self._get_key_for_current_period(user_id, cycle_start_date)
+
+        try:
+            client = await get_redis_client()
+            await client.hset(redis_key, {"usage_limit": str(limit)})
+            logger.info(
+                f"[Repository] Set usage limit {limit} for user '{user_id}' in key '{redis_key}'."
+            )
+        except Exception as e:
+            logger.error(f"Failed to set usage limit for user {user_id}: {e}")
+            raise
+
+    async def check_usage_limit(
+        self, user_id: str, cycle_start_date: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Check if user has exceeded usage limit
+
+        Args:
+            user_id: User ID
+            cycle_start_date: Cycle start date
+
+        Returns:
+            Dictionary containing limit check results
+        """
+        if not user_id:
+            return {
+                "exceeded": False,
+                "current_usage": 0,
+                "limit": 0,
+                "remaining": 0,
+            }
+
+        try:
+            stats = await self.get_usage_stats(user_id, cycle_start_date)
+            current_usage = stats.get("total_tokens", 0)
+
+            redis_key = self._get_key_for_current_period(user_id, cycle_start_date)
+            client = await get_redis_client()
+            limit_str = await client.hget(redis_key, "usage_limit")
+            limit = int(limit_str) if limit_str else 0
+
+            exceeded = limit > 0 and current_usage >= limit
+            remaining = max(0, limit - current_usage) if limit > 0 else float("inf")
+
+            result = {
+                "exceeded": exceeded,
+                "current_usage": current_usage,
+                "limit": limit,
+                "remaining": remaining,
+            }
+
+            logger.debug(f"[Repository] Usage limit check for user '{user_id}': {result}")
+            return result
+
+        except Exception as e:
+            logger.error(f"Failed to check usage limit for user {user_id}: {e}")
+            return {
+                "exceeded": False,
+                "current_usage": 0,
+                "limit": 0,
+                "remaining": 0,
+            }
+
+
+# Create a singleton for global application use
+token_usage_repo = TokenUsageRepository()

aiecs/ws/socket_server.py

@@ -0,0 +1,52 @@
+import socketio
+from aiecs.config.config import get_settings
+
+settings = get_settings()
+# In production, this should be set to specific origins
+# For example: ["https://your-frontend-domain.com"]
+allowed_origins = (
+    settings.cors_allowed_origins.split(",")
+    if hasattr(settings, "cors_allowed_origins")
+    else ["http://express-gateway:3001"]
+)
+
+# Allow all origins for development (more permissive)
+# In production, you should use specific origins
+# Explicitly set async_mode to 'asgi' for compatibility with uvicorn
+sio = socketio.AsyncServer(cors_allowed_origins="*", async_mode="asgi")
+# We no longer create a FastAPI app or combined ASGI app here
+# The FastAPI app will be created in main.py and the Socket.IO server will
+# be mounted there
+
+# Store connected clients by user ID
+connected_clients = {}
+
+
+@sio.event
+async def connect(sid, environ, auth=None):
+    print(f"Client connected: {sid}")
+
+
+@sio.event
+async def disconnect(sid):
+    print(f"Client disconnected: {sid}")
+    for user, socket_id in list(connected_clients.items()):
+        if socket_id == sid:
+            del connected_clients[user]
+
+
+@sio.event
+async def register(sid, data):
+    user_id = data.get("user_id")
+    if user_id:
+        connected_clients[user_id] = sid
+        print(f"Registered user {user_id} on SID {sid}")
+
+
+# Send progress update to user
+
+
+async def push_progress(user_id: str, data: dict):
+    sid = connected_clients.get(user_id)
+    if sid:
+        await sio.emit("progress", data, to=sid)