aiecs 1.0.0__py3-none-any.whl

aiecs/utils/token_usage_repository.py

@@ -0,0 +1,279 @@
+from datetime import datetime
+from typing import Optional, Dict, Any
+import logging
+from ..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,41 @@
+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)