""" 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)}")