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

aiecs/llm/client_factory.py

@@ -1,33 +1,121 @@
 import logging
-from typing import Dict, Any, Optional, Union, List
+from typing import Dict, Any, Optional, Union, List, TYPE_CHECKING
 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
+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 .clients.openai_compatible_mixin import StreamChunk
+from .callbacks.custom_callbacks import CustomAsyncCallbackHandler
+
+if TYPE_CHECKING:
+    from .protocols import LLMClientProtocol
 
 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] = {}
+    _custom_clients: Dict[str, "LLMClientProtocol"] = {}
+
+    @classmethod
+    def register_custom_provider(cls, name: str, client: "LLMClientProtocol") -> None:
+        """
+        Register a custom LLM client provider.
+
+        This allows registration of custom LLM clients that implement the LLMClientProtocol
+        without inheriting from BaseLLMClient. Custom providers can be retrieved by name
+        using get_client().
+
+        Args:
+            name: Custom provider name (e.g., "my-llm", "llama-local", "custom-gpt")
+            client: Client implementing LLMClientProtocol
+
+        Raises:
+            ValueError: If client doesn't implement LLMClientProtocol
+            ValueError: If name conflicts with standard AIProvider enum values
+
+        Example:
+            ```python
+            # Register custom LLM client
+            custom_client = MyCustomLLMClient()
+            LLMClientFactory.register_custom_provider("my-llm", custom_client)
+
+            # Use custom client
+            client = LLMClientFactory.get_client("my-llm")
+            response = await client.generate_text(messages)
+            ```
+        """
+        # Import here to avoid circular dependency
+        from .protocols import LLMClientProtocol
+
+        # Validate protocol compliance
+        if not isinstance(client, LLMClientProtocol):
+            raise ValueError(
+                f"Client must implement LLMClientProtocol. "
+                f"Required methods: generate_text, stream_text, close, get_embeddings. "
+                f"Required attribute: provider_name"
+            )
+
+        # Prevent conflicts with standard provider names
+        try:
+            AIProvider(name)
+            raise ValueError(
+                f"Custom provider name '{name}' conflicts with standard AIProvider enum. "
+                f"Please use a different name."
+            )
+        except ValueError as e:
+            # If ValueError is raised because name is not in enum, that's good
+            if "conflicts with standard AIProvider" in str(e):
+                raise
+            # Otherwise, name is not in enum, proceed with registration
+
+        cls._custom_clients[name] = client
+        logger.info(f"Registered custom LLM provider: {name}")
 
     @classmethod
-    def get_client(cls, provider: Union[str, AIProvider]) -> BaseLLMClient:
-        """Get or create a client for the specified provider"""
+    def get_client(cls, provider: Union[str, AIProvider]) -> Union[BaseLLMClient, "LLMClientProtocol"]:
+        """
+        Get or create a client for the specified provider.
+
+        Supports both standard AIProvider enum values and custom provider names
+        registered via register_custom_provider().
+
+        Args:
+            provider: AIProvider enum or custom provider name string
+
+        Returns:
+            LLM client (BaseLLMClient for standard providers, LLMClientProtocol for custom)
+
+        Raises:
+            ValueError: If provider is unknown (not standard and not registered)
+        """
+        # Check custom providers first
+        if isinstance(provider, str) and provider in cls._custom_clients:
+            return cls._custom_clients[provider]
+
+        # Handle standard providers
         if isinstance(provider, str):
             try:
                 provider = AIProvider(provider)
             except ValueError:
-                raise ValueError(f"Unsupported provider: {provider}")
+                raise ValueError(
+                    f"Unknown provider: {provider}. "
+                    f"Standard providers: {[p.value for p in AIProvider]}. "
+                    f"Custom providers: {list(cls._custom_clients.keys())}. "
+                    f"Register custom providers with LLMClientFactory.register_custom_provider()"
+                )
 
         if provider not in cls._clients:
             cls._clients[provider] = cls._create_client(provider)
@@ -41,6 +129,8 @@ class LLMClientFactory:
             return OpenAIClient()
         elif provider == AIProvider.VERTEX:
             return VertexAIClient()
+        elif provider == AIProvider.GOOGLEAI:
+            return GoogleAIClient()
         elif provider == AIProvider.XAI:
             return XAIClient()
         else:
@@ -48,7 +138,8 @@ class LLMClientFactory:
 
     @classmethod
     async def close_all(cls):
-        """Close all active clients"""
+        """Close all active clients (both standard and custom)"""
+        # Close standard clients
         for client in cls._clients.values():
             try:
                 await client.close()
@@ -56,11 +147,34 @@ class LLMClientFactory:
                 logger.error(f"Error closing client {client.provider_name}: {e}")
         cls._clients.clear()
 
+        # Close custom clients
+        for name, client in cls._custom_clients.items():
+            try:
+                await client.close()
+            except Exception as e:
+                logger.error(f"Error closing custom client {name}: {e}")
+        cls._custom_clients.clear()
+
     @classmethod
     async def close_client(cls, provider: Union[str, AIProvider]):
-        """Close a specific client"""
+        """Close a specific client (standard or custom)"""
+        # Check if it's a custom provider
+        if isinstance(provider, str) and provider in cls._custom_clients:
+            try:
+                await cls._custom_clients[provider].close()
+                del cls._custom_clients[provider]
+                logger.info(f"Closed custom client: {provider}")
+            except Exception as e:
+                logger.error(f"Error closing custom client {provider}: {e}")
+            return
+
+        # Handle standard providers
         if isinstance(provider, str):
-            provider = AIProvider(provider)
+            try:
+                provider = AIProvider(provider)
+            except ValueError:
+                logger.warning(f"Unknown provider to close: {provider}")
+                return
 
         if provider in cls._clients:
             try:
@@ -69,6 +183,25 @@ class LLMClientFactory:
             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"""
 
@@ -80,19 +213,19 @@ class LLMClientManager:
         if not context:
             return None, None
 
-        metadata = context.get('metadata', {})
+        metadata = context.get("metadata", {})
 
         # First, check for aiPreference in metadata
-        ai_preference = metadata.get('aiPreference', {})
+        ai_preference = metadata.get("aiPreference", {})
         if isinstance(ai_preference, dict):
-            provider = ai_preference.get('provider')
-            model = ai_preference.get('model')
+            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')
+        provider = metadata.get("provider")
+        model = metadata.get("model")
         return provider, model
 
     async def generate_text(
@@ -104,7 +237,7 @@ class LLMClientManager:
         temperature: float = 0.7,
         max_tokens: Optional[int] = None,
         callbacks: Optional[List[CustomAsyncCallbackHandler]] = None,
-        **kwargs
+        **kwargs,
     ) -> LLMResponse:
         """
         Generate text using context-aware provider selection
@@ -139,7 +272,12 @@ class LLMClientManager:
             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)
+                    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}")
 
@@ -153,7 +291,7 @@ class LLMClientManager:
                 model=final_model,
                 temperature=temperature,
                 max_tokens=max_tokens,
-                **kwargs
+                **kwargs,
             )
 
             # Execute on_llm_end callbacks
@@ -167,11 +305,16 @@ class LLMClientManager:
                     "prompt_tokens": response.prompt_tokens,
                     "completion_tokens": response.completion_tokens,
                     "cost_estimate": response.cost_estimate,
-                    "response_time": response.response_time
+                    "response_time": response.response_time,
                 }
                 for callback in callbacks:
                     try:
-                        await callback.on_llm_end(response_dict, provider=final_provider, model=final_model, **kwargs)
+                        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}")
 
@@ -183,7 +326,12 @@ class LLMClientManager:
             if callbacks:
                 for callback in callbacks:
                     try:
-                        await callback.on_llm_error(e, provider=final_provider, model=final_model, **kwargs)
+                        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}")
 
@@ -199,7 +347,7 @@ class LLMClientManager:
         temperature: float = 0.7,
         max_tokens: Optional[int] = None,
         callbacks: Optional[List[CustomAsyncCallbackHandler]] = None,
-        **kwargs
+        **kwargs,
     ):
         """
         Stream text generation using context-aware provider selection
@@ -234,7 +382,12 @@ class LLMClientManager:
             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)
+                    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}")
 
@@ -246,14 +399,20 @@ class LLMClientManager:
             collected_content = ""
 
             # Stream text
-            async for chunk in await client.stream_text(
+            # Note: stream_text is an async generator, not a coroutine,
+            # so we iterate directly without await
+            async for chunk in client.stream_text(
                 messages=messages,
                 model=final_model,
                 temperature=temperature,
                 max_tokens=max_tokens,
-                **kwargs
+                **kwargs,
             ):
-                collected_content += chunk
+                # Handle StreamChunk objects (when return_chunks=True or function calling)
+                if hasattr(chunk, 'content') and chunk.content:
+                    collected_content += chunk.content
+                elif isinstance(chunk, str):
+                    collected_content += chunk
                 yield chunk
 
             # Create a response object for callbacks (streaming doesn't return LLMResponse directly)
@@ -263,7 +422,7 @@ class LLMClientManager:
                 content=collected_content,
                 provider=str(final_provider),
                 model=final_model or "unknown",
-                tokens_used=estimated_tokens
+                tokens_used=estimated_tokens,
             )
 
             # Execute on_llm_end callbacks
@@ -277,11 +436,16 @@ class LLMClientManager:
                     "prompt_tokens": stream_response.prompt_tokens,
                     "completion_tokens": stream_response.completion_tokens,
                     "cost_estimate": stream_response.cost_estimate,
-                    "response_time": stream_response.response_time
+                    "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)
+                        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}")
 
@@ -290,7 +454,12 @@ class LLMClientManager:
             if callbacks:
                 for callback in callbacks:
                     try:
-                        await callback.on_llm_error(e, provider=final_provider, model=final_model, **kwargs)
+                        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}")
 
@@ -301,14 +470,19 @@ class LLMClientManager:
         """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,
@@ -317,11 +491,21 @@ async def generate_text(
     temperature: float = 0.7,
     max_tokens: Optional[int] = None,
     callbacks: Optional[List[CustomAsyncCallbackHandler]] = None,
-    **kwargs
+    **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)
+    return await manager.generate_text(
+        messages,
+        provider,
+        model,
+        context,
+        temperature,
+        max_tokens,
+        callbacks,
+        **kwargs,
+    )
+
 
 async def stream_text(
     messages: Union[str, list[LLMMessage]],
@@ -331,9 +515,18 @@ async def stream_text(
     temperature: float = 0.7,
     max_tokens: Optional[int] = None,
     callbacks: Optional[List[CustomAsyncCallbackHandler]] = None,
-    **kwargs
+    **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):
+    async for chunk in manager.stream_text(
+        messages,
+        provider,
+        model,
+        context,
+        temperature,
+        max_tokens,
+        callbacks,
+        **kwargs,
+    ):
         yield chunk

aiecs/llm/client_resolver.py

@@ -0,0 +1,155 @@
+"""
+LLM Client Resolution Helper
+
+Provides convenient helper functions for resolving LLM clients from provider names
+with caching support for improved performance.
+"""
+
+import logging
+from typing import Optional, Union, Dict, TYPE_CHECKING
+from aiecs.llm.client_factory import LLMClientFactory, AIProvider
+
+if TYPE_CHECKING:
+    from aiecs.llm.protocols import LLMClientProtocol
+
+logger = logging.getLogger(__name__)
+
+
+# Cache for resolved clients to avoid repeated factory calls
+_client_cache: Dict[str, "LLMClientProtocol"] = {}
+
+
+def resolve_llm_client(
+    provider: Union[str, AIProvider],
+    model: Optional[str] = None,
+    use_cache: bool = True,
+) -> "LLMClientProtocol":
+    """
+    Resolve an LLM client from a provider name with optional caching.
+
+    This helper function provides a convenient interface for resolving LLM clients
+    from provider names, supporting both standard AIProvider enum values and custom
+    provider names registered via LLMClientFactory.register_custom_provider().
+
+    The function includes built-in caching to avoid repeated factory calls for the
+    same provider, improving performance when the same client is requested multiple times.
+
+    Args:
+        provider: AIProvider enum or custom provider name string
+        model: Optional model name (for logging/debugging purposes)
+        use_cache: Whether to use cached clients (default: True)
+
+    Returns:
+        LLM client implementing LLMClientProtocol
+
+    Raises:
+        ValueError: If provider is unknown (not standard and not registered)
+
+    Example:
+        ```python
+        from aiecs.llm.client_resolver import resolve_llm_client
+        from aiecs.llm.client_factory import LLMClientFactory
+
+        # Resolve standard provider
+        client = resolve_llm_client("OpenAI", model="gpt-4")
+
+        # Register and resolve custom provider
+        LLMClientFactory.register_custom_provider("my-llm", custom_client)
+        client = resolve_llm_client("my-llm", model="custom-model")
+
+        # Use the client
+        response = await client.generate_text(messages, model="gpt-4")
+        ```
+
+    Note:
+        - Caching is based on provider name only, not model name
+        - Custom providers registered after caching will require cache clearing
+        - Use `clear_client_cache()` to clear the cache if needed
+    """
+    # Convert provider to string for cache key
+    cache_key = str(provider) if isinstance(provider, AIProvider) else provider
+
+    # Check cache first if caching is enabled
+    if use_cache and cache_key in _client_cache:
+        logger.debug(f"Using cached client for provider: {cache_key}")
+        return _client_cache[cache_key]
+
+    # Resolve client from factory
+    try:
+        client = LLMClientFactory.get_client(provider)
+        
+        # Log resolution
+        if model:
+            logger.info(f"Resolved LLM client for provider: {cache_key}, model: {model}")
+        else:
+            logger.info(f"Resolved LLM client for provider: {cache_key}")
+
+        # Cache the client if caching is enabled
+        # Cast to LLMClientProtocol since BaseLLMClient implements the protocol
+        if use_cache:
+            from typing import cast
+            _client_cache[cache_key] = cast("LLMClientProtocol", client)
+            logger.debug(f"Cached client for provider: {cache_key}")
+
+        # Cast return value to match return type annotation
+        from typing import cast
+        return cast("LLMClientProtocol", client)
+
+    except ValueError as e:
+        logger.error(f"Failed to resolve LLM client for provider: {cache_key}")
+        raise
+
+
+def clear_client_cache(provider: Optional[Union[str, AIProvider]] = None) -> None:
+    """
+    Clear the client resolution cache.
+
+    Args:
+        provider: Optional provider to clear from cache.
+                 If None, clears entire cache.
+
+    Example:
+        ```python
+        from aiecs.llm.client_resolver import clear_client_cache
+
+        # Clear specific provider
+        clear_client_cache("OpenAI")
+
+        # Clear entire cache
+        clear_client_cache()
+        ```
+    """
+    global _client_cache
+
+    if provider is None:
+        # Clear entire cache
+        count = len(_client_cache)
+        _client_cache.clear()
+        logger.info(f"Cleared entire client cache ({count} entries)")
+    else:
+        # Clear specific provider
+        cache_key = str(provider) if isinstance(provider, AIProvider) else provider
+        if cache_key in _client_cache:
+            del _client_cache[cache_key]
+            logger.info(f"Cleared cache for provider: {cache_key}")
+        else:
+            logger.debug(f"Provider not in cache: {cache_key}")
+
+
+def get_cached_providers() -> list[str]:
+    """
+    Get list of providers currently in the cache.
+
+    Returns:
+        List of provider names that have cached clients
+
+    Example:
+        ```python
+        from aiecs.llm.client_resolver import get_cached_providers
+
+        providers = get_cached_providers()
+        print(f"Cached providers: {providers}")
+        ```
+    """
+    return list(_client_cache.keys())
+

aiecs/llm/clients/__init__.py

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