llms.py

"""LLM providers for datafast using LiteLLM.

This module provides classes for different LLM providers (OpenAI, Anthropic, Gemini, Mistral)
with a unified interface using LiteLLM under the hood.
"""

from typing import Any, Type, TypeVar
from abc import ABC, abstractmethod
import os
import time
import traceback
import warnings
from loguru import logger

# Pydantic
from pydantic import BaseModel

# LiteLLM
import litellm
from litellm.exceptions import RateLimitError

# Internal imports
from .llm_utils import get_messages
from .tracing import (
    build_trace_metadata,
    load_env_once,
    maybe_configure_langfuse_tracing,
)

# Type aliases for Python 3.10+
Message = dict[str, str]
Messages = list[Message]
T = TypeVar('T', bound=BaseModel)


class LLMProvider(ABC):
    """Abstract base class for LLM providers."""

    def __init__(
        self,
        model_id: str,
        api_key: str | None = None,
        temperature: float | None = None,
        max_completion_tokens: int | None = None,
        top_p: float | None = None,
        frequency_penalty: float | None = None,
        rpm_limit: int | None = None,
        timeout: int | None = None,
        ):
        """Initialize the LLM provider with common parameters.

        Args:
            model_id: The model identifier
            api_key: API key (if None, will get from environment)
            temperature: The sampling temperature to be used, between 0 and 2. Higher values like 0.8 produce more random outputs, while lower values like 0.2 make outputs more focused and deterministic
            max_completion_tokens: An upper bound for the number of tokens that can be generated for a completion, including visible output tokens and reasoning tokens.
            top_p: Nucleus sampling parameter (0.0 to 1.0)
            frequency_penalty: Penalty for token frequency (-2.0 to 2.0)
        """
        self.model_id = model_id
        load_env_once()
        maybe_configure_langfuse_tracing(load_env=False)
        self.api_key = api_key or self._get_api_key()

        # Set generation parameters
        self.temperature = temperature
        self.max_completion_tokens = max_completion_tokens
        self.top_p = top_p
        self.frequency_penalty = frequency_penalty

        # Rate limiting
        self.rpm_limit = rpm_limit
        self._request_timestamps: list[float] = []

        # timeout
        self.timeout = timeout

        # Configure environment with API key if needed
        self._configure_env()
        # Log successful initialization
        logger.info(f"Initialized {self.provider_name} | Model: {self.model_id}")

    def _build_request_metadata(
        self,
        metadata: dict[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Build default tracing metadata for provider-level calls."""
        return build_trace_metadata(
            model=self,
            component="provider.generate",
            trace_name=f"datafast.{self.provider_name}",
            metadata=metadata,
        )

    @property
    @abstractmethod
    def provider_name(self) -> str:
        """Return the provider name used by LiteLLM."""
        pass

    @property
    @abstractmethod
    def env_key_name(self) -> str:
        """Return the environment variable name for API key."""
        pass

    def _get_api_key(self) -> str:
        """Get API key from environment variables."""
        api_key = os.getenv(self.env_key_name)
        if not api_key:
            logger.error(
                f"Missing API key | Set {self.env_key_name} environment variable"
            )
            raise ValueError(
                f"{self.env_key_name} environment variable not set. "
                f"Please set it or provide an API key when initializing the provider."
            )
        return api_key

    def _configure_env(self) -> None:
        """Configure environment variables for API key."""
        if self.api_key:
            os.environ[self.env_key_name] = self.api_key

    def _get_model_string(self) -> str:
        """Get the full model string for LiteLLM."""
        return f"{self.provider_name}/{self.model_id}"

    def _respect_rate_limit(self) -> None:
        """Block execution to ensure we do not exceed the rpm_limit."""
        if self.rpm_limit is None:
            return
        current = time.monotonic()
        # Keep only timestamps within the last minute
        self._request_timestamps = [
            ts for ts in self._request_timestamps if current - ts < 60]
        
        # Be more conservative - wait if we're at 90% of the limit
        conservative_limit = max(1, int(self.rpm_limit * 0.9))
        
        if len(self._request_timestamps) < conservative_limit:
            return
        
        # Need to wait until the earliest request is outside the 60-second window
        earliest = self._request_timestamps[0]
        # Add a 2s margin to avoid accidental rate limit exceedance
        sleep_time = 62 - (current - earliest)
        if sleep_time > 0:
            logger.warning(
                f"Rate limit approaching | Requests: {len(self._request_timestamps)}/{self.rpm_limit} | "
                f"Waiting {sleep_time:.1f}s"
            )
            time.sleep(sleep_time)
            # Clean up old timestamps after waiting
            current = time.monotonic()
            self._request_timestamps = [
                ts for ts in self._request_timestamps if current - ts < 60]

    @staticmethod
    def _strip_code_fences(content: str) -> str:
        """Strip markdown code fences from content if present.
        
        Args:
            content: The content string that may contain code fences
            
        Returns:
            Content with code fences removed
        """
        if not content:
            return content
        
        content = content.strip()
        
        # Check for code fences with optional language identifier
        if content.startswith('```'):
            # Find the end of the first line (language identifier)
            first_newline = content.find('\n')
            if first_newline != -1:
                content = content[first_newline + 1:]
            else:
                # No newline after opening fence, remove just the fence
                content = content[3:]
        
        # Remove closing fence
        if content.endswith('```'):
            content = content[:-3]
        
        return content.strip()

    def generate(
        self,
        prompt: str | list[str] | None = None,
        messages: list[Messages] | Messages | None = None,
        response_format: Type[T] | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> str | list[str] | T | list[T]:
        """
        Generate responses from the LLM using single or batch inference.

        Args:
            prompt: Single text prompt (str) or list of text prompts for batch processing
            messages: Single message list or list of message lists for batch processing
            response_format: Optional Pydantic model class for structured output
            metadata: Optional LiteLLM metadata for tracing / observability

        Returns:
            Single string/model or list of strings/models depending on input type.

        Raises:
            ValueError: If neither prompt nor messages is provided, or if both are provided.
            RuntimeError: If there's an error during generation.
        """
        # Validate inputs
        if prompt is None and messages is None:
            raise ValueError("Either prompts or messages must be provided")
        if prompt is not None and messages is not None:
            raise ValueError("Provide either prompts or messages, not both")

        # Determine if this is a single input or batch input
        single_input = False
        batch_prompts = None
        batch_messages = None

        if prompt is not None:
            if isinstance(prompt, str):
                # Single prompt - convert to batch
                batch_prompts = [prompt]
                single_input = True
            elif isinstance(prompt, list):
                # Already a list of prompts
                batch_prompts = prompt
                single_input = False
            else:
                raise ValueError("prompt must be a string or list of strings")

        if messages is not None:
            if isinstance(messages, list) and len(messages) > 0:
                # Check if it's a single message list or batch
                if isinstance(messages[0], dict):
                    # Single message list - convert to batch
                    batch_messages = [messages]
                    single_input = True
                elif isinstance(messages[0], list):
                    # Already a batch of message lists
                    batch_messages = messages
                    single_input = False
                else:
                    raise ValueError("Invalid messages format")
            else:
                raise ValueError("messages cannot be empty")

        try:
            # Append JSON formatting instructions if response_format is provided
            json_instructions = (
                "\nReturn only valid JSON. To do so, don't include ```json ``` markdown "
                "or code fences around the JSON. Use double quotes for all keys and values. "
                "Escape internal quotes and newlines (use \\n). Do not include trailing commas."
            )
            
            # Convert batch prompts to messages if needed
            batch_to_send = []
            if batch_prompts is not None:
                for one_prompt in batch_prompts:
                    # Append JSON instructions to prompt if response_format is provided
                    modified_prompt = one_prompt + json_instructions if response_format is not None else one_prompt
                    batch_to_send.append(get_messages(modified_prompt))
            else:
                batch_to_send = batch_messages
                # Append JSON instructions to the last user message if response_format is provided
                if response_format is not None:
                    for message_list in batch_to_send:
                        for msg in reversed(message_list):
                            if msg.get("role") == "user":
                                msg["content"] += json_instructions
                                break

            # Enforce rate limit per batch
            self._respect_rate_limit()

            # Prepare completion parameters for batch
            completion_params = {
                "model": self._get_model_string(),
                "messages": batch_to_send,
                "temperature": self.temperature,
                "max_tokens": self.max_completion_tokens,
                "top_p": self.top_p,
                "frequency_penalty": self.frequency_penalty,
                "timeout": self.timeout,
                "metadata": self._build_request_metadata(metadata),
            }
            if response_format is not None:
                completion_params["response_format"] = response_format

            # Call LiteLLM completion with retry on rate limit.
            # OpenRouter accepts single message requests via completion(), but
            # rejects the same payload when wrapped in batch_completion().
            max_retries = 3
            retry_delay = 5  # Start with 5 seconds
            response = None

            for attempt in range(max_retries):
                try:
                    if len(batch_to_send) == 1:
                        response = [litellm.completion(
                            **{**completion_params, "messages": batch_to_send[0]}
                        )]
                    else:
                        response = litellm.batch_completion(**completion_params)
                    break  # Success, exit retry loop
                except RateLimitError:
                    if attempt < max_retries - 1:
                        wait_time = retry_delay * (2 ** attempt)  # Exponential backoff
                        logger.warning(
                            f"Rate limit hit | Provider: {self.provider_name} | Model: {self.model_id} | "
                            f"Attempt {attempt + 1}/{max_retries} | Waiting {wait_time}s before retry"
                        )
                        time.sleep(wait_time)
                    else:
                        logger.error(
                            f"Rate limit exceeded after {max_retries} attempts | "
                            f"Provider: {self.provider_name} | Model: {self.model_id}"
                        )
                        raise

            if response is None:
                raise RuntimeError("Failed to get response after retries")

            # Record timestamp for rate limiting (one timestamp per batch item)
            if self.rpm_limit is not None:
                current_time = time.monotonic()
                for _ in range(len(batch_to_send)):
                    self._request_timestamps.append(current_time)

            # Extract content from each response
            results = []
            for idx, one_response in enumerate(response):
                if isinstance(one_response, Exception):
                    if isinstance(one_response, RateLimitError):
                        logger.warning(
                            "Rate limit error in batch item | Provider: %s | Model: %s | Item: %d",
                            self.provider_name,
                            self.model_id,
                            idx,
                        )
                    raise RuntimeError(
                        f"Batch item {idx} failed during generation: {one_response}"
                    ) from one_response

                if not getattr(one_response, "choices", None):
                    raise RuntimeError(
                        f"Unexpected response type from LiteLLM batch completion at item {idx}: {type(one_response).__name__}"
                    )

                content = one_response.choices[0].message.content
                
                if response_format is not None:
                    # Strip code fences before validation
                    content = self._strip_code_fences(content)
                    try:
                        results.append(
                            response_format.model_validate_json(content))
                    except Exception as validation_error:
                        # Show the content that failed to parse for debugging
                        content_preview = content[:200] + "..." if len(content) > 200 else content
                        logger.warning(
                            f"JSON parsing failed, skipping response | "
                            f"Model: {self.model_id} | "
                            f"Format: {response_format.__name__} | "
                            f"Content preview: {content_preview}"
                        )
                        raise ValueError(
                            f"Failed to parse JSON response into {response_format.__name__}.\n"
                            f"Validation error: {validation_error}\n"
                            f"Content received (first 200 chars):\n{content_preview}"
                        ) from validation_error
                else:
                    # Strip leading/trailing whitespace for text responses
                    results.append(content.strip() if content else content)

            # Return single result for backward compatibility
            if single_input and len(results) == 1:
                return results[0]
            return results

        except Exception as e:
            error_trace = traceback.format_exc()
            logger.error(
                f"Generation failed | Provider: {self.provider_name} | "
                f"Model: {self.model_id} | Error: {str(e)}"
            )
            raise RuntimeError(
                f"Error generating batch response with {self.provider_name}:\n{error_trace}"
            )


class OpenAIProvider(LLMProvider):
    """OpenAI provider using litellm.responses endpoint.
    
    Note: This provider uses the new responses endpoint which has different
    parameter support compared to the standard completion endpoint:
    - temperature, top_p, and frequency_penalty are not supported
    - Uses text_format instead of response_format
    - Supports reasoning parameter for controlling reasoning effort
    - Does not support batch operations (will process sequentially with warning)
    """

    @property
    def provider_name(self) -> str:
        return "openai"

    @property
    def env_key_name(self) -> str:
        return "OPENAI_API_KEY"

    def __init__(
        self,
        model_id: str = "gpt-5-mini-2025-08-07",
        api_key: str | None = None,
        max_completion_tokens: int | None = None,
        reasoning_effort: str = "low",
        temperature: float | None = None,
        top_p: float | None = None,
        frequency_penalty: float | None = None,
        timeout: int | None = None,
    ):
        """Initialize the OpenAI provider.

        Args:
            model_id: The model ID (defaults to gpt-5-mini)
            api_key: API key (if None, will get from environment)
            max_completion_tokens: An upper bound for the number of tokens that can be generated for a completion, including visible output tokens and reasoning tokens.
            reasoning_effort: Reasoning effort level - "low", "medium", or "high" (defaults to "low")
            temperature: DEPRECATED - Not supported by responses endpoint
            top_p: DEPRECATED - Not supported by responses endpoint
            frequency_penalty: DEPRECATED - Not supported by responses endpoint
            timeout: Request timeout in seconds
        """
        # Warn about deprecated parameters
        if temperature is not None:
            warnings.warn(
                "temperature parameter is not supported by OpenAI responses endpoint and will be ignored",
                UserWarning,
                stacklevel=2
            )
        if top_p is not None:
            warnings.warn(
                "top_p parameter is not supported by OpenAI responses endpoint and will be ignored",
                UserWarning,
                stacklevel=2
            )
        if frequency_penalty is not None:
            warnings.warn(
                "frequency_penalty parameter is not supported by OpenAI responses endpoint and will be ignored",
                UserWarning,
                stacklevel=2
            )
        
        # Store reasoning effort
        self.reasoning_effort = reasoning_effort
        
        # Call parent init with None for unsupported params
        super().__init__(
            model_id=model_id,
            api_key=api_key,
            temperature=None,
            max_completion_tokens=max_completion_tokens,
            top_p=None,
            frequency_penalty=None,
            timeout=timeout,
        )
    
    def generate(
        self,
        prompt: str | list[str] | None = None,
        messages: list[Messages] | Messages | None = None,
        response_format: Type[T] | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> str | list[str] | T | list[T]:
        """
        Generate responses from the LLM using the responses endpoint.
        
        Note: Batch operations are processed sequentially as the responses endpoint
        does not support native batching.

        Args:
            prompt: Single text prompt (str) or list of text prompts for batch processing
            messages: Single message list or list of message lists for batch processing
            response_format: Optional Pydantic model class for structured output
            metadata: Optional LiteLLM metadata for tracing / observability

        Returns:
            Single string/model or list of strings/models depending on input type.

        Raises:
            ValueError: If neither prompt nor messages is provided, or if both are provided.
            RuntimeError: If there's an error during generation.
        """
        # Validate inputs
        if prompt is None and messages is None:
            raise ValueError("Either prompts or messages must be provided")
        if prompt is not None and messages is not None:
            raise ValueError("Provide either prompts or messages, not both")

        # Determine if this is a single input or batch input
        single_input = False
        batch_prompts = None
        batch_messages = None

        if prompt is not None:
            if isinstance(prompt, str):
                # Single prompt - convert to batch
                batch_prompts = [prompt]
                single_input = True
            elif isinstance(prompt, list):
                # Already a list of prompts
                batch_prompts = prompt
                single_input = False
            else:
                raise ValueError("prompt must be a string or list of strings")

        if messages is not None:
            if isinstance(messages, list) and len(messages) > 0:
                # Check if it's a single message list or batch
                if isinstance(messages[0], dict):
                    # Single message list - convert to batch
                    batch_messages = [messages]
                    single_input = True
                elif isinstance(messages[0], list):
                    # Already a batch of message lists
                    batch_messages = messages
                    single_input = False
                else:
                    raise ValueError("Invalid messages format")
            else:
                raise ValueError("messages cannot be empty")

        try:
            # Convert batch prompts to messages if needed
            batch_to_send = []
            if batch_prompts is not None:
                for one_prompt in batch_prompts:
                    batch_to_send.append([{"role": "user", "content": one_prompt}])
            else:
                batch_to_send = batch_messages

            # Warn if batch processing is being used
            if len(batch_to_send) > 1:
                warnings.warn(
                    f"OpenAI responses endpoint does not support batch operations. "
                    f"Processing {len(batch_to_send)} requests sequentially.",
                    UserWarning,
                    stacklevel=2
                )

            # Process each request sequentially
            results = []
            for message_list in batch_to_send:
                # Enforce rate limit per request
                self._respect_rate_limit()

                # Prepare completion parameters
                completion_params = {
                    "model": self._get_model_string(),
                    "input": message_list,
                    "reasoning": {"effort": self.reasoning_effort},
                    "metadata": self._build_request_metadata(metadata),
                }
                
                # Add max_output_tokens if specified
                if self.max_completion_tokens is not None:
                    completion_params["max_output_tokens"] = self.max_completion_tokens
                
                # Add text_format if response_format is provided
                if response_format is not None:
                    completion_params["text_format"] = response_format

                # Call LiteLLM responses endpoint
                response = litellm.responses(**completion_params)

                # Record timestamp for rate limiting
                if self.rpm_limit is not None:
                    self._request_timestamps.append(time.monotonic())

                # Extract content from response
                # Response structure: response.output[1].content[0].text
                content = response.output[1].content[0].text
                
                if response_format is not None:
                    # Strip code fences before validation
                    content = self._strip_code_fences(content)
                    try:
                        results.append(response_format.model_validate_json(content))
                    except Exception as validation_error:
                        # Show the content that failed to parse for debugging
                        content_preview = content[:200] + "..." if len(content) > 200 else content
                        logger.warning(
                            f"JSON parsing failed, skipping response | "
                            f"Model: {self.model_id} | "
                            f"Format: {response_format.__name__} | "
                            f"Content preview: {content_preview}"
                        )
                        raise ValueError(
                            f"Failed to parse JSON response into {response_format.__name__}.\n"
                            f"Validation error: {validation_error}\n"
                            f"Content received (first 200 chars):\n{content_preview}"
                        ) from validation_error
                else:
                    # Strip leading/trailing whitespace for text responses
                    results.append(content.strip() if content else content)

            # Return single result for backward compatibility
            if single_input and len(results) == 1:
                return results[0]
            return results

        except Exception as e:
            error_trace = traceback.format_exc()
            logger.error(
                f"Generation failed | Provider: {self.provider_name} | "
                f"Model: {self.model_id} | Error: {str(e)}"
            )
            raise RuntimeError(
                f"Error generating response with {self.provider_name}:\n{error_trace}"
            )


class AnthropicProvider(LLMProvider):
    """Anthropic provider using litellm."""

    @property
    def provider_name(self) -> str:
        return "anthropic"

    @property
    def env_key_name(self) -> str:
        return "ANTHROPIC_API_KEY"

    def __init__(
        self,
        model_id: str = "claude-haiku-4-5-20251001",
        api_key: str | None = None,
        temperature: float | None = None,
        max_completion_tokens: int | None = None,
        timeout: int | None = None,
        # top_p: float | None = None, # Not properly supported by anthropic models 4.5
        # frequency_penalty: float | None = None,  # Not supported by anthropic models 4.5
    ):
        """Initialize the Anthropic provider.

        Args:
            model_id: The model ID (defaults to claude-haiku-4-5-20251001)
            api_key: API key (if None, will get from environment)
            temperature: Temperature for generation (0.0 to 1.0)
            max_completion_tokens: Maximum tokens to generate
            timeout: Request timeout in seconds
            top_p: Nucleus sampling parameter (0.0 to 1.0)
        """
        super().__init__(
            model_id=model_id,
            api_key=api_key,
            temperature=temperature,
            max_completion_tokens=max_completion_tokens,
            timeout=timeout,
        )


class GeminiProvider(LLMProvider):
    """Google Gemini provider using litellm."""

    @property
    def provider_name(self) -> str:
        return "gemini"

    @property
    def env_key_name(self) -> str:
        return "GEMINI_API_KEY"

    def __init__(
        self,
        model_id: str = "gemini-2.0-flash",
        api_key: str | None = None,
        temperature: float | None = None,
        max_completion_tokens: int | None = None,
        top_p: float | None = None,
        frequency_penalty: float | None = None,
        rpm_limit: int | None = None,
        timeout: int | None = None,
    ):
        """Initialize the Gemini provider.

        Args:
            model_id: The model ID (defaults to gemini-2.0-flash)
            api_key: API key (if None, will get from environment)
            temperature: Temperature for generation (0.0 to 1.0)
            max_completion_tokens: Maximum tokens to generate
            top_p: Nucleus sampling parameter (0.0 to 1.0)
            frequency_penalty: Penalty for token frequency (-2.0 to 2.0)
            timeout: Request timeout in seconds
        """
        super().__init__(
            model_id=model_id,
            api_key=api_key,
            temperature=temperature,
            max_completion_tokens=max_completion_tokens,
            top_p=top_p,
            frequency_penalty=frequency_penalty,
            rpm_limit=rpm_limit,
            timeout=timeout,
        )


class OllamaProvider(LLMProvider):
    """Ollama provider using litellm.

    Note: Ollama typically doesn't require an API key as it's usually run locally.
    """

    @property
    def provider_name(self) -> str:
        return "ollama_chat"

    @property
    def env_key_name(self) -> str:
        return "OLLAMA_API_BASE"

    def _get_api_key(self) -> str:
        """Override to handle Ollama not requiring an API key.

        Returns an empty string since Ollama typically doesn't need an API key.
        OLLAMA_API_BASE can be used to set a custom base URL.
        """
        return ""

    def __init__(
        self,
        model_id: str = "gemma3:4b",
        temperature: float | None = None,
        max_completion_tokens: int | None = None,
        top_p: float | None = None,
        frequency_penalty: float | None = None,
        api_base: str | None = None,
        rpm_limit: int | None = None,
        timeout: int | None = None,
    ):
        """Initialize the Ollama provider.

        Args:
            model_id: The model ID (defaults to llama3)
            temperature: Temperature for generation (0.0 to 1.0)
            max_completion_tokens: Maximum tokens to generate
            top_p: Nucleus sampling parameter (0.0 to 1.0)
            frequency_penalty: Penalty for token frequency (-2.0 to 2.0)
            api_base: Base URL for Ollama API (e.g., "http://localhost:11434")
            timeout: Request timeout in seconds
        """
        # Set API base URL if provided
        if api_base:
            os.environ["OLLAMA_API_BASE"] = api_base

        super().__init__(
            model_id=model_id,
            api_key="",  # Pass empty string since parent class requires this parameter
            temperature=temperature,
            max_completion_tokens=max_completion_tokens,
            top_p=top_p,
            frequency_penalty=frequency_penalty,
            rpm_limit=rpm_limit,
            timeout=timeout,
        )


class OpenRouterProvider(LLMProvider):
    """OpenRouter provider using litellm"""

    @property
    def provider_name(self) -> str:
        return "openrouter"
    
    @property
    def env_key_name(self) -> str:
        return "OPENROUTER_API_KEY"
    
    def __init__(
            self,
            model_id: str = "openai/gpt-5-mini",  # for default model
            api_key: str | None = None,
            temperature: float | None = None,
            max_completion_tokens: int | None = None,
            top_p: float | None = None,
            frequency_penalty: float | None = None,
            timeout: int | None = None,
    ):
        """Initialize the OpenRouter provider.

        Args:
            model_id: The model ID (defaults to openai/gpt-5-mini)
            api_key: API key (if None, will get from environment)
            temperature: Temperature for generation (0.0 to 1.0)
            max_completion_tokens: Maximum tokens to generate
            top_p: Nucleus sampling parameter (0.0 to 1.0)
            frequency_penalty: Penalty for token frequency (-2.0 to 2.0)
            timeout: Request timeout in seconds
        """
        super().__init__(
            model_id = model_id,
            api_key = api_key,
            temperature = temperature,
            max_completion_tokens = max_completion_tokens,
            top_p = top_p,
            frequency_penalty = frequency_penalty,
            timeout = timeout,
        )


class MistralProvider(LLMProvider):
    """Mistral AI provider using litellm."""

    @property
    def provider_name(self) -> str:
        return "mistral"

    @property
    def env_key_name(self) -> str:
        return "MISTRAL_API_KEY"

    def __init__(
        self,
        model_id: str = "mistral-small-latest",
        api_key: str | None = None,
        temperature: float | None = None,
        max_completion_tokens: int | None = None,
        top_p: float | None = None,
        frequency_penalty: float | None = None,
        rpm_limit: int | None = None,
        timeout: int | None = None,
    ):
        """Initialize the Mistral provider.

        Args:
            model_id: The model ID (defaults to mistral-small-latest)
            api_key: API key (if None, will get from MISTRAL_API_KEY env var)
            temperature: Temperature for generation (0.0 to 1.0)
            max_completion_tokens: Maximum tokens to generate
            top_p: Nucleus sampling parameter (0.0 to 1.0)
            frequency_penalty: Penalty for token frequency (-2.0 to 2.0)
            rpm_limit: Requests per minute limit for rate limiting
            timeout: Request timeout in seconds
        """
        super().__init__(
            model_id=model_id,
            api_key=api_key,
            temperature=temperature,
            max_completion_tokens=max_completion_tokens,
            top_p=top_p,
            frequency_penalty=frequency_penalty,
            rpm_limit=rpm_limit,
            timeout=timeout,
        )


def openai(model_id: str = "gpt-5-mini-2025-08-07", **kwargs) -> OpenAIProvider:
    """Create an OpenAI provider instance."""
    return OpenAIProvider(model_id=model_id, **kwargs)


def anthropic(
    model_id: str = "claude-haiku-4-5-20251001",
    **kwargs,
) -> AnthropicProvider:
    """Create an Anthropic provider instance."""
    return AnthropicProvider(model_id=model_id, **kwargs)


def gemini(model_id: str = "gemini-2.0-flash", **kwargs) -> GeminiProvider:
    """Create a Gemini provider instance."""
    return GeminiProvider(model_id=model_id, **kwargs)


def ollama(model_id: str = "gemma3:4b", **kwargs) -> OllamaProvider:
    """Create an Ollama provider instance."""
    return OllamaProvider(model_id=model_id, **kwargs)


def openrouter(
    model_id: str = "openai/gpt-5-mini",
    **kwargs,
) -> OpenRouterProvider:
    """Create an OpenRouter provider instance."""
    return OpenRouterProvider(model_id=model_id, **kwargs)


def mistral(model_id: str = "mistral-small-latest", **kwargs) -> MistralProvider:
    """Create a Mistral provider instance."""
    return MistralProvider(model_id=model_id, **kwargs)


__all__ = [
    "LLMProvider",
    "OpenAIProvider",
    "AnthropicProvider",
    "GeminiProvider",
    "OllamaProvider",
    "OpenRouterProvider",
    "MistralProvider",
    "openai",
    "anthropic",
    "gemini",
    "ollama",
    "openrouter",
    "mistral",
]