pablan/backend/app/llm/client.py

"""The ONLY code that talks to LLM endpoints.

Exactly three functions: chat_stream, chat_json, embed. Three model roles
(chat / utility / embedding), each base_url + api_key + model from settings.
The openai SDK is used purely as a client for OpenAI-compatible endpoints
(llama.cpp locally, cloud APIs in production).

Logging policy: metadata only — prompts and responses are logged ONLY at
DEBUG level behind PABLAN_DEBUG_LOG_PROMPTS=true (never in production).
LLMError messages are sanitized and never contain content.
"""

import logging
import time
from collections.abc import AsyncIterator
from functools import lru_cache
from typing import Any, Literal, TypeVar

import httpx
from openai import AsyncOpenAI
from pydantic import BaseModel, ValidationError

from app.config import get_settings
from app.llm.errors import LLMError, llm_error
from app.llm.gate import slot
from app.llm.overrides import env_defaults, get_config
from app.metrics import metrics

logger = logging.getLogger("pablan.llm")

Role = Literal["chat", "utility", "embedding"]
ChatMessage = dict[str, str]
T = TypeVar("T", bound=BaseModel)

# Turn off a reasoning model's hidden thinking. Latency-critical calls (a
# refinement fires on a typing pause) want the answer, not the deliberation:
# ~1s instead of ~10s with no quality loss on mechanical rewrites. Endpoints
# and templates that do not know the parameter ignore it.
NO_THINKING = {"chat_template_kwargs": {"enable_thinking": False}}

_RETRY_INSTRUCTION = (
    "Your previous reply did not match the required JSON schema. "
    "Reply again with ONLY valid JSON matching the schema — no prose."
)


def _http_client_factory() -> httpx.AsyncClient | None:
    """Tests override this to inject an ASGI transport."""
    return None


def role_config(role: Role) -> tuple[str, str, str]:
    """Effective endpoint config: the DB row, seeded from `.env` at first
    start (see app/llm/overrides.py — "bootstrap, then DB").

    The `or env` fallbacks are a safety net, not the model: they cover the
    window before `load_config()` has run (early startup, tests that never
    touch the table) and a field an admin blanked. In a bootstrapped
    instance the stored value always wins.
    """
    stored = get_config(role)
    env = env_defaults(role)
    return (
        stored.base_url or env.base_url or "",
        stored.api_key or env.api_key or "",
        stored.model or env.model or "",
    )


def _build_client(base_url: str, api_key: str) -> AsyncOpenAI:
    kwargs: dict[str, Any] = {
        "base_url": base_url,
        "api_key": api_key,
        "timeout": get_settings().llm_timeout_seconds,
        # One SDK retry: llama.cpp closes idle keep-alive connections, and
        # the first call on a stale connection fails with APIConnectionError.
        # (SDK-internal retries are not separately metered.)
        "max_retries": 1,
    }
    http_client = _http_client_factory()
    if http_client is not None:
        kwargs["http_client"] = http_client
    return AsyncOpenAI(**kwargs)


@lru_cache(maxsize=None)
def _client_for(role: Role) -> AsyncOpenAI:
    base_url, api_key, _ = role_config(role)
    return _build_client(base_url, api_key)


def rebuild_clients() -> None:
    """Apply changed endpoint config without a restart: the cached clients
    hold the old base_url and key, so they must go."""
    _client_for.cache_clear()


async def probe(
    role: Role,
    *,
    base_url: str | None = None,
    api_key: str | None = None,
    model: str | None = None,
) -> None:
    """Smallest possible call against a candidate config, so an admin can
    test an endpoint before saving it. Raises LLMError on failure."""
    effective_url, effective_key, effective_model = role_config(role)
    client = _build_client(base_url or effective_url, api_key or effective_key)
    target = model or effective_model
    started = time.monotonic()
    try:
        if role == "embedding":
            await client.embeddings.create(model=target, input=["ping"])
        else:
            stream = await client.chat.completions.create(
                model=target,
                messages=[{"role": "user", "content": "ping"}],
                max_tokens=1,
                stream=True,
            )
            async for _ in stream:
                break
    except Exception as exc:
        raise llm_error("probe", role, exc, started) from None


async def list_models(
    role: Role,
    *,
    base_url: str | None = None,
    api_key: str | None = None,
) -> list[str]:
    """Ask an endpoint what it serves (`GET /v1/models`).

    Server-side on purpose: the credentials must never leave the backend,
    and the browser has no business talking to the model endpoint at all.

    Not every OpenAI-compatible server implements the route, so a failure
    here is ordinary rather than exceptional — the caller degrades to a
    free-text model field. Raises LLMError so the caller can distinguish
    "no such route" from "wrong credentials".
    """
    effective_url, effective_key, _ = role_config(role)
    client = _build_client(base_url or effective_url, api_key or effective_key)
    started = time.monotonic()
    try:
        page = await client.models.list()
    except Exception as exc:
        raise llm_error("list_models", role, exc, started) from None
    # Ids only, sorted for a stable dropdown. Model ids are configuration,
    # not content, so they may be returned and logged by count.
    return sorted({model.id for model in page.data if getattr(model, "id", None)})


def _record(
    role: Role,
    kind: str,
    status: str,
    started: float,
    usage: Any = None,
    **extra_fields: Any,
) -> None:
    duration = time.monotonic() - started
    metrics.inc("llm_calls_total", {"role": role, "kind": kind, "status": status})
    metrics.observe("llm_call_seconds", duration, {"role": role, "kind": kind})
    extra: dict[str, Any] = {
        "event": "llm_call",
        "role": role,
        "kind": kind,
        "status": status,
        "duration_ms": round(duration * 1000),
        **extra_fields,
    }
    if usage is not None:
        prompt_tokens = getattr(usage, "prompt_tokens", None)
        completion_tokens = getattr(usage, "completion_tokens", None)
        if prompt_tokens:
            metrics.inc(
                "llm_tokens_total", {"role": role, "direction": "prompt"}, prompt_tokens
            )
            extra["prompt_tokens"] = prompt_tokens
        if completion_tokens:
            metrics.inc(
                "llm_tokens_total",
                {"role": role, "direction": "completion"},
                completion_tokens,
            )
            extra["completion_tokens"] = completion_tokens
    logger.info("llm call", extra=extra)


def _debug_log_content(label: str, content: Any) -> None:
    if get_settings().debug_log_prompts:
        logger.debug("llm content", extra={"label": label, "content": content})


async def chat_stream(
    messages: list[ChatMessage],
    *,
    role: Role = "chat",
    temperature: float | None = None,
    max_tokens: int | None = None,
    extra_body: dict[str, Any] | None = None,
) -> AsyncIterator[str]:
    """Stream a chat completion as text deltas.

    `extra_body` is passed through to the endpoint verbatim — used to reach
    non-standard OpenAI-compatible parameters such as
    `{"chat_template_kwargs": {"enable_thinking": False}}`, which turns off a
    reasoning model's hidden thinking for latency-critical calls. Only
    `delta.content` is ever yielded, so a reasoning channel never leaks into
    the output regardless.
    """
    base_url, _, model = role_config(role)
    _debug_log_content("chat_stream.messages", messages)
    options: dict[str, Any] = {}
    if temperature is not None:
        options["temperature"] = temperature
    if max_tokens is not None:
        options["max_tokens"] = max_tokens
    if extra_body is not None:
        options["extra_body"] = extra_body

    started = time.monotonic()
    status = "ok"
    usage = None
    try:
        # The slot is held until the last token: a streaming completion
        # occupies its server slot for its whole life (app/llm/gate.py).
        async with slot(base_url, role):
            stream = await _client_for(role).chat.completions.create(
                model=model,
                messages=messages,  # type: ignore[arg-type]
                stream=True,
                stream_options={"include_usage": True},
                **options,
            )
            async for chunk in stream:
                if chunk.usage is not None:
                    usage = chunk.usage
                if chunk.choices and chunk.choices[0].delta.content:
                    yield chunk.choices[0].delta.content
    except GeneratorExit:
        status = "aborted"
        raise
    except LLMError:
        status = "error"
        raise
    except Exception as exc:
        status = "error"
        raise llm_error("chat_stream", role, exc, started) from None
    finally:
        _record(role, "chat_stream", status, started, usage)


async def chat_json(
    messages: list[ChatMessage],
    schema: type[T],
    *,
    role: Role = "utility",
    temperature: float = 0.0,
    max_tokens: int | None = None,
    extra_body: dict[str, Any] | None = None,
) -> T:
    """Structured output: response_format JSON schema + validation + one retry.

    `extra_body` is passed through verbatim (e.g.
    `{"chat_template_kwargs": {"enable_thinking": False}}` to skip a reasoning
    model's hidden thinking on latency-sensitive utility calls)."""
    base_url, _, model = role_config(role)
    response_format = {
        "type": "json_schema",
        "json_schema": {
            "name": schema.__name__,
            "schema": schema.model_json_schema(),
            "strict": True,
        },
    }
    options: dict[str, Any] = {"temperature": temperature}
    if max_tokens is not None:
        options["max_tokens"] = max_tokens
    if extra_body is not None:
        options["extra_body"] = extra_body

    attempt_messages = list(messages)
    for attempt in (1, 2):
        _debug_log_content("chat_json.messages", attempt_messages)
        started = time.monotonic()
        usage = None
        try:
            async with slot(base_url, role):
                response = await _client_for(role).chat.completions.create(
                    model=model,
                    messages=attempt_messages,  # type: ignore[arg-type]
                    response_format=response_format,  # type: ignore[arg-type]
                    **options,
                )
            usage = response.usage
            content = response.choices[0].message.content or ""
            result = schema.model_validate_json(content)
            _record(role, "chat_json", "ok", started, usage, attempt=attempt)
            return result
        except ValidationError:
            _record(role, "chat_json", "invalid", started, usage, attempt=attempt)
            _debug_log_content("chat_json.invalid_response", content)
            attempt_messages = [
                *attempt_messages,
                {"role": "assistant", "content": content},
                {"role": "user", "content": _RETRY_INSTRUCTION},
            ]
        except LLMError:
            _record(role, "chat_json", "error", started, usage, attempt=attempt)
            raise
        except Exception as exc:
            _record(role, "chat_json", "error", started, usage, attempt=attempt)
            raise llm_error("chat_json", role, exc, started, attempt=attempt) from None

    raise LLMError(
        f"chat_json failed (role={role}): response did not match schema "
        f"{schema.__name__} after retry",
        role=role,
        kind="chat_json",
        status="invalid",
        cause_type="ValidationError",
        duration_ms=round((time.monotonic() - started) * 1000),
        attempt=2,
    )


async def embed(texts: list[str], *, role: Role = "embedding") -> list[list[float]]:
    """Embed a batch of texts; order of results matches the input order."""
    base_url, _, model = role_config(role)
    started = time.monotonic()
    try:
        async with slot(base_url, role):
            response = await _client_for(role).embeddings.create(
                model=model, input=texts
            )
    except LLMError:
        _record(role, "embed", "error", started, batch_size=len(texts))
        raise
    except Exception as exc:
        _record(role, "embed", "error", started, batch_size=len(texts))
        raise llm_error("embed", role, exc, started) from None
    _record(role, "embed", "ok", started, response.usage, batch_size=len(texts))
    ordered = sorted(response.data, key=lambda item: item.index)
    return [item.embedding for item in ordered]