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

aiecs/infrastructure/persistence/redis_client.py

@@ -5,10 +5,11 @@ import os
 
 logger = logging.getLogger(__name__)
 
+
 class RedisClient:
     """Redis client singleton for sharing across different caching strategies"""
 
-    def __init__(self):
+    def __init__(self) -> None:
         self._client: Optional[redis.Redis] = None
         self._connection_pool: Optional[redis.ConnectionPool] = None
 
@@ -16,10 +17,10 @@ class RedisClient:
         """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')
+            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(
@@ -29,7 +30,7 @@ class RedisClient:
                 password=redis_password,
                 decode_responses=True,
                 max_connections=20,
-                retry_on_timeout=True
+                retry_on_timeout=True,
             )
 
             # Create Redis client
@@ -62,27 +63,69 @@ class RedisClient:
     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)
+        return await client.hincrby(name, key, amount)  # type: ignore[misc]
 
     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)
+        return await client.hget(name, key)  # type: ignore[misc]
 
     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, mapping: dict) -> int:
-        """Set hash fields"""
+        return await client.hgetall(name)  # type: ignore[misc]
+
+    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()
-        return await client.hset(name, mapping=mapping)
+
+        if mapping is not None:
+            # Multiple fields mode
+            return await client.hset(name, mapping=mapping)  # type: ignore[misc]
+        elif key is not None and value is not None:
+            # Single field mode
+            return await client.hset(name, key=key, value=value)  # type: ignore[misc]
+        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)
+        return await client.expire(name, time)  # type: ignore[misc]
 
     async def exists(self, name: str) -> bool:
         """Check if key exists"""
@@ -99,11 +142,11 @@ class RedisClient:
             logger.error(f"Redis ping failed: {e}")
             return False
 
-    async def info(self, section: str = None) -> dict:
+    async def info(self, section: Optional[str] = None) -> dict:
         """Get Redis server information"""
         try:
             client = await self.get_client()
-            return await client.info(section)
+            return await client.info(section)  # type: ignore[misc]
         except Exception as e:
             logger.error(f"Redis info failed: {e}")
             return {}
@@ -112,16 +155,16 @@ class RedisClient:
         """Delete one or more keys"""
         try:
             client = await self.get_client()
-            return await client.delete(*keys)
+            return await client.delete(*keys)  # type: ignore[misc]
         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:
+    async def set(self, key: str, value: str, ex: Optional[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)
+            return await client.set(key, value, ex=ex)  # type: ignore[misc]
         except Exception as e:
             logger.error(f"Redis set failed for key {key}: {e}")
             return False
@@ -130,17 +173,21 @@ class RedisClient:
         """Get value by key"""
         try:
             client = await self.get_client()
-            return await client.get(key)
+            return await client.get(key)  # type: ignore[misc]
         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.
+# 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
@@ -148,13 +195,19 @@ async def initialize_redis_client():
         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:

aiecs/llm/__init__.py

@@ -3,52 +3,100 @@ 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 all main components
-from .base_client import (
+# Import from organized subpackages
+from .clients import (
     BaseLLMClient,
+    CacheControl,
     LLMMessage,
     LLMResponse,
     LLMClientError,
     ProviderNotAvailableError,
-    RateLimitError
+    RateLimitError,
+    OpenAIClient,
+    VertexAIClient,
+    GoogleAIClient,
+    XAIClient,
 )
 
+from .protocols import LLMClientProtocol
+
 from .client_factory import (
     AIProvider,
     LLMClientFactory,
     LLMClientManager,
     get_llm_manager,
     generate_text,
-    stream_text
+    stream_text,
 )
 
-from .openai_client import OpenAIClient
-from .vertex_client import VertexAIClient
-from .xai_client import XAIClient
+from .config import (
+    ModelCostConfig,
+    ModelCapabilities,
+    ModelDefaultParams,
+    ModelConfig,
+    ProviderConfig,
+    LLMModelsConfig,
+    LLMConfigLoader,
+    get_llm_config_loader,
+    get_llm_config,
+    reload_llm_config,
+)
+
+from .callbacks import CustomAsyncCallbackHandler
+
+from .client_resolver import (
+    resolve_llm_client,
+    clear_client_cache,
+    get_cached_providers,
+)
 
 __all__ = [
     # Base classes and types
-    'BaseLLMClient',
-    'LLMMessage',
-    'LLMResponse',
-    'LLMClientError',
-    'ProviderNotAvailableError',
-    'RateLimitError',
-    'AIProvider',
-
+    "BaseLLMClient",
+    "CacheControl",
+    "LLMMessage",
+    "LLMResponse",
+    "LLMClientError",
+    "ProviderNotAvailableError",
+    "RateLimitError",
+    "AIProvider",
+    # Protocols
+    "LLMClientProtocol",
     # Factory and manager
-    'LLMClientFactory',
-    'LLMClientManager',
-    'get_llm_manager',
-
+    "LLMClientFactory",
+    "LLMClientManager",
+    "get_llm_manager",
     # Individual clients
-    'OpenAIClient',
-    'VertexAIClient',
-    'XAIClient',
-
+    "OpenAIClient",
+    "VertexAIClient",
+    "GoogleAIClient",
+    "XAIClient",
     # Convenience functions
-    'generate_text',
-    'stream_text',
+    "generate_text",
+    "stream_text",
+    # Client resolution helpers
+    "resolve_llm_client",
+    "clear_client_cache",
+    "get_cached_providers",
+    # 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/{custom_callbacks.py → callbacks/custom_callbacks.py}

@@ -1,12 +1,13 @@
 from typing import Any, List, Optional
 import logging
+import time
 
 # Import the base callback handler from utils
-from ..utils.base_callback import CustomAsyncCallbackHandler
+from aiecs.utils.base_callback import CustomAsyncCallbackHandler
+
 # Import LLM types for internal use only
-from .base_client import LLMMessage, LLMResponse
 # Import token usage repository
-from ..utils.token_usage_repository import token_usage_repo
+from aiecs.utils.token_usage_repository import token_usage_repo
 
 logger = logging.getLogger(__name__)
 
@@ -22,12 +23,13 @@ class RedisTokenCallbackHandler(CustomAsyncCallbackHandler):
             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
+        self.start_time: Optional[float] = None
+        self.messages: Optional[List[dict]] = 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
 
@@ -39,6 +41,7 @@ class RedisTokenCallbackHandler(CustomAsyncCallbackHandler):
             # 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")
 
@@ -47,11 +50,7 @@ class RedisTokenCallbackHandler(CustomAsyncCallbackHandler):
 
             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
-                )
+                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:
@@ -65,6 +64,7 @@ class RedisTokenCallbackHandler(CustomAsyncCallbackHandler):
         """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:
@@ -82,13 +82,14 @@ class DetailedRedisTokenCallbackHandler(CustomAsyncCallbackHandler):
             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.start_time: Optional[float] = None
+        self.messages: Optional[List[dict]] = 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
 
@@ -103,6 +104,7 @@ class DetailedRedisTokenCallbackHandler(CustomAsyncCallbackHandler):
             # 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")
 
@@ -119,7 +121,7 @@ class DetailedRedisTokenCallbackHandler(CustomAsyncCallbackHandler):
                     self.user_id,
                     prompt_tokens,
                     completion_tokens,
-                    self.cycle_start_date
+                    self.cycle_start_date,
                 )
 
                 logger.info(f"[DetailedCallback] Recorded detailed tokens for user '{self.user_id}': prompt={prompt_tokens}, completion={completion_tokens}")
@@ -134,6 +136,7 @@ class DetailedRedisTokenCallbackHandler(CustomAsyncCallbackHandler):
         """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:
@@ -141,7 +144,7 @@ class DetailedRedisTokenCallbackHandler(CustomAsyncCallbackHandler):
 
     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)
+        total_chars = sum(len(msg.get("content", "")) for msg in messages)
         # Rough estimation: 4 characters ≈ 1 token
         return total_chars // 4
 
@@ -153,14 +156,14 @@ class DetailedRedisTokenCallbackHandler(CustomAsyncCallbackHandler):
             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
+        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
+        tokens_used = response.get("tokens_used") or 0
         if tokens_used > 0:
             # Use previously estimated prompt tokens
             prompt_tokens = self.prompt_tokens
@@ -168,7 +171,7 @@ class DetailedRedisTokenCallbackHandler(CustomAsyncCallbackHandler):
             return prompt_tokens, completion_tokens
 
         # If no token information, try to estimate from response content
-        content = response.get('content', '')
+        content = response.get("content", "")
         if content:
             completion_tokens = len(content) // 4
             prompt_tokens = self.prompt_tokens
@@ -219,10 +222,14 @@ def create_token_callback(user_id: str, cycle_start_date: Optional[str] = None)
     """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:
+
+def create_composite_callback(
+    *handlers: CustomAsyncCallbackHandler,
+) -> CompositeCallbackHandler:
     """Create a composite callback handler"""
     return CompositeCallbackHandler(list(handlers))