aiecs 1.5.1__py3-none-any.whl

aiecs/infrastructure/persistence/redis_client.py

@@ -0,0 +1,220 @@
+import redis.asyncio as redis
+import logging
+from typing import Optional
+import os
+
+logger = logging.getLogger(__name__)
+
+
+class RedisClient:
+    """Redis client singleton for sharing across different caching strategies"""
+
+    def __init__(self):
+        self._client: Optional[redis.Redis] = None
+        self._connection_pool: Optional[redis.ConnectionPool] = None
+
+    async def initialize(self):
+        """Initialize Redis client"""
+        try:
+            # Get Redis configuration from environment variables
+            redis_host = os.getenv("REDIS_HOST", "localhost")
+            redis_port = int(os.getenv("REDIS_PORT", 6379))
+            redis_db = int(os.getenv("REDIS_DB", 0))
+            redis_password = os.getenv("REDIS_PASSWORD")
+
+            # Create connection pool
+            self._connection_pool = redis.ConnectionPool(
+                host=redis_host,
+                port=redis_port,
+                db=redis_db,
+                password=redis_password,
+                decode_responses=True,
+                max_connections=20,
+                retry_on_timeout=True,
+            )
+
+            # Create Redis client
+            self._client = redis.Redis(connection_pool=self._connection_pool)
+
+            # Test connection
+            await self._client.ping()
+            logger.info(
+                f"Redis client initialized successfully: {redis_host}:{redis_port}/{redis_db}"
+            )
+
+        except Exception as e:
+            logger.error(f"Failed to initialize Redis client: {e}")
+            raise
+
+    async def get_client(self) -> redis.Redis:
+        """Get Redis client instance"""
+        if self._client is None:
+            raise RuntimeError("Redis client not initialized. Call initialize() first.")
+        return self._client
+
+    async def close(self):
+        """Close Redis connection"""
+        if self._client:
+            await self._client.close()
+            self._client = None
+        if self._connection_pool:
+            await self._connection_pool.disconnect()
+            self._connection_pool = None
+        logger.info("Redis client closed")
+
+    async def hincrby(self, name: str, key: str, amount: int = 1) -> int:
+        """Atomically increment hash field"""
+        client = await self.get_client()
+        return await client.hincrby(name, key, amount)
+
+    async def hget(self, name: str, key: str) -> Optional[str]:
+        """Get hash field value"""
+        client = await self.get_client()
+        return await client.hget(name, key)
+
+    async def hgetall(self, name: str) -> dict:
+        """Get all hash fields"""
+        client = await self.get_client()
+        return await client.hgetall(name)
+
+    async def hset(
+        self,
+        name: str,
+        key: Optional[str] = None,
+        value: Optional[str] = None,
+        mapping: Optional[dict] = None,
+    ) -> int:
+        """Set hash fields
+
+        Supports two calling patterns:
+        1. hset(name, key, value) - Set single field (positional)
+        2. hset(name, key=key, value=value) - Set single field (keyword)
+        3. hset(name, mapping={...}) - Set multiple fields
+
+        Args:
+            name: Redis hash key name
+            key: Field name (for single field set)
+            value: Field value (for single field set)
+            mapping: Dictionary of field-value pairs (for multiple fields)
+
+        Returns:
+            Number of fields that were added
+
+        Raises:
+            ValueError: If neither (key, value) nor mapping is provided
+
+        Examples:
+            # Single field with positional args
+            await redis_client.hset("myhash", "field1", "value1")
+
+            # Single field with keyword args
+            await redis_client.hset("myhash", key="field1", value="value1")
+
+            # Multiple fields with mapping
+            await redis_client.hset("myhash", mapping={"field1": "value1", "field2": "value2"})
+        """
+        client = await self.get_client()
+
+        if mapping is not None:
+            # Multiple fields mode
+            return await client.hset(name, mapping=mapping)
+        elif key is not None and value is not None:
+            # Single field mode
+            return await client.hset(name, key=key, value=value)
+        else:
+            raise ValueError(
+                "Either provide (key, value) or mapping parameter. "
+                f"Got: key={key}, value={value}, mapping={mapping}"
+            )
+
+    async def expire(self, name: str, time: int) -> bool:
+        """Set expiration time"""
+        client = await self.get_client()
+        return await client.expire(name, time)
+
+    async def exists(self, name: str) -> bool:
+        """Check if key exists"""
+        client = await self.get_client()
+        return bool(await client.exists(name))
+
+    async def ping(self) -> bool:
+        """Test Redis connection"""
+        try:
+            client = await self.get_client()
+            result = await client.ping()
+            return result
+        except Exception as e:
+            logger.error(f"Redis ping failed: {e}")
+            return False
+
+    async def info(self, section: str = None) -> dict:
+        """Get Redis server information"""
+        try:
+            client = await self.get_client()
+            return await client.info(section)
+        except Exception as e:
+            logger.error(f"Redis info failed: {e}")
+            return {}
+
+    async def delete(self, *keys) -> int:
+        """Delete one or more keys"""
+        try:
+            client = await self.get_client()
+            return await client.delete(*keys)
+        except Exception as e:
+            logger.error(f"Redis delete failed: {e}")
+            return 0
+
+    async def set(self, key: str, value: str, ex: int = None) -> bool:
+        """Set a key-value pair with optional expiration"""
+        try:
+            client = await self.get_client()
+            return await client.set(key, value, ex=ex)
+        except Exception as e:
+            logger.error(f"Redis set failed for key {key}: {e}")
+            return False
+
+    async def get(self, key: str) -> Optional[str]:
+        """Get value by key"""
+        try:
+            client = await self.get_client()
+            return await client.get(key)
+        except Exception as e:
+            logger.error(f"Redis get failed for key {key}: {e}")
+            return None
+
+
+# ✅ Key changes:
+# 1. No longer create instance immediately.
+# 2. Define a global variable with initial value None. This variable will
+# be populated by lifespan.
+redis_client: Optional[RedisClient] = None
+
+# 3. Provide an initialization function for lifespan to call
+
+
+async def initialize_redis_client():
+    """Create and initialize global Redis client instance at application startup."""
+    global redis_client
+    if redis_client is None:
+        redis_client = RedisClient()
+        await redis_client.initialize()
+
+
+# 4. Provide a close function for lifespan to call
+
+
+async def close_redis_client():
+    """Close global Redis client instance at application shutdown."""
+    if redis_client:
+        await redis_client.close()
+
+
+# For backward compatibility, keep get_redis_client function
+
+
+async def get_redis_client() -> RedisClient:
+    """Get global Redis client instance"""
+    if redis_client is None:
+        raise RuntimeError("Redis client not initialized. Call initialize_redis_client() first.")
+    return redis_client

aiecs/llm/__init__.py

@@ -0,0 +1,86 @@
+"""
+LLM Package - Modular AI Provider Architecture
+
+This package provides a unified interface to multiple AI providers through
+individual client implementations and a factory pattern.
+
+Package Structure:
+- clients/: LLM client implementations
+- config/: Configuration management
+- callbacks/: Callback handlers
+- utils/: Utility functions and scripts
+"""
+
+# Import from organized subpackages
+from .clients import (
+    BaseLLMClient,
+    LLMMessage,
+    LLMResponse,
+    LLMClientError,
+    ProviderNotAvailableError,
+    RateLimitError,
+    OpenAIClient,
+    VertexAIClient,
+    GoogleAIClient,
+    XAIClient,
+)
+
+from .client_factory import (
+    AIProvider,
+    LLMClientFactory,
+    LLMClientManager,
+    get_llm_manager,
+    generate_text,
+    stream_text,
+)
+
+from .config import (
+    ModelCostConfig,
+    ModelCapabilities,
+    ModelDefaultParams,
+    ModelConfig,
+    ProviderConfig,
+    LLMModelsConfig,
+    LLMConfigLoader,
+    get_llm_config_loader,
+    get_llm_config,
+    reload_llm_config,
+)
+
+from .callbacks import CustomAsyncCallbackHandler
+
+__all__ = [
+    # Base classes and types
+    "BaseLLMClient",
+    "LLMMessage",
+    "LLMResponse",
+    "LLMClientError",
+    "ProviderNotAvailableError",
+    "RateLimitError",
+    "AIProvider",
+    # Factory and manager
+    "LLMClientFactory",
+    "LLMClientManager",
+    "get_llm_manager",
+    # Individual clients
+    "OpenAIClient",
+    "VertexAIClient",
+    "GoogleAIClient",
+    "XAIClient",
+    # Convenience functions
+    "generate_text",
+    "stream_text",
+    # Configuration management
+    "ModelCostConfig",
+    "ModelCapabilities",
+    "ModelDefaultParams",
+    "ModelConfig",
+    "ProviderConfig",
+    "LLMModelsConfig",
+    "LLMConfigLoader",
+    "get_llm_config_loader",
+    "get_llm_config",
+    "reload_llm_config",
+    # Callbacks
+    "CustomAsyncCallbackHandler",
+]

aiecs/llm/callbacks/__init__.py

@@ -0,0 +1,11 @@
+"""
+LLM Callbacks.
+
+This package contains callback handlers for LLM operations.
+"""
+
+from .custom_callbacks import CustomAsyncCallbackHandler
+
+__all__ = [
+    "CustomAsyncCallbackHandler",
+]

aiecs/llm/callbacks/custom_callbacks.py

@@ -0,0 +1,264 @@
+from typing import Any, List, Optional
+import logging
+
+# Import the base callback handler from utils
+from aiecs.utils.base_callback import CustomAsyncCallbackHandler
+
+# Import LLM types for internal use only
+# Import token usage repository
+from aiecs.utils.token_usage_repository import token_usage_repo
+
+logger = logging.getLogger(__name__)
+
+
+class RedisTokenCallbackHandler(CustomAsyncCallbackHandler):
+    """
+    Concrete token recording callback handler.
+    Responsible for recording token usage after LLM calls by delegating to the repository.
+    """
+
+    def __init__(self, user_id: str, cycle_start_date: Optional[str] = None):
+        if not user_id:
+            raise ValueError("user_id must be provided for RedisTokenCallbackHandler")
+        self.user_id = user_id
+        self.cycle_start_date = cycle_start_date
+        self.start_time = None
+        self.messages = None
+
+    async def on_llm_start(self, messages: List[dict], **kwargs: Any) -> None:
+        """Triggered when LLM call starts"""
+        import time
+
+        self.start_time = time.time()
+        self.messages = messages
+
+        logger.info(
+            f"[Callback] LLM call started for user '{self.user_id}' with {len(messages)} messages"
+        )
+
+    async def on_llm_end(self, response: dict, **kwargs: Any) -> None:
+        """Triggered when LLM call ends successfully"""
+        try:
+            # Record call duration
+            if self.start_time:
+                import time
+
+                call_duration = time.time() - self.start_time
+                logger.info(
+                    f"[Callback] LLM call completed for user '{self.user_id}' in {call_duration:.2f}s"
+                )
+
+            # Extract token usage from response dictionary
+            tokens_used = response.get("tokens_used")
+
+            if tokens_used and tokens_used > 0:
+                # Delegate recording work to repository
+                await token_usage_repo.increment_total_usage(
+                    self.user_id, tokens_used, self.cycle_start_date
+                )
+
+                logger.info(f"[Callback] Recorded {tokens_used} tokens for user '{self.user_id}'")
+            else:
+                logger.warning(
+                    f"[Callback] No token usage data available for user '{self.user_id}'"
+                )
+
+        except Exception as e:
+            logger.error(f"[Callback] Failed to record token usage for user '{self.user_id}': {e}")
+            # Don't re-raise exception to avoid affecting main LLM call flow
+
+    async def on_llm_error(self, error: Exception, **kwargs: Any) -> None:
+        """Triggered when LLM call encounters an error"""
+        if self.start_time:
+            import time
+
+            call_duration = time.time() - self.start_time
+            logger.error(
+                f"[Callback] LLM call failed for user '{self.user_id}' after {call_duration:.2f}s: {error}"
+            )
+        else:
+            logger.error(f"[Callback] LLM call failed for user '{self.user_id}': {error}")
+
+
+class DetailedRedisTokenCallbackHandler(CustomAsyncCallbackHandler):
+    """
+    Detailed token recording callback handler.
+    Records separate prompt and completion token usage in addition to total usage.
+    """
+
+    def __init__(self, user_id: str, cycle_start_date: Optional[str] = None):
+        if not user_id:
+            raise ValueError("user_id must be provided for DetailedRedisTokenCallbackHandler")
+        self.user_id = user_id
+        self.cycle_start_date = cycle_start_date
+        self.start_time = None
+        self.messages = None
+        self.prompt_tokens = 0
+
+    async def on_llm_start(self, messages: List[dict], **kwargs: Any) -> None:
+        """Triggered when LLM call starts"""
+        import time
+
+        self.start_time = time.time()
+        self.messages = messages
+
+        # Estimate input token count
+        self.prompt_tokens = self._estimate_prompt_tokens(messages)
+
+        logger.info(
+            f"[DetailedCallback] LLM call started for user '{self.user_id}' with estimated {self.prompt_tokens} prompt tokens"
+        )
+
+    async def on_llm_end(self, response: dict, **kwargs: Any) -> None:
+        """Triggered when LLM call ends successfully"""
+        try:
+            # Record call duration
+            if self.start_time:
+                import time
+
+                call_duration = time.time() - self.start_time
+                logger.info(
+                    f"[DetailedCallback] LLM call completed for user '{self.user_id}' in {call_duration:.2f}s"
+                )
+
+            # Extract detailed token information from response
+            prompt_tokens, completion_tokens = self._extract_detailed_tokens(response)
+
+            # Ensure we have valid integers (not None)
+            prompt_tokens = prompt_tokens or 0
+            completion_tokens = completion_tokens or 0
+
+            if prompt_tokens > 0 or completion_tokens > 0:
+                # Use detailed token recording method
+                await token_usage_repo.increment_detailed_usage(
+                    self.user_id,
+                    prompt_tokens,
+                    completion_tokens,
+                    self.cycle_start_date,
+                )
+
+                logger.info(
+                    f"[DetailedCallback] Recorded detailed tokens for user '{self.user_id}': prompt={prompt_tokens}, completion={completion_tokens}"
+                )
+            else:
+                logger.warning(
+                    f"[DetailedCallback] No detailed token usage data available for user '{self.user_id}'"
+                )
+
+        except Exception as e:
+            logger.error(
+                f"[DetailedCallback] Failed to record detailed token usage for user '{self.user_id}': {e}"
+            )
+            # Don't re-raise exception to avoid affecting main LLM call flow
+
+    async def on_llm_error(self, error: Exception, **kwargs: Any) -> None:
+        """Triggered when LLM call encounters an error"""
+        if self.start_time:
+            import time
+
+            call_duration = time.time() - self.start_time
+            logger.error(
+                f"[DetailedCallback] LLM call failed for user '{self.user_id}' after {call_duration:.2f}s: {error}"
+            )
+        else:
+            logger.error(f"[DetailedCallback] LLM call failed for user '{self.user_id}': {error}")
+
+    def _estimate_prompt_tokens(self, messages: List[dict]) -> int:
+        """Estimate token count for input messages"""
+        total_chars = sum(len(msg.get("content", "")) for msg in messages)
+        # Rough estimation: 4 characters ≈ 1 token
+        return total_chars // 4
+
+    def _extract_detailed_tokens(self, response: dict) -> tuple[int, int]:
+        """
+        Extract detailed token information from response dictionary
+
+        Returns:
+            tuple: (prompt_tokens, completion_tokens)
+        """
+        # If response has detailed token information, use it first
+        prompt_tokens = response.get("prompt_tokens") or 0
+        completion_tokens = response.get("completion_tokens") or 0
+
+        if prompt_tokens > 0 and completion_tokens > 0:
+            return prompt_tokens, completion_tokens
+
+        # If only total token count is available, try to allocate
+        tokens_used = response.get("tokens_used") or 0
+        if tokens_used > 0:
+            # Use previously estimated prompt tokens
+            prompt_tokens = self.prompt_tokens
+            completion_tokens = max(0, tokens_used - prompt_tokens)
+            return prompt_tokens, completion_tokens
+
+        # If no token information, try to estimate from response content
+        content = response.get("content", "")
+        if content:
+            completion_tokens = len(content) // 4
+            prompt_tokens = self.prompt_tokens
+            return prompt_tokens, completion_tokens
+
+        return 0, 0
+
+
+class CompositeCallbackHandler(CustomAsyncCallbackHandler):
+    """
+    Composite callback handler that can execute multiple callback handlers simultaneously
+    """
+
+    def __init__(self, handlers: List[CustomAsyncCallbackHandler]):
+        self.handlers = handlers or []
+
+    def add_handler(self, handler: CustomAsyncCallbackHandler):
+        """Add a callback handler"""
+        self.handlers.append(handler)
+
+    async def on_llm_start(self, messages: List[dict], **kwargs: Any) -> None:
+        """Execute start callbacks for all handlers"""
+        for handler in self.handlers:
+            try:
+                await handler.on_llm_start(messages, **kwargs)
+            except Exception as e:
+                logger.error(
+                    f"Error in callback handler {type(handler).__name__}.on_llm_start: {e}"
+                )
+
+    async def on_llm_end(self, response: dict, **kwargs: Any) -> None:
+        """Execute end callbacks for all handlers"""
+        for handler in self.handlers:
+            try:
+                await handler.on_llm_end(response, **kwargs)
+            except Exception as e:
+                logger.error(f"Error in callback handler {type(handler).__name__}.on_llm_end: {e}")
+
+    async def on_llm_error(self, error: Exception, **kwargs: Any) -> None:
+        """Execute error callbacks for all handlers"""
+        for handler in self.handlers:
+            try:
+                await handler.on_llm_error(error, **kwargs)
+            except Exception as e:
+                logger.error(
+                    f"Error in callback handler {type(handler).__name__}.on_llm_error: {e}"
+                )
+
+
+# Convenience functions for creating common callback handlers
+def create_token_callback(
+    user_id: str, cycle_start_date: Optional[str] = None
+) -> RedisTokenCallbackHandler:
+    """Create a basic token recording callback handler"""
+    return RedisTokenCallbackHandler(user_id, cycle_start_date)
+
+
+def create_detailed_token_callback(
+    user_id: str, cycle_start_date: Optional[str] = None
+) -> DetailedRedisTokenCallbackHandler:
+    """Create a detailed token recording callback handler"""
+    return DetailedRedisTokenCallbackHandler(user_id, cycle_start_date)
+
+
+def create_composite_callback(
+    *handlers: CustomAsyncCallbackHandler,
+) -> CompositeCallbackHandler:
+    """Create a composite callback handler"""
+    return CompositeCallbackHandler(list(handlers))