src/generation/llm_client.py

"""Unified async/sync LLM client via LiteLLM.

Supports Groq, OpenAI, and OpenRouter via provider-prefixed model strings:
  groq/llama-3.1-8b-instant
  openai/gpt-4o-mini
  openrouter/meta-llama/llama-3.3-70b-instruct

`agenerate` is the primary async entry point used by the benchmark pipeline.
`generate` is a sync wrapper for single-query and query-transform paths.
Model strings without a "/" prefix are assumed to be Groq models.

Rate limiting
─────────────
`ProviderRateLimiter` is a sliding-window async rate limiter keyed by provider
prefix ("groq", "openai", etc.).  `agenerate` acquires a slot before every
call so the benchmark pipeline never exceeds the configured RPM ceiling even
when many coroutines are in flight simultaneously.

LiteLLM `num_retries=3` handles transient 429s that slip through (e.g. burst
spikes or shared-key contention) with exponential back-off and respects the
provider's Retry-After header.
"""

import asyncio
import logging
import os
import re
import time

import litellm
from litellm import acompletion, completion
from litellm.exceptions import RateLimitError as LiteLLMRateLimitError

from src import settings

log = logging.getLogger(__name__)

# LiteLLM global retry disabled — we manage retries ourselves so we can
# honour the exact wait time Groq sends in the 429 error body.
litellm.num_retries = 0

# ---------------------------------------------------------------------------
# Rate limits — provider-level RPM and model-level overrides
#
# Model-level limits are needed when TPM is the binding constraint rather
# than RPM.  llama-3.3-70b-versatile has only 12k TPM on the Groq free tier;
# each judge call uses ~3-4k tokens, so the effective safe rate is ~3 RPM
# regardless of the 30 RPM ceiling.
# ---------------------------------------------------------------------------
_MODEL_RPM: dict[str, int] = {
    "groq/llama-3.3-70b-versatile": 3,   # 12k TPM / ~3.5k tokens per call
    "groq/llama-3.1-70b-versatile": 3,
    "openai/gpt-4o":                5,   # 30k TPM / ~5k tokens per judge call
    "openai/gpt-4o-mini":          40,   # 200k TPM / ~5k tokens per call
}
_PROVIDER_RPM: dict[str, int] = {
    "groq":       28,   # Groq free tier: 30 RPM; 2 headroom
    "openai":    490,   # OpenAI tier-1: 500 RPM
    "openrouter": 18,   # OpenRouter free: ~20 RPM
}
_DEFAULT_RPM = 20

_MAX_RETRIES = 3  # maximum retries on RateLimitError before giving up


class ProviderRateLimiter:
    """Sliding-window async rate limiter.

    Tracks the timestamps of the last `rpm` calls in a 60-second window.
    If the window is full, waits until the oldest call drops out before
    allowing a new one.  One instance per provider, shared across all
    concurrent benchmark coroutines.
    """

    def __init__(self, rpm: int):
        self._rpm = rpm
        self._timestamps: list[float] = []
        self._lock = asyncio.Lock()

    async def acquire(self) -> None:
        async with self._lock:
            now = time.monotonic()
            # Drop timestamps older than 60 s
            self._timestamps = [t for t in self._timestamps if now - t < 60.0]
            if len(self._timestamps) >= self._rpm:
                # Wait until the oldest slot expires
                wait = 60.0 - (now - self._timestamps[0])
                if wait > 0:
                    await asyncio.sleep(wait)
                # Re-prune after sleeping
                now = time.monotonic()
                self._timestamps = [t for t in self._timestamps if now - t < 60.0]
            self._timestamps.append(time.monotonic())


_limiters: dict[str, ProviderRateLimiter] = {}

DEFAULT_MODEL = "groq/llama-3.1-8b-instant"

# Matches RAGBench section 7.3 / 7.7 "LONG" generation prompt.
NEGATIVE_RESPONSE = (
    "The documents are missing some of the information required to answer the question."
)

PROMPT_TEMPLATE = """\
You are a chatbot providing answers to user queries. You will be given one or more context documents. \
Use the information in the documents to answer the question.

If the documents do not provide enough information for you to answer the question, then say \
"{negative_response}" Don't quote things not in the documents. Don't try to make up an answer.

Context Documents:
{context}

Question: {question}"""


def _resolve_model(model: str) -> str:
    """Add groq/ prefix when no provider prefix is present."""
    return model if "/" in model else f"groq/{model}"


def _get_limiter(model: str) -> ProviderRateLimiter:
    """Return the shared ProviderRateLimiter for this model.

    Checks model-level overrides first (for TPM-constrained models like
    llama-3.3-70b-versatile), then falls back to provider-level RPM.
    """
    resolved = _resolve_model(model)
    if resolved not in _limiters:
        rpm = _MODEL_RPM.get(resolved) or _PROVIDER_RPM.get(resolved.split("/")[0], _DEFAULT_RPM)
        _limiters[resolved] = ProviderRateLimiter(rpm)
    return _limiters[resolved]


def _is_request_too_large(exc: Exception) -> bool:
    """Return True when the request itself exceeds the model's token limit.

    Unlike a transient TPM rate limit, a 'request too large' error cannot be
    resolved by waiting — the prompt must be truncated before retrying.
    """
    msg = str(exc).lower()
    return "too large" in msg or "must be reduced" in msg


def _parse_retry_after(exc: Exception) -> float:
    """Extract the wait time from a Groq/OpenAI 429 error message.

    Groq/OpenAI send: 'Please try again in 12.235s.'
    Falls back to 20 s if the pattern is not found.
    """
    m = re.search(r"try again in ([\d.]+)s", str(exc), re.IGNORECASE)
    return float(m.group(1)) + 2.0 if m else 20.0


def _sync_api_keys() -> None:
    """Push dotenv-loaded keys into os.environ so LiteLLM picks them up."""
    if settings.GROQ_API_KEY:
        os.environ.setdefault("GROQ_API_KEY", settings.GROQ_API_KEY)
    if settings.OPENAI_API_KEY:
        os.environ.setdefault("OPENAI_API_KEY", settings.OPENAI_API_KEY)
    if settings.OPENROUTER_API_KEY:
        os.environ.setdefault("OPENROUTER_API_KEY", settings.OPENROUTER_API_KEY)


async def agenerate(
    prompt: str,
    model: str = DEFAULT_MODEL,
    max_tokens: int = 256,
    temperature: float = 0.0,
) -> str:
    """Async LLM call — use this from the benchmark pipeline.

    Acquires a rate-limiter slot before each attempt, then retries on 429
    sleeping for the exact duration Groq specifies in the error body.
    """
    _sync_api_keys()
    resolved = _resolve_model(model)
    for attempt in range(_MAX_RETRIES + 1):
        await _get_limiter(model).acquire()
        try:
            resp = await acompletion(
                model=resolved,
                messages=[{"role": "user", "content": prompt}],
                temperature=temperature,
                max_tokens=max_tokens,
                num_retries=0,
            )
            return resp.choices[0].message.content.strip()
        except LiteLLMRateLimitError as exc:
            if _is_request_too_large(exc):
                log.error("agenerate prompt too large for model %s — truncation needed", resolved)
                raise
            if attempt >= _MAX_RETRIES:
                raise
            wait = _parse_retry_after(exc)
            log.warning("agenerate rate-limited (attempt %d/%d), sleeping %.1fs",
                        attempt + 1, _MAX_RETRIES, wait)
            await asyncio.sleep(wait)


def generate(
    prompt: str,
    model: str = DEFAULT_MODEL,
    max_tokens: int = 256,
    temperature: float = 0.0,
) -> str:
    """Sync LLM call — use this from query-transform and single-query paths."""
    _sync_api_keys()
    resp = completion(
        model=_resolve_model(model),
        messages=[{"role": "user", "content": prompt}],
        temperature=temperature,
        max_tokens=max_tokens,
        num_retries=3,
    )
    return resp.choices[0].message.content.strip()