aiecs 1.5.1__py3-none-any.whl

aiecs/llm/client_factory.py

@@ -0,0 +1,420 @@
+import logging
+from typing import Dict, Any, Optional, Union, List
+from enum import Enum
+
+from .clients.base_client import BaseLLMClient, LLMMessage, LLMResponse
+from .clients.openai_client import OpenAIClient
+from .clients.vertex_client import VertexAIClient
+from .clients.googleai_client import GoogleAIClient
+from .clients.xai_client import XAIClient
+from .callbacks.custom_callbacks import CustomAsyncCallbackHandler
+
+logger = logging.getLogger(__name__)
+
+
+class AIProvider(str, Enum):
+    OPENAI = "OpenAI"
+    VERTEX = "Vertex"
+    GOOGLEAI = "GoogleAI"
+    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.GOOGLEAI:
+            return GoogleAIClient()
+        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}")
+
+    @classmethod
+    def reload_config(cls):
+        """
+        Reload LLM models configuration.
+
+        This reloads the configuration from the YAML file, allowing for
+        hot-reloading of model settings without restarting the application.
+        """
+        try:
+            from aiecs.llm.config import reload_llm_config
+
+            config = reload_llm_config()
+            logger.info(f"Reloaded LLM configuration: {len(config.providers)} providers")
+            return config
+        except Exception as e:
+            logger.error(f"Failed to reload LLM configuration: {e}")
+            raise
+
+
+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

aiecs/llm/clients/__init__.py

@@ -0,0 +1,33 @@
+"""
+LLM Client implementations.
+
+This package contains all LLM provider client implementations.
+"""
+
+from .base_client import (
+    BaseLLMClient,
+    LLMMessage,
+    LLMResponse,
+    LLMClientError,
+    ProviderNotAvailableError,
+    RateLimitError,
+)
+from .openai_client import OpenAIClient
+from .vertex_client import VertexAIClient
+from .googleai_client import GoogleAIClient
+from .xai_client import XAIClient
+
+__all__ = [
+    # Base classes
+    "BaseLLMClient",
+    "LLMMessage",
+    "LLMResponse",
+    "LLMClientError",
+    "ProviderNotAvailableError",
+    "RateLimitError",
+    # Client implementations
+    "OpenAIClient",
+    "VertexAIClient",
+    "GoogleAIClient",
+    "XAIClient",
+]

aiecs/llm/clients/base_client.py

@@ -0,0 +1,193 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any, Optional, List, AsyncGenerator
+from dataclasses import dataclass
+import logging
+
+logger = logging.getLogger(__name__)
+
+# Lazy import to avoid circular dependency
+
+
+def _get_config_loader():
+    """Lazy import of config loader to avoid circular dependency"""
+    from aiecs.llm.config import get_llm_config_loader
+
+    return get_llm_config_loader()
+
+
+@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
+    # Added for backward compatibility
+    metadata: Optional[Dict[str, Any]] = 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"""
+
+
+class ProviderNotAvailableError(LLMClientError):
+    """Raised when a provider is not available or misconfigured"""
+
+
+class RateLimitError(LLMClientError):
+    """Raised when rate limit is exceeded"""
+
+
+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"""
+
+    @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"""
+
+    @abstractmethod
+    async def close(self):
+        """Clean up resources"""
+
+    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.
+
+        DEPRECATED: Use _estimate_cost_from_config instead for config-based cost estimation.
+        This method is kept for backward compatibility.
+        """
+        if model in token_costs:
+            costs = token_costs[model]
+            return (input_tokens * costs["input"] + output_tokens * costs["output"]) / 1000
+        return 0.0
+
+    def _estimate_cost_from_config(
+        self, model_name: str, input_tokens: int, output_tokens: int
+    ) -> float:
+        """
+        Estimate the cost using configuration-based pricing.
+
+        Args:
+            model_name: Name of the model
+            input_tokens: Number of input tokens
+            output_tokens: Number of output tokens
+
+        Returns:
+            Estimated cost in USD
+        """
+        try:
+            loader = _get_config_loader()
+            model_config = loader.get_model_config(self.provider_name, model_name)
+
+            if model_config and model_config.costs:
+                input_cost = (input_tokens * model_config.costs.input) / 1000
+                output_cost = (output_tokens * model_config.costs.output) / 1000
+                return input_cost + output_cost
+            else:
+                self.logger.warning(
+                    f"No cost configuration found for model {model_name} "
+                    f"in provider {self.provider_name}"
+                )
+                return 0.0
+        except Exception as e:
+            self.logger.warning(f"Failed to estimate cost from config: {e}")
+            return 0.0
+
+    def _get_model_config(self, model_name: str):
+        """
+        Get model configuration from the config loader.
+
+        Args:
+            model_name: Name of the model
+
+        Returns:
+            ModelConfig if found, None otherwise
+        """
+        try:
+            loader = _get_config_loader()
+            return loader.get_model_config(self.provider_name, model_name)
+        except Exception as e:
+            self.logger.warning(f"Failed to get model config: {e}")
+            return None
+
+    def _get_default_model(self) -> Optional[str]:
+        """
+        Get the default model for this provider from configuration.
+
+        Returns:
+            Default model name if configured, None otherwise
+        """
+        try:
+            loader = _get_config_loader()
+            return loader.get_default_model(self.provider_name)
+        except Exception as e:
+            self.logger.warning(f"Failed to get default model: {e}")
+            return None