shared/llm_client.py

"""
LLM client for text and JSON completions.

Provides a unified interface for LLM calls using litellm, which supports
multiple providers (OpenAI, Anthropic, etc.).

Configuration via environment variables:
- LLM_MODEL: Model to use (default: gpt-4o-mini)
- LLM_PROVIDER: Provider prefix for litellm (default: openai)
- OPENAI_API_KEY: Required for OpenAI models
- ANTHROPIC_API_KEY: Required for Anthropic models
"""
import os
import json
import asyncio
from typing import Any, Dict, List, Optional
from litellm import completion, acompletion

# Default model and provider from environment
# Supports both formats:
#   - DEFAULT_LLM_MODEL=provider/model (e.g., gemini/gemini-3-flash-preview)
#   - LLM_PROVIDER + LLM_MODEL separately
def _parse_default_model():
    default_llm = os.getenv("DEFAULT_LLM_MODEL", "")
    if "/" in default_llm:
        parts = default_llm.split("/", 1)
        return parts[1], parts[0]
    return os.getenv("LLM_MODEL", "gpt-4o-mini"), os.getenv("LLM_PROVIDER", "openai")

DEFAULT_MODEL, DEFAULT_PROVIDER = _parse_default_model()


class LLMError(Exception):
    """Generic error from LLM provider."""
    pass


class LLMParseError(Exception):
    """Error parsing JSON response from LLM."""
    pass


def _validate_api_key(provider: str) -> None:
    """
    Validate that the required API key is set for the given provider.

    Args:
        provider: The LLM provider (openai, anthropic, etc.)

    Raises:
        EnvironmentError: If the required API key is not found
    """
    key_mapping = {
        "openai": "OPENAI_API_KEY",
        "anthropic": "ANTHROPIC_API_KEY",
        "azure": "AZURE_API_KEY",
        "xai": "XAI_API_KEY",
    }

    env_var = key_mapping.get(provider)
    if env_var and not os.getenv(env_var):
        raise EnvironmentError(f"{env_var} not found in environment variables.")


def _get_full_model_name(model: str, provider: str) -> str:
    """
    Get the full model name with provider prefix for litellm.

    Args:
        model: Base model name (e.g., 'gpt-4o-mini', 'claude-3-haiku')
        provider: Provider name (e.g., 'openai', 'anthropic')

    Returns:
        Full model name for litellm (e.g., 'openai/gpt-4o-mini')
    """
    # If model already has provider prefix, return as-is
    if "/" in model:
        return model

    # Add provider prefix for non-OpenAI providers
    if provider != "openai":
        return f"{provider}/{model}"

    return model


def llm_complete(
    prompt: str,
    system_prompt: str = "You are a helpful financial assistant.",
    model: str | None = None,
    provider: str | None = None,
    temperature: float = 1.0
) -> str:
    """
    Get a simple text completion from the LLM.

    Args:
        prompt: The user prompt to send
        system_prompt: System instructions for the LLM
        model: Model name (defaults to LLM_MODEL env var or gpt-4o-mini)
        provider: Provider name (defaults to LLM_PROVIDER env var or openai)
        temperature: Sampling temperature (0.0 = deterministic)

    Returns:
        The LLM's text response

    Raises:
        LLMError: If the completion fails
        EnvironmentError: If required API key is missing
    """
    model = model or DEFAULT_MODEL
    provider = provider or DEFAULT_PROVIDER

    _validate_api_key(provider)

    try:
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ]

        full_model = _get_full_model_name(model, provider)

        response = completion(
            model=full_model,
            messages=messages,
            temperature=temperature
        )

        return response.choices[0].message.content.strip()

    except Exception as e:
        raise LLMError(f"LLM completion failed: {str(e)}")


async def llm_complete_async(
    prompt: str = "",
    system_prompt: str = "You are a helpful financial assistant.",
    model: str | None = None,
    provider: str | None = None,
    temperature: float = 0.2,
    messages: List[Dict[str, str]] | None = None,
    max_tokens: int | None = None,
    max_retries: int = 3,
) -> str:
    """
    Async version of llm_complete with retry logic. Supports multi-turn conversations.

    Args:
        prompt: The user prompt (ignored if messages is provided)
        system_prompt: System instructions (ignored if messages is provided)
        model: Model name (defaults to LLM_MODEL env var)
        provider: Provider name (defaults to LLM_PROVIDER env var)
        temperature: Sampling temperature
        messages: Optional pre-built messages list for multi-turn conversations
        max_tokens: Optional max tokens in response
        max_retries: Number of retry attempts on failure

    Returns:
        The LLM's text response

    Raises:
        LLMError: If all retries fail
    """
    model = model or DEFAULT_MODEL
    provider = provider or DEFAULT_PROVIDER

    _validate_api_key(provider)

    if messages is None:
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ]

    full_model = _get_full_model_name(model, provider)

    completion_kwargs: Dict[str, Any] = {
        "model": full_model,
        "messages": messages,
        "temperature": temperature,
    }
    if max_tokens is not None:
        completion_kwargs["max_tokens"] = max_tokens

    last_error = None
    for attempt in range(max_retries):
        try:
            response = await acompletion(**completion_kwargs)
            return response.choices[0].message.content.strip()

        except Exception as e:
            last_error = LLMError(f"LLM completion failed: {str(e)}")
            await asyncio.sleep(2 ** attempt)  # Exponential backoff

    raise last_error or LLMError("LLM request failed after retries")


def llm_complete_json(
    prompt: str,
    system_prompt: str = "You are a data extraction assistant. Output valid JSON only.",
    model: str | None = None,
    provider: str | None = None,
    temperature: float = 1.0,
    response_model: type | None = None,
) -> Dict[str, Any]:
    """
    Get a JSON response from the LLM.

    Ensures the output is parsed into a Python dictionary. Optionally validates
    against a Pydantic model.

    Args:
        prompt: The user prompt to send
        system_prompt: System instructions for the LLM
        model: Model name (defaults to LLM_MODEL env var or gpt-4o-mini)
        provider: Provider name (defaults to LLM_PROVIDER env var or openai)
        temperature: Sampling temperature (0.0 = deterministic)
        response_model: Optional Pydantic model class to validate/parse response

    Returns:
        Parsed JSON as dictionary, or Pydantic model instance if response_model provided

    Raises:
        LLMError: If the completion fails
        LLMParseError: If JSON parsing fails
        EnvironmentError: If required API key is missing
    """
    model = model or DEFAULT_MODEL
    provider = provider or DEFAULT_PROVIDER

    _validate_api_key(provider)

    try:
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ]

        full_model = _get_full_model_name(model, provider)

        response = completion(
            model=full_model,
            messages=messages,
            temperature=temperature,
            response_format={"type": "json_object"}
        )

        content = response.choices[0].message.content.strip()

        try:
            data = json.loads(content)
        except json.JSONDecodeError:
            raise LLMParseError(f"Failed to parse JSON response: {content}")

        # If a response model is provided, validate and return instance
        if response_model is not None:
            try:
                return response_model(**data)
            except Exception as e:
                raise LLMParseError(f"Failed to validate response against {response_model.__name__}: {e}")

        return data

    except LLMParseError:
        raise
    except Exception as e:
        raise LLMError(f"LLM JSON completion failed: {str(e)}")