backend/llm/openai_client.py

"""Thin OpenAI client wrapper used across the v2 backend.

Exposes chat helpers over ``client.chat.completions.create`` with:

* Configurable concurrency gate (``MAX_CONCURRENT_LLM_CALLS``)
* Exponential backoff retries on HTTP 429 rate limits
* Async wrappers (``*_async``) that acquire an ``asyncio.Semaphore`` before
  delegating to the sync implementation in a worker thread

When no API key is configured, ``is_available`` is False and callers fall back to
deterministic behaviour so the system runs (and tests pass) offline.
"""

from __future__ import annotations

import asyncio
import json
import logging
import random
import re
import threading
import time
from collections.abc import Callable
from contextlib import asynccontextmanager, contextmanager
from typing import Any, TypeVar

from backend.config import settings

logger = logging.getLogger(__name__)

T = TypeVar("T")

_client = None
_thread_llm_semaphore: threading.Semaphore | None = None
_async_llm_semaphore: asyncio.Semaphore | None = None
_semaphore_limit: int | None = None
_semaphore_lock = threading.Lock()


def is_available() -> bool:
    """Whether an OpenAI API key is configured."""
    return bool((settings.openai_api_key or "").strip())


def _default_timeout() -> float:
    return float(settings.openai_request_timeout_seconds)


def pipeline_timeout() -> float:
    """Strict timeout for generation-pipeline LLM / embedding calls."""
    return float(settings.openai_pipeline_timeout_seconds)


def max_concurrent_llm_calls() -> int:
    return max(1, int(settings.max_concurrent_llm_calls))


def _resolve_timeout(timeout: float | None) -> float:
    return float(timeout if timeout is not None else pipeline_timeout())


def _reset_rate_limit_state() -> None:
    global _thread_llm_semaphore, _async_llm_semaphore, _semaphore_limit
    with _semaphore_lock:
        _thread_llm_semaphore = None
        _async_llm_semaphore = None
        _semaphore_limit = None


def _ensure_semaphores() -> None:
    global _thread_llm_semaphore, _async_llm_semaphore, _semaphore_limit
    limit = max_concurrent_llm_calls()
    with _semaphore_lock:
        if _thread_llm_semaphore is None or _semaphore_limit != limit:
            _thread_llm_semaphore = threading.Semaphore(limit)
            _async_llm_semaphore = None
            _semaphore_limit = limit


def _get_thread_semaphore() -> threading.Semaphore:
    _ensure_semaphores()
    assert _thread_llm_semaphore is not None
    return _thread_llm_semaphore


def _get_async_semaphore() -> asyncio.Semaphore:
    global _async_llm_semaphore
    _ensure_semaphores()
    if _async_llm_semaphore is None:
        _async_llm_semaphore = asyncio.Semaphore(max_concurrent_llm_calls())
    return _async_llm_semaphore


@contextmanager
def _sync_llm_slot():
    sem = _get_thread_semaphore()
    sem.acquire()
    try:
        yield
    finally:
        sem.release()


@asynccontextmanager
async def _async_llm_slot():
    async with _get_async_semaphore():
        yield


def _get_client():
    global _client
    if _client is None:
        from openai import OpenAI

        _client = OpenAI(
            api_key=settings.openai_api_key,
            timeout=_default_timeout(),
        )
    return _client


def reset_client() -> None:
    """Reset the cached client and concurrency gates (tests / config reloads)."""
    global _client
    _client = None
    _reset_rate_limit_state()


def _is_rate_limit_error(exc: BaseException) -> bool:
    status = getattr(exc, "status_code", None)
    if status == 429:
        return True
    try:
        from openai import APIStatusError, RateLimitError
    except ImportError:
        return False
    return isinstance(exc, (RateLimitError,)) or (
        isinstance(exc, APIStatusError) and exc.status_code == 429
    )


def _retry_after_seconds(exc: BaseException) -> float | None:
    response = getattr(exc, "response", None)
    headers = getattr(response, "headers", None) if response is not None else None
    if not headers:
        return None
    raw = headers.get("retry-after") or headers.get("Retry-After")
    if raw is None:
        return None
    try:
        return float(raw)
    except (TypeError, ValueError):
        return None


def _run_with_rate_limit_retry(
    operation: Callable[[], T],
    *,
    gate_held: bool = False,
) -> T:
    """Execute *operation* under the concurrency gate with 429 backoff retries."""
    max_retries = max(0, int(settings.openai_rate_limit_max_retries))
    base_delay = float(settings.openai_rate_limit_backoff_base_seconds)
    max_delay = float(settings.openai_rate_limit_backoff_max_seconds)

    attempt = 0
    while True:
        try:
            if gate_held:
                return operation()
            with _sync_llm_slot():
                return operation()
        except Exception as exc:  # noqa: BLE001 — classify provider rate limits
            if not _is_rate_limit_error(exc) or attempt >= max_retries:
                raise
            attempt += 1
            retry_after = _retry_after_seconds(exc)
            if retry_after is not None:
                delay = min(retry_after, max_delay)
            else:
                delay = min(base_delay * (2 ** (attempt - 1)), max_delay)
                delay *= 0.5 + random.random() * 0.5
            logger.warning(
                "OpenAI rate limit (429); retry %s/%s in %.2fs",
                attempt,
                max_retries,
                delay,
            )
            time.sleep(delay)


async def _run_async_llm(
    sync_fn: Callable[..., T],
    /,
    *args: Any,
    **kwargs: Any,
) -> T:
    """Acquire the async concurrency gate, then run *sync_fn* in a worker thread."""
    async with _async_llm_slot():
        kwargs["_gate_held"] = True
        return await asyncio.to_thread(sync_fn, *args, **kwargs)


def _parse_llm_json(content: str) -> dict:
    """Parse an LLM JSON payload, stripping optional markdown code fences."""
    text = (content or "").strip()
    if text.startswith("```"):
        text = re.sub(r"^```(?:json)?\s*", "", text, flags=re.IGNORECASE)
        text = re.sub(r"\s*```$", "", text)
        text = text.strip()
    try:
        parsed = json.loads(text)
    except json.JSONDecodeError:
        logger.warning("LLM response is not valid JSON after fence stripping.")
        return {}
    return parsed if isinstance(parsed, dict) else {}


def chat_text(
    messages: list[dict],
    *,
    model: str | None = None,
    temperature: float = 0.0,
    max_tokens: int | None = None,
    timeout: float | None = None,
    _gate_held: bool = False,
) -> str:
    """Return the assistant text for a chat completion."""
    def _call() -> str:
        resp = _get_client().chat.completions.create(
            model=model or settings.mapping_model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            timeout=_resolve_timeout(timeout),
        )
        return (resp.choices[0].message.content or "").strip()

    return _run_with_rate_limit_retry(_call, gate_held=_gate_held)


async def chat_text_async(
    messages: list[dict],
    *,
    model: str | None = None,
    temperature: float = 0.0,
    max_tokens: int | None = None,
    timeout: float | None = None,
) -> str:
    """Async chat completion with concurrency gate + 429 retry."""
    return await _run_async_llm(
        chat_text,
        messages,
        model=model,
        temperature=temperature,
        max_tokens=max_tokens,
        timeout=timeout,
    )


def chat_vision_json(
    messages: list[dict],
    *,
    model: str | None = None,
    max_tokens: int | None = None,
    timeout: float | None = None,
    _gate_held: bool = False,
) -> dict:
    """Vision-capable chat completion returning parsed JSON."""
    def _call() -> dict:
        kwargs: dict[str, Any] = {
            "model": model or settings.vision_model,
            "messages": messages,
            "temperature": 0.0,
            "max_tokens": max_tokens or settings.vision_max_tokens,
            "response_format": {"type": "json_object"},
            "timeout": _resolve_timeout(timeout),
        }
        resp = _get_client().chat.completions.create(**kwargs)
        content = (resp.choices[0].message.content or "").strip()
        return _parse_llm_json(content)

    return _run_with_rate_limit_retry(_call, gate_held=_gate_held)


async def chat_vision_json_async(
    messages: list[dict],
    *,
    model: str | None = None,
    max_tokens: int | None = None,
    timeout: float | None = None,
) -> dict:
    return await _run_async_llm(
        chat_vision_json,
        messages,
        model=model,
        max_tokens=max_tokens,
        timeout=timeout,
    )


def chat_json(
    messages: list[dict[str, str]],
    *,
    model: str | None = None,
    temperature: float = 0.0,
    max_tokens: int | None = None,
    timeout: float | None = None,
    _gate_held: bool = False,
) -> dict:
    """Return the parsed JSON object for a chat completion in JSON mode."""
    def _call() -> dict:
        resp = _get_client().chat.completions.create(
            model=model or settings.mapping_model,
            messages=messages,
            temperature=temperature,
            max_tokens=max_tokens,
            response_format={"type": "json_object"},
            timeout=_resolve_timeout(timeout),
        )
        content = (resp.choices[0].message.content or "").strip()
        return _parse_llm_json(content)

    return _run_with_rate_limit_retry(_call, gate_held=_gate_held)


async def chat_json_async(
    messages: list[dict[str, str]],
    *,
    model: str | None = None,
    temperature: float = 0.0,
    max_tokens: int | None = None,
    timeout: float | None = None,
) -> dict:
    """Async JSON chat completion with concurrency gate + 429 retry."""
    return await _run_async_llm(
        chat_json,
        messages,
        model=model,
        temperature=temperature,
        max_tokens=max_tokens,
        timeout=timeout,
    )