aiecs 1.0.0__py3-none-any.whl

aiecs/infrastructure/persistence/redis_client.py

@@ -0,0 +1,162 @@
+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, mapping: dict) -> int:
+        """Set hash fields"""
+        client = await self.get_client()
+        return await client.hset(name, 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,54 @@
+"""
+LLM Package - Modular AI Provider Architecture
+
+This package provides a unified interface to multiple AI providers through
+individual client implementations and a factory pattern.
+"""
+
+# Import all main components
+from .base_client import (
+    BaseLLMClient,
+    LLMMessage,
+    LLMResponse,
+    LLMClientError,
+    ProviderNotAvailableError,
+    RateLimitError
+)
+
+from .client_factory import (
+    AIProvider,
+    LLMClientFactory,
+    LLMClientManager,
+    get_llm_manager,
+    generate_text,
+    stream_text
+)
+
+from .openai_client import OpenAIClient
+from .vertex_client import VertexAIClient
+from .xai_client import XAIClient
+
+__all__ = [
+    # Base classes and types
+    'BaseLLMClient',
+    'LLMMessage',
+    'LLMResponse',
+    'LLMClientError',
+    'ProviderNotAvailableError',
+    'RateLimitError',
+    'AIProvider',
+
+    # Factory and manager
+    'LLMClientFactory',
+    'LLMClientManager',
+    'get_llm_manager',
+
+    # Individual clients
+    'OpenAIClient',
+    'VertexAIClient',
+    'XAIClient',
+
+    # Convenience functions
+    'generate_text',
+    'stream_text',
+]

aiecs/llm/base_client.py

@@ -0,0 +1,99 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Optional, List, AsyncGenerator
+from dataclasses import dataclass
+import time
+import logging
+
+logger = logging.getLogger(__name__)
+
+@dataclass
+class LLMMessage:
+    role: str  # "system", "user", "assistant"
+    content: str
+
+@dataclass
+class LLMResponse:
+    content: str
+    provider: str
+    model: str
+    tokens_used: Optional[int] = None
+    prompt_tokens: Optional[int] = None
+    completion_tokens: Optional[int] = None
+    cost_estimate: Optional[float] = None
+    response_time: Optional[float] = None
+
+    def __post_init__(self):
+        """Ensure consistency of token data"""
+        # If there are detailed token information but no total, calculate the total
+        if self.prompt_tokens is not None and self.completion_tokens is not None and self.tokens_used is None:
+            self.tokens_used = self.prompt_tokens + self.completion_tokens
+
+        # If only total is available but no detailed information, try to estimate (cannot accurately allocate in this case)
+        elif self.tokens_used is not None and self.prompt_tokens is None and self.completion_tokens is None:
+            # In this case we cannot accurately allocate, keep as is
+            pass
+
+class LLMClientError(Exception):
+    """Base exception for LLM client errors"""
+    pass
+
+class ProviderNotAvailableError(LLMClientError):
+    """Raised when a provider is not available or misconfigured"""
+    pass
+
+class RateLimitError(LLMClientError):
+    """Raised when rate limit is exceeded"""
+    pass
+
+class BaseLLMClient(ABC):
+    """Abstract base class for all LLM provider clients"""
+
+    def __init__(self, provider_name: str):
+        self.provider_name = provider_name
+        self.logger = logging.getLogger(f"{__name__}.{provider_name}")
+
+    @abstractmethod
+    async def generate_text(
+        self,
+        messages: List[LLMMessage],
+        model: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        **kwargs
+    ) -> LLMResponse:
+        """Generate text using the provider's API"""
+        pass
+
+    @abstractmethod
+    async def stream_text(
+        self,
+        messages: List[LLMMessage],
+        model: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        **kwargs
+    ) -> AsyncGenerator[str, None]:
+        """Stream text generation using the provider's API"""
+        pass
+
+    @abstractmethod
+    async def close(self):
+        """Clean up resources"""
+        pass
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()
+
+    def _count_tokens_estimate(self, text: str) -> int:
+        """Rough token count estimation (4 chars ≈ 1 token for English)"""
+        return len(text) // 4
+
+    def _estimate_cost(self, model: str, input_tokens: int, output_tokens: int, token_costs: Dict) -> float:
+        """Estimate the cost of the API call"""
+        if model in token_costs:
+            costs = token_costs[model]
+            return (input_tokens * costs["input"] + output_tokens * costs["output"]) / 1000
+        return 0.0

aiecs/llm/client_factory.py

@@ -0,0 +1,339 @@
+import logging
+from typing import Dict, Any, Optional, Union, List
+from enum import Enum
+
+from .base_client import BaseLLMClient, LLMMessage, LLMResponse
+from .openai_client import OpenAIClient
+from .vertex_client import VertexAIClient
+from .xai_client import XAIClient
+from ..utils.base_callback import CustomAsyncCallbackHandler
+
+logger = logging.getLogger(__name__)
+
+class AIProvider(str, Enum):
+    OPENAI = "OpenAI"
+    VERTEX = "Vertex"
+    XAI = "xAI"
+
+class LLMClientFactory:
+    """Factory for creating and managing LLM provider clients"""
+
+    _clients: Dict[AIProvider, BaseLLMClient] = {}
+
+    @classmethod
+    def get_client(cls, provider: Union[str, AIProvider]) -> BaseLLMClient:
+        """Get or create a client for the specified provider"""
+        if isinstance(provider, str):
+            try:
+                provider = AIProvider(provider)
+            except ValueError:
+                raise ValueError(f"Unsupported provider: {provider}")
+
+        if provider not in cls._clients:
+            cls._clients[provider] = cls._create_client(provider)
+
+        return cls._clients[provider]
+
+    @classmethod
+    def _create_client(cls, provider: AIProvider) -> BaseLLMClient:
+        """Create a new client instance for the provider"""
+        if provider == AIProvider.OPENAI:
+            return OpenAIClient()
+        elif provider == AIProvider.VERTEX:
+            return VertexAIClient()
+        elif provider == AIProvider.XAI:
+            return XAIClient()
+        else:
+            raise ValueError(f"Unsupported provider: {provider}")
+
+    @classmethod
+    async def close_all(cls):
+        """Close all active clients"""
+        for client in cls._clients.values():
+            try:
+                await client.close()
+            except Exception as e:
+                logger.error(f"Error closing client {client.provider_name}: {e}")
+        cls._clients.clear()
+
+    @classmethod
+    async def close_client(cls, provider: Union[str, AIProvider]):
+        """Close a specific client"""
+        if isinstance(provider, str):
+            provider = AIProvider(provider)
+
+        if provider in cls._clients:
+            try:
+                await cls._clients[provider].close()
+                del cls._clients[provider]
+            except Exception as e:
+                logger.error(f"Error closing client {provider}: {e}")
+
+class LLMClientManager:
+    """High-level manager for LLM operations with context-aware provider selection"""
+
+    def __init__(self):
+        self.factory = LLMClientFactory()
+
+    def _extract_ai_preference(self, context: Optional[Dict[str, Any]]) -> tuple[Optional[str], Optional[str]]:
+        """Extract AI provider and model from context"""
+        if not context:
+            return None, None
+
+        metadata = context.get('metadata', {})
+
+        # First, check for aiPreference in metadata
+        ai_preference = metadata.get('aiPreference', {})
+        if isinstance(ai_preference, dict):
+            provider = ai_preference.get('provider')
+            model = ai_preference.get('model')
+            if provider is not None:
+                return provider, model
+
+        # Fallback to direct provider/model in metadata
+        provider = metadata.get('provider')
+        model = metadata.get('model')
+        return provider, model
+
+    async def generate_text(
+        self,
+        messages: Union[str, list[LLMMessage]],
+        provider: Optional[Union[str, AIProvider]] = None,
+        model: Optional[str] = None,
+        context: Optional[Dict[str, Any]] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        callbacks: Optional[List[CustomAsyncCallbackHandler]] = None,
+        **kwargs
+    ) -> LLMResponse:
+        """
+        Generate text using context-aware provider selection
+
+        Args:
+            messages: Either a string prompt or list of LLMMessage objects
+            provider: AI provider to use (can be overridden by context)
+            model: Specific model to use (can be overridden by context)
+            context: TaskContext or dict containing aiPreference
+            temperature: Sampling temperature (0.0 to 2.0)
+            max_tokens: Maximum tokens to generate
+            callbacks: List of callback handlers to execute during LLM calls
+            **kwargs: Additional provider-specific parameters
+
+        Returns:
+            LLMResponse object with generated text and metadata
+        """
+        # Extract provider/model from context if available
+        context_provider, context_model = self._extract_ai_preference(context)
+
+        # Use context preferences if available, otherwise use provided values
+        final_provider = context_provider or provider or AIProvider.OPENAI
+        final_model = context_model or model
+
+        # Convert string prompt to messages format
+        if isinstance(messages, str):
+            messages = [LLMMessage(role="user", content=messages)]
+
+        # Execute on_llm_start callbacks
+        if callbacks:
+            # Convert LLMMessage objects to dictionaries for callbacks
+            messages_dict = [{"role": msg.role, "content": msg.content} for msg in messages]
+            for callback in callbacks:
+                try:
+                    await callback.on_llm_start(messages_dict, provider=final_provider, model=final_model, **kwargs)
+                except Exception as e:
+                    logger.error(f"Error in callback on_llm_start: {e}")
+
+        try:
+            # Get the appropriate client
+            client = self.factory.get_client(final_provider)
+
+            # Generate text
+            response = await client.generate_text(
+                messages=messages,
+                model=final_model,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                **kwargs
+            )
+
+            # Execute on_llm_end callbacks
+            if callbacks:
+                # Convert LLMResponse object to dictionary for callbacks
+                response_dict = {
+                    "content": response.content,
+                    "provider": response.provider,
+                    "model": response.model,
+                    "tokens_used": response.tokens_used,
+                    "prompt_tokens": response.prompt_tokens,
+                    "completion_tokens": response.completion_tokens,
+                    "cost_estimate": response.cost_estimate,
+                    "response_time": response.response_time
+                }
+                for callback in callbacks:
+                    try:
+                        await callback.on_llm_end(response_dict, provider=final_provider, model=final_model, **kwargs)
+                    except Exception as e:
+                        logger.error(f"Error in callback on_llm_end: {e}")
+
+            logger.info(f"Generated text using {final_provider}/{response.model}")
+            return response
+
+        except Exception as e:
+            # Execute on_llm_error callbacks
+            if callbacks:
+                for callback in callbacks:
+                    try:
+                        await callback.on_llm_error(e, provider=final_provider, model=final_model, **kwargs)
+                    except Exception as callback_error:
+                        logger.error(f"Error in callback on_llm_error: {callback_error}")
+
+            # Re-raise the original exception
+            raise
+
+    async def stream_text(
+        self,
+        messages: Union[str, list[LLMMessage]],
+        provider: Optional[Union[str, AIProvider]] = None,
+        model: Optional[str] = None,
+        context: Optional[Dict[str, Any]] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        callbacks: Optional[List[CustomAsyncCallbackHandler]] = None,
+        **kwargs
+    ):
+        """
+        Stream text generation using context-aware provider selection
+
+        Args:
+            messages: Either a string prompt or list of LLMMessage objects
+            provider: AI provider to use (can be overridden by context)
+            model: Specific model to use (can be overridden by context)
+            context: TaskContext or dict containing aiPreference
+            temperature: Sampling temperature (0.0 to 2.0)
+            max_tokens: Maximum tokens to generate
+            callbacks: List of callback handlers to execute during LLM calls
+            **kwargs: Additional provider-specific parameters
+
+        Yields:
+            str: Incremental text chunks
+        """
+        # Extract provider/model from context if available
+        context_provider, context_model = self._extract_ai_preference(context)
+
+        # Use context preferences if available, otherwise use provided values
+        final_provider = context_provider or provider or AIProvider.OPENAI
+        final_model = context_model or model
+
+        # Convert string prompt to messages format
+        if isinstance(messages, str):
+            messages = [LLMMessage(role="user", content=messages)]
+
+        # Execute on_llm_start callbacks
+        if callbacks:
+            # Convert LLMMessage objects to dictionaries for callbacks
+            messages_dict = [{"role": msg.role, "content": msg.content} for msg in messages]
+            for callback in callbacks:
+                try:
+                    await callback.on_llm_start(messages_dict, provider=final_provider, model=final_model, **kwargs)
+                except Exception as e:
+                    logger.error(f"Error in callback on_llm_start: {e}")
+
+        try:
+            # Get the appropriate client
+            client = self.factory.get_client(final_provider)
+
+            # Collect streamed content for token counting
+            collected_content = ""
+
+            # Stream text
+            async for chunk in await client.stream_text(
+                messages=messages,
+                model=final_model,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                **kwargs
+            ):
+                collected_content += chunk
+                yield chunk
+
+            # Create a response object for callbacks (streaming doesn't return LLMResponse directly)
+            # We need to estimate token usage for streaming responses
+            estimated_tokens = len(collected_content) // 4  # Rough estimation
+            stream_response = LLMResponse(
+                content=collected_content,
+                provider=str(final_provider),
+                model=final_model or "unknown",
+                tokens_used=estimated_tokens
+            )
+
+            # Execute on_llm_end callbacks
+            if callbacks:
+                # Convert LLMResponse object to dictionary for callbacks
+                response_dict = {
+                    "content": stream_response.content,
+                    "provider": stream_response.provider,
+                    "model": stream_response.model,
+                    "tokens_used": stream_response.tokens_used,
+                    "prompt_tokens": stream_response.prompt_tokens,
+                    "completion_tokens": stream_response.completion_tokens,
+                    "cost_estimate": stream_response.cost_estimate,
+                    "response_time": stream_response.response_time
+                }
+                for callback in callbacks:
+                    try:
+                        await callback.on_llm_end(response_dict, provider=final_provider, model=final_model, **kwargs)
+                    except Exception as e:
+                        logger.error(f"Error in callback on_llm_end: {e}")
+
+        except Exception as e:
+            # Execute on_llm_error callbacks
+            if callbacks:
+                for callback in callbacks:
+                    try:
+                        await callback.on_llm_error(e, provider=final_provider, model=final_model, **kwargs)
+                    except Exception as callback_error:
+                        logger.error(f"Error in callback on_llm_error: {callback_error}")
+
+            # Re-raise the original exception
+            raise
+
+    async def close(self):
+        """Close all clients"""
+        await self.factory.close_all()
+
+# Global instance for easy access
+_llm_manager = LLMClientManager()
+
+async def get_llm_manager() -> LLMClientManager:
+    """Get the global LLM manager instance"""
+    return _llm_manager
+
+# Convenience functions for backward compatibility
+async def generate_text(
+    messages: Union[str, list[LLMMessage]],
+    provider: Optional[Union[str, AIProvider]] = None,
+    model: Optional[str] = None,
+    context: Optional[Dict[str, Any]] = None,
+    temperature: float = 0.7,
+    max_tokens: Optional[int] = None,
+    callbacks: Optional[List[CustomAsyncCallbackHandler]] = None,
+    **kwargs
+) -> LLMResponse:
+    """Generate text using the global LLM manager"""
+    manager = await get_llm_manager()
+    return await manager.generate_text(messages, provider, model, context, temperature, max_tokens, callbacks, **kwargs)
+
+async def stream_text(
+    messages: Union[str, list[LLMMessage]],
+    provider: Optional[Union[str, AIProvider]] = None,
+    model: Optional[str] = None,
+    context: Optional[Dict[str, Any]] = None,
+    temperature: float = 0.7,
+    max_tokens: Optional[int] = None,
+    callbacks: Optional[List[CustomAsyncCallbackHandler]] = None,
+    **kwargs
+):
+    """Stream text using the global LLM manager"""
+    manager = await get_llm_manager()
+    async for chunk in manager.stream_text(messages, provider, model, context, temperature, max_tokens, callbacks, **kwargs):
+        yield chunk