src/id/llm/client.py

"""Provider-agnostic LLM client wrapping the OpenAI SDK (Section 4).

All three providers (OpenAI, OpenRouter, custom) are OpenAI-Chat-Completions
compatible, so a single implementation handles them via a custom ``base_url``.
Every call is routed by *tier name*, retried with backoff on transient errors,
and logged to the token ledger.
"""

from __future__ import annotations

import json
import re
import time
from dataclasses import dataclass
from typing import Any

from openai import APIConnectionError, APIStatusError, APITimeoutError, OpenAI, RateLimitError

from ..config import Config
from ..models import UsageRecord
from .usage import UsageLedger

_FENCE_RE = re.compile(r"^\s*```(?:json)?\s*|\s*```\s*$", re.IGNORECASE)
_TRANSIENT = (APIConnectionError, APITimeoutError, RateLimitError)


class LLMError(RuntimeError):
    """Raised on non-recoverable LLM failures (auth/config/exhausted retries)."""


class JSONParseError(LLMError):
    """Raised when a response that must be JSON cannot be parsed."""


@dataclass
class LLMResponse:
    text: str
    prompt_tokens: int
    completion_tokens: int
    total_tokens: int
    retries: int


class LLMClient:
    def __init__(
        self,
        config: Config,
        usage: UsageLedger | None = None,
        *,
        world_id: str = "",
        session_id: str = "",
    ) -> None:
        self.config = config
        self.usage = usage
        self.world_id = world_id
        self.session_id = session_id
        self._clients: dict[str, OpenAI] = {}

    def bind(self, *, world_id: str = "", session_id: str = "",
             usage: UsageLedger | None = None) -> LLMClient:
        """Return a shallow copy with updated logging context."""
        c = LLMClient(
            self.config,
            usage or self.usage,
            world_id=world_id or self.world_id,
            session_id=session_id or self.session_id,
        )
        c._clients = self._clients  # reuse pooled SDK clients
        return c

    def _client_for(self, provider: str) -> OpenAI:
        if provider not in self._clients:
            if provider == "local":
                # Self-contained, no-API mode: run a small model in-process via
                # llama.cpp behind an OpenAI-shaped adapter.
                from .local import LocalLlamaClient

                self._clients[provider] = LocalLlamaClient()  # type: ignore[assignment]
                return self._clients[provider]
            pcfg = self.config.providers[provider]
            self._clients[provider] = OpenAI(
                base_url=pcfg.base_url,
                api_key=pcfg.api_key(),
                default_headers=pcfg.default_headers or None,
                timeout=self.config.engine.request_timeout,
                max_retries=0,  # we manage retries ourselves for logging
            )
        return self._clients[provider]

    # -- core call ----------------------------------------------------------

    def complete(
        self,
        *,
        tier: str,
        task: str,
        system: str | None = None,
        user: str,
        messages: list[dict[str, str]] | None = None,
        json_mode: bool = False,
        max_tokens: int | None = None,
    ) -> LLMResponse:
        """Call chat completions for ``tier``, logging usage under ``task``."""
        tcfg, pcfg = self.config.resolve_tier(tier)
        client = self._client_for(tcfg.provider)

        msgs: list[dict[str, str]] = []
        if messages is not None:
            msgs = list(messages)
        else:
            if system:
                msgs.append({"role": "system", "content": system})
            msgs.append({"role": "user", "content": user})

        kwargs: dict[str, Any] = {
            "model": tcfg.model,
            "messages": msgs,
            "temperature": tcfg.temperature,
        }
        if tcfg.top_p is not None:
            kwargs["top_p"] = tcfg.top_p
        eff_max = max_tokens or tcfg.max_tokens
        if eff_max is not None:
            kwargs["max_tokens"] = eff_max
        if json_mode:
            kwargs["response_format"] = {"type": "json_object"}

        retries = 0
        last_exc: Exception | None = None
        max_retries = self.config.engine.max_retries
        while retries <= max_retries:
            try:
                resp = client.chat.completions.create(**kwargs)
                text = resp.choices[0].message.content or ""
                usage = resp.usage
                pt = getattr(usage, "prompt_tokens", 0) or 0
                ct = getattr(usage, "completion_tokens", 0) or 0
                tt = getattr(usage, "total_tokens", 0) or (pt + ct)
                self._log(task, tier, tcfg.provider, tcfg.model, pt, ct, tt,
                          ok=True, retries=retries)
                return LLMResponse(text, pt, ct, tt, retries)
            except _TRANSIENT as exc:  # transient -> backoff + retry
                last_exc = exc
                if retries >= max_retries:
                    break
                time.sleep(min(2 ** retries, 8) + 0.1)
                retries += 1
            except APIStatusError as exc:  # 4xx/5xx -> fail loudly
                self._log(task, tier, tcfg.provider, tcfg.model, 0, 0, 0,
                          ok=False, retries=retries)
                raise LLMError(
                    f"{task}: API error {exc.status_code} from {tcfg.provider}: "
                    f"{getattr(exc, 'message', exc)}"
                ) from exc

        self._log(task, tier, tcfg.provider, tcfg.model, 0, 0, 0,
                  ok=False, retries=retries)
        raise LLMError(f"{task}: exhausted retries against {tcfg.provider}: {last_exc}")

    # -- JSON helper --------------------------------------------------------

    def complete_json(
        self,
        *,
        tier: str,
        task: str,
        system: str | None = None,
        user: str,
        max_tokens: int | None = None,
    ) -> tuple[Any, LLMResponse]:
        """Call and parse a JSON response, stripping ``` fences defensively."""
        resp = self.complete(
            tier=tier, task=task, system=system, user=user,
            json_mode=True, max_tokens=max_tokens,
        )
        try:
            return parse_json(resp.text), resp
        except JSONParseError:
            # Some endpoints ignore json_mode; retry once without it then parse.
            resp = self.complete(
                tier=tier, task=task, system=system, user=user,
                json_mode=False, max_tokens=max_tokens,
            )
            return parse_json(resp.text), resp

    def _log(self, task: str, tier: str, provider: str, model: str,
             pt: int, ct: int, tt: int, *, ok: bool, retries: int) -> None:
        if self.usage is None:
            return
        self.usage.record(UsageRecord(
            world_id=self.world_id, session_id=self.session_id, task=task,
            tier=tier, provider=provider, model=model,
            prompt_tokens=pt, completion_tokens=ct, total_tokens=tt,
            ok=ok, retries=retries,
        ))


def strip_fences(text: str) -> str:
    s = text.strip()
    if s.startswith("```"):
        # remove leading and trailing fence lines
        lines = s.splitlines()
        if lines and lines[0].lstrip().startswith("```"):
            lines = lines[1:]
        if lines and lines[-1].strip().startswith("```"):
            lines = lines[:-1]
        s = "\n".join(lines)
    return s.strip()


def parse_json(text: str) -> Any:
    """Parse JSON from a model response, tolerating fences and surrounding prose."""
    candidate = strip_fences(text)
    try:
        return json.loads(candidate)
    except json.JSONDecodeError:
        pass
    # Fall back: grab the first balanced {...} or [...] span.
    for opener, closer in (("{", "}"), ("[", "]")):
        start = candidate.find(opener)
        end = candidate.rfind(closer)
        if start != -1 and end > start:
            try:
                return json.loads(candidate[start : end + 1])
            except json.JSONDecodeError:
                continue
    raise JSONParseError(f"could not parse JSON from response: {text[:200]!r}")