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

aiecs/llm/clients/xai_client.py

@@ -0,0 +1,201 @@
+from openai import AsyncOpenAI
+from aiecs.config.config import get_settings
+from aiecs.llm.clients.base_client import (
+    BaseLLMClient,
+    LLMMessage,
+    LLMResponse,
+    ProviderNotAvailableError,
+    RateLimitError,
+)
+from aiecs.llm.clients.openai_compatible_mixin import (
+    OpenAICompatibleFunctionCallingMixin,
+    StreamChunk,
+)
+from tenacity import (
+    retry,
+    stop_after_attempt,
+    wait_exponential,
+    retry_if_exception_type,
+)
+import logging
+from typing import Dict, Optional, List, AsyncGenerator, cast, Any
+
+# 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()
+
+
+logger = logging.getLogger(__name__)
+
+
+class XAIClient(BaseLLMClient, OpenAICompatibleFunctionCallingMixin):
+    """xAI (Grok) provider client"""
+
+    def __init__(self) -> None:
+        super().__init__("xAI")
+        self.settings = get_settings()
+        self._openai_client: Optional[AsyncOpenAI] = None
+        self._model_map: Optional[Dict[str, str]] = None
+
+    def _get_openai_client(self) -> AsyncOpenAI:
+        """Lazy initialization of OpenAI client for XAI"""
+        if not self._openai_client:
+            api_key = self._get_api_key()
+            self._openai_client = AsyncOpenAI(
+                api_key=api_key,
+                base_url="https://api.x.ai/v1",
+                timeout=360.0,  # Override default timeout with longer timeout for reasoning models
+            )
+        return self._openai_client
+
+    def _get_api_key(self) -> str:
+        """Get API key with backward compatibility"""
+        # Support both xai_api_key and grok_api_key for backward compatibility
+        api_key = getattr(self.settings, "xai_api_key", None) or getattr(self.settings, "grok_api_key", None)
+        if not api_key:
+            raise ProviderNotAvailableError("xAI API key not configured")
+        return api_key
+
+    def _get_model_map(self) -> Dict[str, str]:
+        """Get model mappings from configuration"""
+        if self._model_map is None:
+            try:
+                loader = _get_config_loader()
+                provider_config = loader.get_provider_config("xAI")
+                if provider_config and provider_config.model_mappings:
+                    self._model_map = provider_config.model_mappings
+                else:
+                    self._model_map = {}
+            except Exception as e:
+                self.logger.warning(f"Failed to load model mappings from config: {e}")
+                self._model_map = {}
+        return self._model_map
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential(multiplier=1, min=4, max=10),
+        retry=retry_if_exception_type((Exception, RateLimitError)),
+    )
+    async def generate_text(
+        self,
+        messages: List[LLMMessage],
+        model: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        functions: Optional[List[Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: Optional[Any] = None,
+        **kwargs,
+    ) -> LLMResponse:
+        """
+        Generate text using xAI API via OpenAI library (supports all Grok models).
+        
+        xAI API is OpenAI-compatible, so it supports Function Calling.
+        """
+        # Check API key availability
+        api_key = self._get_api_key()
+        if not api_key:
+            raise ProviderNotAvailableError("xAI API key is not configured.")
+
+        client = self._get_openai_client()
+
+        # Get model name from config if not provided
+        selected_model = model or self._get_default_model() or "grok-4"
+
+        # Get model mappings from config
+        model_map = self._get_model_map()
+        api_model = model_map.get(selected_model, selected_model)
+
+        try:
+            # Use mixin method for Function Calling support
+            response = await self._generate_text_with_function_calling(
+                client=client,
+                messages=messages,
+                model=api_model,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                functions=functions,
+                tools=tools,
+                tool_choice=tool_choice,
+                **kwargs,
+            )
+            
+            # Override provider and model name for xAI
+            response.provider = self.provider_name
+            response.model = selected_model
+            response.cost_estimate = 0.0  # xAI pricing not available yet
+            
+            return response
+
+        except Exception as e:
+            if "rate limit" in str(e).lower() or "429" in str(e):
+                raise RateLimitError(f"xAI rate limit exceeded: {str(e)}")
+            logger.error(f"xAI API error: {str(e)}")
+            raise
+
+    async def stream_text(  # type: ignore[override]
+        self,
+        messages: List[LLMMessage],
+        model: Optional[str] = None,
+        temperature: float = 0.7,
+        max_tokens: Optional[int] = None,
+        functions: Optional[List[Dict[str, Any]]] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,
+        tool_choice: Optional[Any] = None,
+        return_chunks: bool = False,
+        **kwargs,
+    ) -> AsyncGenerator[Any, None]:
+        """
+        Stream text using xAI API via OpenAI library (supports all Grok models).
+        
+        xAI API is OpenAI-compatible, so it supports Function Calling.
+        
+        Args:
+            return_chunks: If True, returns StreamChunk objects with tool_calls info; if False, returns str tokens only
+        """
+        # Check API key availability
+        api_key = self._get_api_key()
+        if not api_key:
+            raise ProviderNotAvailableError("xAI API key is not configured.")
+
+        client = self._get_openai_client()
+
+        # Get model name from config if not provided
+        selected_model = model or self._get_default_model() or "grok-4"
+
+        # Get model mappings from config
+        model_map = self._get_model_map()
+        api_model = model_map.get(selected_model, selected_model)
+
+        try:
+            # Use mixin method for Function Calling support
+            async for chunk in self._stream_text_with_function_calling(
+                client=client,
+                messages=messages,
+                model=api_model,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                functions=functions,
+                tools=tools,
+                tool_choice=tool_choice,
+                return_chunks=return_chunks,
+                **kwargs,
+            ):
+                yield chunk
+
+        except Exception as e:
+            if "rate limit" in str(e).lower() or "429" in str(e):
+                raise RateLimitError(f"xAI rate limit exceeded: {str(e)}")
+            logger.error(f"xAI API streaming error: {str(e)}")
+            raise
+
+    async def close(self):
+        """Clean up resources"""
+        if self._openai_client:
+            await self._openai_client.close()
+            self._openai_client = None

aiecs/llm/config/__init__.py

@@ -0,0 +1,51 @@
+"""
+LLM Configuration management.
+
+This package handles configuration loading, validation, and management
+for all LLM providers and models.
+"""
+
+from .model_config import (
+    ModelCostConfig,
+    ModelCapabilities,
+    ModelDefaultParams,
+    ModelConfig,
+    ProviderConfig,
+    LLMModelsConfig,
+)
+from .config_loader import (
+    LLMConfigLoader,
+    get_llm_config_loader,
+    get_llm_config,
+    reload_llm_config,
+)
+from .config_validator import (
+    ConfigValidationError,
+    validate_cost_config,
+    validate_model_config,
+    validate_provider_config,
+    validate_llm_config,
+    validate_config_file,
+)
+
+__all__ = [
+    # Configuration models
+    "ModelCostConfig",
+    "ModelCapabilities",
+    "ModelDefaultParams",
+    "ModelConfig",
+    "ProviderConfig",
+    "LLMModelsConfig",
+    # Config loader
+    "LLMConfigLoader",
+    "get_llm_config_loader",
+    "get_llm_config",
+    "reload_llm_config",
+    # Validation
+    "ConfigValidationError",
+    "validate_cost_config",
+    "validate_model_config",
+    "validate_provider_config",
+    "validate_llm_config",
+    "validate_config_file",
+]

aiecs/llm/config/config_loader.py

@@ -0,0 +1,272 @@
+"""
+Configuration loader for LLM models.
+
+This module provides a singleton configuration loader that loads and manages
+LLM model configurations from YAML files with support for hot-reloading.
+"""
+
+import logging
+import os
+from pathlib import Path
+from typing import Optional
+import yaml
+from threading import Lock
+
+from aiecs.llm.config.model_config import (
+    LLMModelsConfig,
+    ProviderConfig,
+    ModelConfig,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class LLMConfigLoader:
+    """
+    Singleton configuration loader for LLM models.
+
+    Supports:
+    - Loading configuration from YAML files
+    - Hot-reloading (manual refresh)
+    - Thread-safe access
+    - Caching for performance
+    """
+
+    _instance: Optional["LLMConfigLoader"] = None
+    _lock = Lock()
+    _config_lock = Lock()
+    _initialized: bool = False
+
+    def __new__(cls):
+        """Ensure singleton instance"""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = super().__new__(cls)
+                    cls._instance._initialized = False
+        return cls._instance
+
+    def __init__(self) -> None:
+        """Initialize the configuration loader"""
+        if self._initialized:
+            return
+
+        self._config: Optional[LLMModelsConfig] = None
+        self._config_path: Optional[Path] = None
+        self._initialized = True
+        logger.info("LLMConfigLoader initialized")
+
+    def _find_config_file(self) -> Path:
+        """
+        Find the configuration file.
+
+        Search order:
+        1. Settings llm_models_config_path
+        2. Environment variable LLM_MODELS_CONFIG
+        3. aiecs/config/llm_models.yaml
+        4. config/llm_models.yaml
+        """
+        # Check settings first
+        try:
+            from aiecs.config.config import get_settings
+
+            settings = get_settings()
+            if settings.llm_models_config_path:
+                path = Path(settings.llm_models_config_path)
+                if path.exists():
+                    logger.info(f"Using LLM config from settings: {path}")
+                    return path
+                else:
+                    logger.warning(f"Settings llm_models_config_path does not exist: {path}")
+        except Exception as e:
+            logger.debug(f"Could not load settings: {e}")
+
+        # Check environment variable
+        env_path = os.environ.get("LLM_MODELS_CONFIG")
+        if env_path:
+            path = Path(env_path)
+            if path.exists():
+                logger.info(f"Using LLM config from environment: {path}")
+                return path
+            else:
+                logger.warning(f"LLM_MODELS_CONFIG path does not exist: {path}")
+
+        # Check standard locations
+        current_dir = Path(__file__).parent.parent  # aiecs/
+
+        # Try aiecs/config/llm_models.yaml
+        config_path1 = current_dir / "config" / "llm_models.yaml"
+        if config_path1.exists():
+            logger.info(f"Using LLM config from: {config_path1}")
+            return config_path1
+
+        # Try config/llm_models.yaml (relative to project root)
+        config_path2 = current_dir.parent / "config" / "llm_models.yaml"
+        if config_path2.exists():
+            logger.info(f"Using LLM config from: {config_path2}")
+            return config_path2
+
+        # Default to the first path even if it doesn't exist
+        logger.warning(f"LLM config file not found, using default path: {config_path1}")
+        return config_path1
+
+    def load_config(self, config_path: Optional[Path] = None) -> LLMModelsConfig:
+        """
+        Load configuration from YAML file.
+
+        Args:
+            config_path: Optional path to configuration file. If not provided,
+                        will search in standard locations.
+
+        Returns:
+            LLMModelsConfig: Loaded and validated configuration
+
+        Raises:
+            FileNotFoundError: If config file doesn't exist
+            ValueError: If config file is invalid
+        """
+        with self._config_lock:
+            if config_path is None:
+                config_path = self._find_config_file()
+            else:
+                config_path = Path(config_path)
+
+            if not config_path.exists():
+                raise FileNotFoundError(f"LLM models configuration file not found: {config_path}\n" f"Please create the configuration file or set LLM_MODELS_CONFIG environment variable.")
+
+            try:
+                with open(config_path, "r", encoding="utf-8") as f:
+                    config_data = yaml.safe_load(f)
+
+                if not config_data:
+                    raise ValueError("Configuration file is empty")
+
+                # Validate and parse using Pydantic
+                self._config = LLMModelsConfig(**config_data)
+                self._config_path = config_path
+
+                logger.info(f"Loaded LLM configuration from {config_path}: " f"{len(self._config.providers)} providers, " f"{sum(len(p.models) for p in self._config.providers.values())} models")
+
+                return self._config
+
+            except yaml.YAMLError as e:
+                raise ValueError(f"Invalid YAML in configuration file: {e}")
+            except Exception as e:
+                raise ValueError(f"Failed to load configuration: {e}")
+
+    def reload_config(self) -> LLMModelsConfig:
+        """
+        Reload configuration from the current config file.
+
+        This supports the hybrid loading mode - configuration is loaded at startup
+        but can be manually refreshed without restarting the application.
+
+        Returns:
+            LLMModelsConfig: Reloaded configuration
+        """
+        logger.info("Reloading LLM configuration...")
+        return self.load_config(self._config_path)
+
+    def get_config(self) -> LLMModelsConfig:
+        """
+        Get the current configuration.
+
+        Loads configuration on first access if not already loaded.
+
+        Returns:
+            LLMModelsConfig: Current configuration
+        """
+        if self._config is None:
+            self.load_config()
+        # After load_config(), _config should never be None
+        if self._config is None:
+            raise RuntimeError("Failed to load LLM configuration")
+        return self._config
+
+    def get_provider_config(self, provider_name: str) -> Optional[ProviderConfig]:
+        """
+        Get configuration for a specific provider.
+
+        Args:
+            provider_name: Name of the provider (e.g., "Vertex", "OpenAI")
+
+        Returns:
+            ProviderConfig if found, None otherwise
+        """
+        config = self.get_config()
+        return config.get_provider_config(provider_name)
+
+    def get_model_config(self, provider_name: str, model_name: str) -> Optional[ModelConfig]:
+        """
+        Get configuration for a specific model.
+
+        Args:
+            provider_name: Name of the provider
+            model_name: Name of the model (or alias)
+
+        Returns:
+            ModelConfig if found, None otherwise
+        """
+        config = self.get_config()
+        return config.get_model_config(provider_name, model_name)
+
+    def get_default_model(self, provider_name: str) -> Optional[str]:
+        """
+        Get the default model name for a provider.
+
+        Args:
+            provider_name: Name of the provider
+
+        Returns:
+            Default model name if found, None otherwise
+        """
+        provider_config = self.get_provider_config(provider_name)
+        if provider_config:
+            return provider_config.default_model
+        return None
+
+    def is_loaded(self) -> bool:
+        """Check if configuration has been loaded"""
+        return self._config is not None
+
+    def get_config_path(self) -> Optional[Path]:
+        """Get the path to the current configuration file"""
+        return self._config_path
+
+
+# Global singleton instance
+_loader = LLMConfigLoader()
+
+
+def get_llm_config_loader() -> LLMConfigLoader:
+    """
+    Get the global LLM configuration loader instance.
+
+    Returns:
+        LLMConfigLoader: Global singleton instance
+    """
+    return _loader
+
+
+def get_llm_config() -> LLMModelsConfig:
+    """
+    Get the current LLM configuration.
+
+    Convenience function that returns the configuration from the global loader.
+
+    Returns:
+        LLMModelsConfig: Current configuration
+    """
+    return _loader.get_config()
+
+
+def reload_llm_config() -> LLMModelsConfig:
+    """
+    Reload the LLM configuration.
+
+    Convenience function that reloads the configuration in the global loader.
+
+    Returns:
+        LLMModelsConfig: Reloaded configuration
+    """
+    return _loader.reload_config()