"""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]