113 lines
3.2 KiB
Python
113 lines
3.2 KiB
Python
"""LiteLLM wrapper for model-agnostic LLM access.
|
|
|
|
Supports OpenAI, Anthropic, Google, Mistral, and Ollama via a unified interface.
|
|
Implements retry logic with exponential backoff for transient failures.
|
|
"""
|
|
|
|
import logging
|
|
import time
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from ingestion.config import IngestionConfig
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Retry configuration
|
|
MAX_RETRIES = 3
|
|
BASE_DELAY_SECONDS = 1.0
|
|
MAX_DELAY_SECONDS = 30.0
|
|
BACKOFF_FACTOR = 2.0
|
|
|
|
|
|
class LLMError(Exception):
|
|
"""Raised when an LLM call fails after all retries."""
|
|
|
|
def __init__(self, provider: str, model: str, message: str):
|
|
self.provider = provider
|
|
self.model = model
|
|
super().__init__(
|
|
f"LLM call failed [{provider}/{model}]: {message}"
|
|
)
|
|
|
|
|
|
class LLMClient:
|
|
"""Model-agnostic LLM client using LiteLLM."""
|
|
|
|
def __init__(self, config: IngestionConfig):
|
|
self.model = config.llm_model_string
|
|
self.provider = config.llm_provider
|
|
self.config = config
|
|
|
|
def complete(
|
|
self,
|
|
messages: list[dict[str, str]],
|
|
response_format: type[BaseModel] | None = None,
|
|
) -> str:
|
|
"""Send a completion request with retry logic.
|
|
|
|
Args:
|
|
messages: List of message dicts with 'role' and 'content' keys.
|
|
response_format: Optional Pydantic model for structured output.
|
|
|
|
Returns:
|
|
The assistant's response content as a string.
|
|
|
|
Raises:
|
|
LLMError: If all retries are exhausted.
|
|
"""
|
|
import litellm
|
|
|
|
last_error: Exception | None = None
|
|
|
|
for attempt in range(1, MAX_RETRIES + 1):
|
|
try:
|
|
kwargs: dict[str, Any] = {
|
|
"model": self.model,
|
|
"messages": messages,
|
|
}
|
|
|
|
if response_format is not None:
|
|
kwargs["response_format"] = response_format
|
|
|
|
response = litellm.completion(**kwargs)
|
|
content = response.choices[0].message.content
|
|
return content or ""
|
|
|
|
except (
|
|
litellm.RateLimitError,
|
|
litellm.Timeout,
|
|
litellm.ServiceUnavailableError,
|
|
litellm.InternalServerError,
|
|
) as e:
|
|
last_error = e
|
|
delay = min(
|
|
BASE_DELAY_SECONDS * (BACKOFF_FACTOR ** (attempt - 1)),
|
|
MAX_DELAY_SECONDS,
|
|
)
|
|
logger.warning(
|
|
"LLM call attempt %d/%d failed (%s: %s), retrying in %.1fs",
|
|
attempt,
|
|
MAX_RETRIES,
|
|
type(e).__name__,
|
|
str(e),
|
|
delay,
|
|
)
|
|
time.sleep(delay)
|
|
|
|
except Exception as e:
|
|
# Non-retryable error
|
|
raise LLMError(
|
|
provider=self.provider,
|
|
model=self.model,
|
|
message=f"{type(e).__name__}: {e}",
|
|
) from e
|
|
|
|
# All retries exhausted
|
|
raise LLMError(
|
|
provider=self.provider,
|
|
model=self.model,
|
|
message=f"All {MAX_RETRIES} retries exhausted. Last error: {last_error}",
|
|
)
|