app/llm/llm_throttle.py

from __future__ import annotations

import asyncio
import random
import threading
import time
from typing import Awaitable, Callable, TypeVar

import structlog
from tenacity import (
    AsyncRetrying,
    RetryCallState,
    retry_if_exception_type,
)

from app.config import settings

try:
    # openai v1: RateLimitError lives in the top-level module
    from openai import RateLimitError
except Exception:  # pragma: no cover
    RateLimitError = Exception  # type: ignore[assignment,misc]

T = TypeVar("T")

_semaphore: asyncio.Semaphore | None = None
_sync_semaphore: threading.Semaphore | None = None

# Configure structlog once on import. This is intentionally minimal: JSON to
# stdout with enough fields for the latency optimisation work.
structlog.configure(
    processors=[
        structlog.processors.TimeStamper(fmt="iso_8601", utc=True),
        structlog.processors.add_log_level,
        structlog.processors.StackInfoRenderer(),
        structlog.processors.format_exc_info,
        structlog.processors.JSONRenderer(),
    ],
    logger_factory=structlog.stdlib.LoggerFactory(),
    cache_logger_on_first_use=True,
)

_log = structlog.get_logger(__name__)


def _get_semaphore() -> asyncio.Semaphore:
    global _semaphore
    if _semaphore is None:
        _semaphore = asyncio.Semaphore(int(settings.max_concurrent_llm_calls))
    return _semaphore


def _get_sync_semaphore() -> threading.Semaphore:
    global _sync_semaphore
    if _sync_semaphore is None:
        _sync_semaphore = threading.Semaphore(int(settings.max_concurrent_llm_calls))
    return _sync_semaphore


def _exp_backoff_with_jitter(retry_state: RetryCallState) -> float:
    """Exponential backoff with +/-20% jitter."""
    # attempt_number starts at 1 for the first retry.
    attempt = retry_state.attempt_number
    base = 1.0 * (2 ** (attempt - 1))
    base = min(base, 60.0)
    jitter = base * random.uniform(-0.2, 0.2)
    return max(0.0, base + jitter)


def make_cache_hit_slot() -> list[bool | None]:
    """Mutable slot callers fill after ``log_openai_cache_usage`` inside ``call``."""
    return [None]


def _resolved_cache_hit(
    cache_hit: bool | None,
    cache_hit_out: list[bool | None] | None,
) -> bool | None:
    if cache_hit_out is not None and len(cache_hit_out) > 0:
        return cache_hit_out[0]
    return cache_hit


async def throttled_llm_call(
    *,
    phase: str,
    section_id: str | None,
    call: Callable[[], Awaitable[T]],
    cache_hit: bool | None = None,
    cache_hit_out: list[bool | None] | None = None,
) -> T:
    """Wrap an LLM call with:
    - global concurrency semaphore
    - exponential backoff with jitter on OpenAI RateLimitError
    - structured JSON logs via structlog
    """
    sem = _get_semaphore()
    start = time.perf_counter()
    last_exc: BaseException | None = None

    async with sem:
        retrying = AsyncRetrying(
            retry=retry_if_exception_type(RateLimitError),
            wait=_exp_backoff_with_jitter,
            reraise=False,
        )

        async for attempt in retrying:
            try:
                result = await call()
                duration_ms = int((time.perf_counter() - start) * 1000)
                hit = _resolved_cache_hit(cache_hit, cache_hit_out)
                _log.info(
                    event="llm_call",
                    duration_ms=duration_ms,
                    phase=phase,
                    section_id=section_id,
                    cache_hit=hit,
                    attempt=attempt.retry_state.attempt_number,
                )
                return result
            except BaseException as exc:  # noqa: BLE001
                last_exc = exc
                duration_ms = int((time.perf_counter() - start) * 1000)
                hit = _resolved_cache_hit(cache_hit, cache_hit_out)
                _log.warning(
                    event="llm_call_rate_limited",
                    duration_ms=duration_ms,
                    phase=phase,
                    section_id=section_id,
                    cache_hit=hit,
                    attempt=attempt.retry_state.attempt_number,
                    exc_type=type(exc).__name__,
                )

        # If we got here, reraise=False and we exhausted retries.
        duration_ms = int((time.perf_counter() - start) * 1000)
        hit = _resolved_cache_hit(cache_hit, cache_hit_out)
        _log.error(
            event="llm_call_exhausted_retries",
            duration_ms=duration_ms,
            phase=phase,
            section_id=section_id,
            cache_hit=hit,
        )
        if last_exc is None:
            raise RuntimeError("LLM call exhausted retries without exception")
        raise last_exc


def throttled_sync_llm_call(
    *,
    phase: str,
    section_id: str | None,
    call: Callable[[], T],
) -> T:
    """Sync OpenAI calls with global concurrency cap and rate-limit retries."""
    from tenacity import Retrying, retry_if_exception_type

    sem = _get_sync_semaphore()
    start = time.perf_counter()
    with sem:
        retrying = Retrying(
            retry=retry_if_exception_type(RateLimitError),
            wait=_exp_backoff_with_jitter,
            reraise=True,
        )
        for attempt in retrying:
            with attempt:
                try:
                    result = call()
                except RateLimitError:
                    duration_ms = int((time.perf_counter() - start) * 1000)
                    _log.warning(
                        event="llm_call_rate_limited",
                        duration_ms=duration_ms,
                        phase=phase,
                        section_id=section_id,
                        cache_hit=None,
                        attempt=attempt.retry_state.attempt_number,
                        exc_type="RateLimitError",
                    )
                    raise
                duration_ms = int((time.perf_counter() - start) * 1000)
                _log.info(
                    event="llm_call",
                    duration_ms=duration_ms,
                    phase=phase,
                    section_id=section_id,
                    cache_hit=None,
                    attempt=attempt.retry_state.attempt_number,
                )
                return result
    raise RuntimeError("unreachable")