rakam-systems-agent 0.1.1rc7__py3-none-any.whl

rakam_systems_agent/components/llm_gateway/gateway_factory.py

@@ -0,0 +1,313 @@
+"""LLM Gateway Factory for provider routing and configuration-driven model selection."""
+from __future__ import annotations
+import os
+from typing import Any, Dict, Optional
+
+from rakam_systems_core.ai_utils import logging
+from rakam_systems_core.ai_core.interfaces.llm_gateway import LLMGateway
+from .openai_gateway import OpenAIGateway
+from .mistral_gateway import MistralGateway
+
+logger = logging.getLogger(__name__)
+
+
+class LLMGatewayFactory:
+    """Factory for creating LLM gateways based on provider and configuration.
+
+    This factory enables:
+    - Configuration-driven provider selection
+    - Automatic routing to the correct gateway
+    - Model string parsing (e.g., "openai:gpt-4o")
+    - Fallback to environment-based configuration
+
+    Example:
+        >>> # Using model string with provider prefix
+        >>> gateway = LLMGatewayFactory.create_gateway("openai:gpt-4o")
+        >>> 
+        >>> # Using explicit provider and model
+        >>> gateway = LLMGatewayFactory.create_gateway_from_config({
+        ...     "provider": "mistral",
+        ...     "model": "mistral-large-latest",
+        ...     "temperature": 0.7
+        ... })
+    """
+
+    # Registry of available providers
+    _PROVIDERS = {
+        "openai": OpenAIGateway,
+        "mistral": MistralGateway,
+    }
+
+    # Default models for each provider
+    _DEFAULT_MODELS = {
+        "openai": "gpt-4o",
+        "mistral": "mistral-large-latest",
+    }
+
+    @classmethod
+    def parse_model_string(cls, model_string: str) -> tuple[str, str]:
+        """Parse a model string into provider and model name.
+
+        Supports formats:
+        - "openai:gpt-4o" -> ("openai", "gpt-4o")
+        - "mistral:mistral-large-latest" -> ("mistral", "mistral-large-latest")
+        - "gpt-4o" -> ("openai", "gpt-4o")  # assumes OpenAI if no prefix
+
+        Args:
+            model_string: Model string to parse
+
+        Returns:
+            Tuple of (provider, model_name)
+
+        Raises:
+            ValueError: If provider is unknown
+        """
+        if ":" in model_string:
+            provider, model = model_string.split(":", 1)
+        else:
+            # Default to OpenAI if no provider specified
+            provider = "openai"
+            model = model_string
+
+        if provider not in cls._PROVIDERS:
+            raise ValueError(
+                f"Unknown provider '{provider}'. Supported providers: {list(cls._PROVIDERS.keys())}"
+            )
+
+        return provider, model
+
+    @classmethod
+    def create_gateway(
+        cls,
+        model_string: Optional[str] = None,
+        temperature: Optional[float] = None,
+        api_key: Optional[str] = None,
+        **kwargs: Any,
+    ) -> LLMGateway:
+        """Create an LLM gateway from a model string.
+
+        Args:
+            model_string: Model string (e.g., "openai:gpt-4o", "mistral:mistral-large-latest")
+                         Falls back to DEFAULT_LLM_MODEL env var
+            temperature: Temperature for generation
+            api_key: API key for the provider (provider-specific env vars used as fallback)
+            **kwargs: Additional parameters passed to gateway constructor
+
+        Returns:
+            Configured LLM gateway instance
+
+        Raises:
+            ValueError: If provider is unknown or configuration is invalid
+        """
+        # Get model string from parameter or environment
+        model_string = model_string or os.getenv(
+            "DEFAULT_LLM_MODEL", "openai:gpt-4o")
+
+        # Parse provider and model
+        provider, model = cls.parse_model_string(model_string)
+
+        # Get gateway class
+        gateway_class = cls._PROVIDERS[provider]
+
+        # Build gateway parameters
+        gateway_params = {
+            "model": model,
+            "default_temperature": temperature or float(os.getenv("DEFAULT_LLM_TEMPERATURE", "0.7")),
+        }
+
+        # Add API key if provided
+        if api_key:
+            gateway_params["api_key"] = api_key
+
+        # Add any additional parameters
+        gateway_params.update(kwargs)
+
+        logger.info(
+            f"Creating {provider} gateway with model={model}, temperature={gateway_params['default_temperature']}"
+        )
+
+        return gateway_class(**gateway_params)
+
+    @classmethod
+    def create_gateway_from_config(
+        cls,
+        config: Dict[str, Any],
+    ) -> LLMGateway:
+        """Create an LLM gateway from a configuration dictionary.
+
+        Expected config format:
+        {
+            "provider": "openai",  # or "mistral"
+            "model": "gpt-4o",
+            "temperature": 0.7,
+            "api_key": "...",  # optional
+            ... # additional provider-specific params
+        }
+
+        Args:
+            config: Configuration dictionary
+
+        Returns:
+            Configured LLM gateway instance
+
+        Raises:
+            ValueError: If provider is unknown or configuration is invalid
+        """
+        provider = config.get("provider")
+        model = config.get("model")
+
+        if not provider:
+            raise ValueError("Configuration must specify 'provider'")
+
+        if provider not in cls._PROVIDERS:
+            raise ValueError(
+                f"Unknown provider '{provider}'. Supported providers: {list(cls._PROVIDERS.keys())}"
+            )
+
+        # Use default model if not specified
+        if not model:
+            model = cls._DEFAULT_MODELS.get(provider)
+            logger.warning(
+                f"No model specified for {provider}, using default: {model}"
+            )
+
+        # Get gateway class
+        gateway_class = cls._PROVIDERS[provider]
+
+        # Extract common parameters
+        gateway_params = {
+            "model": model,
+            "default_temperature": config.get("temperature", 0.7),
+        }
+
+        # Add API key if provided
+        if "api_key" in config:
+            gateway_params["api_key"] = config["api_key"]
+
+        # Add provider-specific parameters
+        provider_specific_keys = {
+            "openai": ["base_url", "organization"],
+            "mistral": [],
+        }
+
+        for key in provider_specific_keys.get(provider, []):
+            if key in config:
+                gateway_params[key] = config[key]
+
+        logger.info(
+            f"Creating {provider} gateway from config with model={model}"
+        )
+
+        return gateway_class(**gateway_params)
+
+    @classmethod
+    def get_default_gateway(cls) -> LLMGateway:
+        """Get a default gateway based on environment configuration.
+
+        Checks environment variables:
+        - DEFAULT_LLM_PROVIDER: Provider name (default: "openai")
+        - DEFAULT_LLM_MODEL: Model name (default: "gpt-4o" for OpenAI)
+        - DEFAULT_LLM_TEMPERATURE: Temperature (default: 0.7)
+
+        Returns:
+            Configured default LLM gateway
+        """
+        provider = os.getenv("DEFAULT_LLM_PROVIDER", "openai")
+        model = os.getenv("DEFAULT_LLM_MODEL",
+                          cls._DEFAULT_MODELS.get(provider, "gpt-4o"))
+        temperature = float(os.getenv("DEFAULT_LLM_TEMPERATURE", "0.7"))
+
+        # If model doesn't have provider prefix, add it
+        if ":" not in model:
+            model_string = f"{provider}:{model}"
+        else:
+            model_string = model
+
+        return cls.create_gateway(
+            model_string=model_string,
+            temperature=temperature,
+        )
+
+    @classmethod
+    def register_provider(
+        cls,
+        provider_name: str,
+        gateway_class: type[LLMGateway],
+        default_model: Optional[str] = None,
+    ) -> None:
+        """Register a new provider gateway.
+
+        This allows extending the factory with custom providers.
+
+        Args:
+            provider_name: Name of the provider (e.g., "custom")
+            gateway_class: Gateway class implementing LLMGateway
+            default_model: Optional default model for this provider
+        """
+        cls._PROVIDERS[provider_name] = gateway_class
+        if default_model:
+            cls._DEFAULT_MODELS[provider_name] = default_model
+
+        logger.info(
+            f"Registered provider '{provider_name}' with gateway class {gateway_class.__name__}"
+        )
+
+    @classmethod
+    def list_providers(cls) -> list[str]:
+        """List all registered providers.
+
+        Returns:
+            List of provider names
+        """
+        return list(cls._PROVIDERS.keys())
+
+    @classmethod
+    def get_default_model(cls, provider: str) -> Optional[str]:
+        """Get the default model for a provider.
+
+        Args:
+            provider: Provider name
+
+        Returns:
+            Default model name or None if provider unknown
+        """
+        return cls._DEFAULT_MODELS.get(provider)
+
+
+# Convenience function for creating gateways
+def get_llm_gateway(
+    model: Optional[str] = None,
+    provider: Optional[str] = None,
+    temperature: Optional[float] = None,
+    **kwargs: Any,
+) -> LLMGateway:
+    """Convenience function to create an LLM gateway.
+
+    Args:
+        model: Model name or full model string (e.g., "gpt-4o" or "openai:gpt-4o")
+        provider: Provider name (optional if model string includes provider)
+        temperature: Temperature for generation
+        **kwargs: Additional parameters
+
+    Returns:
+        Configured LLM gateway
+
+    Example:
+        >>> gateway = get_llm_gateway(model="gpt-4o", provider="openai")
+        >>> gateway = get_llm_gateway(model="openai:gpt-4o")
+    """
+    if provider and model:
+        # Build model string from separate provider and model
+        model_string = f"{provider}:{model}"
+    elif model:
+        # Use model as-is (may already include provider)
+        model_string = model
+    else:
+        # Use default
+        model_string = None
+
+    return LLMGatewayFactory.create_gateway(
+        model_string=model_string,
+        temperature=temperature,
+        **kwargs,
+    )

rakam_systems_agent/components/llm_gateway/mistral_gateway.py

@@ -0,0 +1,287 @@
+"""Mistral LLM Gateway implementation with structured output support."""
+from __future__ import annotations
+import os
+from typing import Any, Dict, Iterator, Optional, Type, TypeVar
+
+from mistralai import Mistral
+from pydantic import BaseModel
+
+from rakam_systems_core.ai_utils import logging
+from rakam_systems_core.ai_core.interfaces.llm_gateway import LLMGateway, LLMRequest, LLMResponse
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T", bound=BaseModel)
+
+
+class MistralGateway(LLMGateway):
+    """Mistral LLM Gateway with support for structured outputs.
+
+    Features:
+    - Text generation
+    - Structured output using JSON mode
+    - Streaming support
+    - Token counting (approximate)
+    - Support for all Mistral models
+
+    Example:
+        >>> gateway = MistralGateway(model="mistral-large-latest", api_key="...")
+        >>> request = LLMRequest(
+        ...     system_prompt="You are a helpful assistant",
+        ...     user_prompt="What is AI?",
+        ...     temperature=0.7
+        ... )
+        >>> response = gateway.generate(request)
+        >>> print(response.content)
+    """
+
+    def __init__(
+        self,
+        name: str = "mistral_gateway",
+        config: Optional[Dict[str, Any]] = None,
+        model: str = "mistral-large-latest",
+        default_temperature: float = 0.7,
+        api_key: Optional[str] = None,
+    ):
+        """Initialize Mistral Gateway.
+
+        Args:
+            name: Gateway name
+            config: Configuration dictionary
+            model: Mistral model name (e.g., "mistral-large-latest", "mistral-small-latest")
+            default_temperature: Default temperature for generation
+            api_key: Mistral API key (falls back to MISTRAL_API_KEY env var)
+        """
+        super().__init__(
+            name=name,
+            config=config,
+            provider="mistral",
+            model=model,
+            default_temperature=default_temperature,
+            api_key=api_key or os.getenv("MISTRAL_API_KEY"),
+        )
+
+        if not self.api_key:
+            raise ValueError(
+                "Mistral API key must be provided via api_key parameter or MISTRAL_API_KEY environment variable"
+            )
+
+        # Initialize Mistral client
+        self.client = Mistral(api_key=self.api_key)
+
+        logger.info(
+            f"Initialized Mistral Gateway with model={self.model}, temperature={self.default_temperature}"
+        )
+
+    def _build_messages(self, request: LLMRequest) -> list[dict]:
+        """Build messages array from request."""
+        messages = []
+
+        if request.system_prompt:
+            messages.append({
+                "role": "system",
+                "content": request.system_prompt
+            })
+
+        messages.append({
+            "role": "user",
+            "content": request.user_prompt
+        })
+
+        return messages
+
+    def generate(self, request: LLMRequest) -> LLMResponse:
+        """Generate a response from Mistral.
+
+        Args:
+            request: Standardized LLM request
+
+        Returns:
+            Standardized LLM response
+        """
+        messages = self._build_messages(request)
+
+        # Prepare API call parameters
+        params = {
+            "model": self.model,
+            "messages": messages,
+            "temperature": request.temperature if request.temperature is not None else self.default_temperature,
+        }
+
+        if request.max_tokens:
+            params["max_tokens"] = request.max_tokens
+
+        # Add extra parameters
+        params.update(request.extra_params)
+
+        logger.debug(
+            f"Calling Mistral API with model={self.model}, temperature={params['temperature']}")
+
+        try:
+            response = self.client.chat.complete(**params)
+
+            # Extract response
+            content = response.choices[0].message.content
+
+            # Build usage information
+            usage = None
+            if response.usage:
+                usage = {
+                    "prompt_tokens": response.usage.prompt_tokens,
+                    "completion_tokens": response.usage.completion_tokens,
+                    "total_tokens": response.usage.total_tokens,
+                }
+
+            llm_response = LLMResponse(
+                content=content,
+                usage=usage,
+                model=response.model,
+                finish_reason=response.choices[0].finish_reason,
+                metadata={
+                    "id": response.id,
+                    "created": response.created,
+                }
+            )
+
+            logger.info(
+                f"Mistral response received: {usage.get('total_tokens', 'unknown') if usage else 'unknown'} tokens, "
+                f"finish_reason={llm_response.finish_reason}"
+            )
+
+            return llm_response
+
+        except Exception as e:
+            logger.error(f"Mistral API error: {str(e)}")
+            raise
+
+    def generate_structured(
+        self,
+        request: LLMRequest,
+        schema: Type[T],
+    ) -> T:
+        """Generate structured output conforming to a Pydantic schema.
+
+        Uses Mistral's JSON mode and parses the response into the schema.
+
+        Args:
+            request: Standardized LLM request
+            schema: Pydantic model class to parse response into
+
+        Returns:
+            Instance of the schema class
+        """
+        import json
+
+        messages = self._build_messages(request)
+
+        # Add schema information to the system prompt
+        schema_json = schema.model_json_schema()
+
+        # Enhance system prompt with schema information
+        enhanced_system = request.system_prompt or ""
+        enhanced_system += f"\n\nYou must respond with valid JSON that matches this schema:\n{json.dumps(schema_json, indent=2)}"
+
+        # Update messages with enhanced system prompt
+        messages = []
+        messages.append({
+            "role": "system",
+            "content": enhanced_system
+        })
+        messages.append({
+            "role": "user",
+            "content": request.user_prompt
+        })
+
+        # Prepare API call parameters
+        params = {
+            "model": self.model,
+            "messages": messages,
+            "temperature": request.temperature if request.temperature is not None else self.default_temperature,
+            "response_format": {"type": "json_object"},
+        }
+
+        if request.max_tokens:
+            params["max_tokens"] = request.max_tokens
+
+        # Add extra parameters
+        params.update(request.extra_params)
+
+        logger.debug(
+            f"Calling Mistral API for structured output with model={self.model}, schema={schema.__name__}"
+        )
+
+        try:
+            response = self.client.chat.complete(**params)
+
+            # Extract and parse JSON response
+            content = response.choices[0].message.content
+            parsed_result = schema.model_validate_json(content)
+
+            logger.info(
+                f"Mistral structured response received: schema={schema.__name__}"
+            )
+
+            return parsed_result
+
+        except Exception as e:
+            logger.error(f"Mistral structured output error: {str(e)}")
+            raise
+
+    def stream(self, request: LLMRequest) -> Iterator[str]:
+        """Stream responses from Mistral.
+
+        Args:
+            request: Standardized LLM request
+
+        Yields:
+            String chunks from the LLM
+        """
+        messages = self._build_messages(request)
+
+        # Prepare API call parameters
+        params = {
+            "model": self.model,
+            "messages": messages,
+            "temperature": request.temperature if request.temperature is not None else self.default_temperature,
+        }
+
+        if request.max_tokens:
+            params["max_tokens"] = request.max_tokens
+
+        # Add extra parameters
+        params.update(request.extra_params)
+
+        logger.debug(f"Streaming from Mistral with model={self.model}")
+
+        try:
+            stream = self.client.chat.stream(**params)
+
+            for chunk in stream:
+                if chunk.data.choices[0].delta.content is not None:
+                    yield chunk.data.choices[0].delta.content
+
+        except Exception as e:
+            logger.error(f"Mistral streaming error: {str(e)}")
+            raise
+
+    def count_tokens(self, text: str, model: Optional[str] = None) -> int:
+        """Count tokens in text.
+
+        Mistral doesn't provide a native tokenization library, so we use approximation.
+
+        Args:
+            text: Text to count tokens for
+            model: Model name (unused for Mistral)
+
+        Returns:
+            Approximate number of tokens in the text
+        """
+        # Approximation: average of 4 characters per token
+        # This is less accurate than tiktoken but reasonable for most use cases
+        token_count = len(text) // 4
+
+        logger.debug(
+            f"Counted ~{token_count} tokens (approximation) for text of length {len(text)} characters"
+        )
+
+        return token_count