# SPDX-License-Identifier: BUSL-1.1 # Copyright (c) 2024-2026 Ali Sadeghi Aghili """Abstract base for all LLM backends. To add a new backend (OpenAI, Anthropic, vLLM, …) implement ``LLMBackend`` and register it in ``llm/sql_agent.py``. """ from __future__ import annotations import asyncio from abc import ABC, abstractmethod from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any if TYPE_CHECKING: # Annotation-only: llm.router imports this module, so a runtime import # here would be circular (see generate_structured's docstring). from llm.router import PromptSegments @dataclass(frozen=True) class SQLGenerationResult: """Outcome of one LLM call.""" sql: str """Cleaned, sanitised SQL string.""" raw_response: str """Verbatim text the model returned (useful for debugging).""" attempt: int = 1 """Which attempt produced this result (1 = first try, 2+ = correction).""" correction_prompts: list[str] = field(default_factory=list) """The correction prompts sent in previous failed attempts.""" llm_meta: dict[str, Any] = field(default_factory=dict) """Best-effort transport metadata for the call that produced this result — whatever :meth:`LLMBackend.generate_with_meta` returned alongside the text (e.g. the OpenAI-compatible response's raw ``usage.prompt_tokens`` / ``usage.completion_tokens`` fields). Empty for backends that don't override :meth:`LLMBackend.generate_with_meta`. Consumed by ``api/runner.py`` to build the ``docs/api-contract-v2.md`` §6 ``llm`` audit block without the backend needing any per-instance mutable state — see :meth:`LLMBackend.generate_with_meta`'s docstring.""" injected_top: int | None = None """The row cap :func:`~security.sql_guard.ensure_top` injected into ``sql``, or ``None`` if the model's own SQL already carried a row-limit clause (nothing was injected). Set by ``SQLAgent._clean_validate_cap`` by comparing ``ensure_top``'s output to its input — ``ensure_top`` itself only returns the capped string, so this is the caller noticing whether that string actually changed. Consumed by ``api/runner.py`` to populate ``docs/api-contract-v2.md`` §4's ``guard.injected_top``: it tells a reader whether the number they are looking at was truncated.""" datasource_selection: dict[str, Any] | None = None """How the data source was chosen when several are configured -- ``{"chosen", "reason", "candidates", "fallback_from"}``, see :meth:`~llm.source_routing.SourceRouting.audit`; ``chosen`` is the source the answer was generated for (the retry's, when the model declined the first). ``None`` for a deployment with one data source. Set by :meth:`~llm.sql_agent.SQLAgent.run` and consumed by ``api/runner.py`` for the audit record.""" value_check: dict[str, Any] | None = None """What the filter value check (``FILTER_VALUE_CHECK``, see :mod:`retrieval.filter_values`) found, as :func:`~retrieval.filter_values.audit_block` writes it for the audit record; ``None`` with the setting off, and when no literal could be checked.""" value_spelling: dict[str, Any] | None = None """What the filter value spelling fix (``FILTER_VALUE_SPELLING``, see :mod:`retrieval.filter_spelling`) did to the statement that ran, as :func:`~retrieval.filter_spelling.audit_block` writes it for the audit record; ``None`` with the setting off, and when no literal was rewritten. ``sql`` is then the rewritten statement, ``raw_response`` still the model's own text.""" class LLMBackend(ABC): """Contract every LLM backend must satisfy. Backends are stateless — all state lives in the prompt. ``generate`` is the only required method; it must: * Accept a complete prompt string. * Return the model's raw text output (no cleaning). * Raise ``RuntimeError`` on unrecoverable transport/API failures. * Raise ``ValueError("OUT_OF_SCOPE")`` when the model signals it cannot answer (pass the sentinel through unchanged). """ @abstractmethod def generate(self, prompt: str) -> str: """Send *prompt* to the model and return raw text.""" ... def generate_with_meta( self, prompt: str, *, max_tokens: int | None = None, # noqa: ARG002 ) -> tuple[str, dict[str, Any]]: """Send *prompt* to the model; return ``(raw_text, meta)``. *meta* is an opt-in, backend-specific dict of transport metadata useful for observability (e.g. token counts, timings, HTTP status) — see ``observability/llm_status.py::build_llm_status`` for the shape a caller typically wants to build from it. This is a **stateless** call: metadata for one call is returned to that call's caller, never stashed on ``self``. ``SQLAgent`` is a long-lived singleton shared across concurrent requests (see ``api/runner.py``'s module docstring), so any backend that cached "the last response" as an instance attribute would race across requests; threading the metadata through the return value instead of instance state is what keeps this safe. The default implementation here makes the metadata path fully opt-in: a backend that only implements :meth:`generate` (the sole abstract method) keeps working unchanged, just with an empty *meta* dict — exactly what :class:`~llm.providers.MockBackend` does, and what :class:`~llm.providers.OpenAIBackend` overrides. Parameters ---------- prompt: Complete prompt string, as passed to :meth:`generate`. max_tokens: A per-call ceiling on the completion length, for a caller that knows its answer is short and wants a runaway one stopped early (the setup wizard's alias call). It is a request, not a contract: this default has no way to apply it and ignores it, and a backend that can (:class:`~llm.providers.OpenAIBackend`) overrides this method. ``None`` leaves the backend's own setting in force. Returns ------- tuple[str, dict[str, Any]] ``(raw_text, meta)`` — *raw_text* is exactly what :meth:`generate` would have returned; *meta* is ``{}`` unless overridden. Examples -------- >>> class EchoBackend(LLMBackend): ... def generate(self, prompt: str) -> str: ... return prompt.upper() >>> text, meta = EchoBackend().generate_with_meta("hi") >>> text, meta ('HI', {}) """ return self.generate(prompt), {} async def agenerate_with_meta(self, prompt: str) -> tuple[str, dict[str, Any]]: """Async counterpart of :meth:`generate_with_meta`. Phase 2 (latency) task 4 moves the API layer to ``async def`` handlers so a slow LLM call no longer ties up one of Starlette's threadpool workers for the duration of the request. A backend that implements a genuine async transport (built on ``httpx.AsyncClient`` with connection keep-alive, say) should override this method directly. The default implementation here — like :meth:`generate_with_meta`'s default — makes async support fully opt-in: a backend that only implements the synchronous :meth:`generate` keeps working unchanged, just via :func:`asyncio.to_thread`, which still frees the event loop (the call runs in a worker thread) even though the transport itself remains blocking underneath. Returns ------- tuple[str, dict[str, Any]] Same contract as :meth:`generate_with_meta`. Examples -------- >>> import asyncio >>> class EchoBackend(LLMBackend): ... def generate(self, prompt: str) -> str: ... return prompt.upper() >>> text, meta = asyncio.run(EchoBackend().agenerate_with_meta("hi")) >>> text, meta ('HI', {}) """ return await asyncio.to_thread(self.generate_with_meta, prompt) def generate_with_meta_segments( self, segments: "PromptSegments" ) -> tuple[str, dict[str, Any]]: """Segment-aware counterpart of :meth:`generate_with_meta`. This is what :meth:`~llm.router.LLMRouter.generate_for_task` calls — never ``generate_with_meta`` directly — so that the router itself never flattens :class:`~llm.router.PromptSegments` before it reaches a backend (Phase 2 task 5's third design point: "prompts enter the router segmented, never as one flat string"). Whether a given backend then chooses to flatten internally is that backend's own decision, not the router's. The default implementation here does flatten (via :meth:`~llm.router.PromptSegments.flatten`) and delegate to :meth:`generate_with_meta` — the correct behaviour for a backend with no segment-aware transport, and in particular for :class:`~llm.providers.OpenAIBackend`: OpenAI-compatible prefix caching (automatic on OpenAI's own API, implicit KV-cache reuse on llama.cpp/vLLM) is byte-identical-prefix-based, exactly like Ollama's before it — flattening the segments in a stable order *is* the correct way to get that benefit, not a compromise. A backend with an explicit, segment-level caching mechanism (e.g. an Anthropic-style ``cache_control`` breakpoint on the static prefix block) would override this method instead of relying on the default. Parameters ---------- segments: A :class:`~llm.router.PromptSegments`, accepted as a duck-typed object exposing ``.flatten() -> str`` for the same circular-import reason documented on :meth:`generate_structured`. Returns ------- tuple[str, dict[str, Any]] Same contract as :meth:`generate_with_meta`. Examples -------- >>> class EchoBackend(LLMBackend): ... def generate(self, prompt: str) -> str: ... return prompt.upper() >>> class FakeSegments: ... def flatten(self) -> str: ... return "hi" >>> text, meta = EchoBackend().generate_with_meta_segments(FakeSegments()) >>> text, meta ('HI', {}) """ return self.generate_with_meta(segments.flatten()) def warm_prefix(self, prefix: str, *, timeout: float | None = None) -> dict[str, Any]: """Prime the model server's prefix cache with *prefix*; return what it reported. Sends *prefix* as the whole prompt of a request that generates one token, so a server that caches prompt prefixes (vLLM with ``--enable-prefix-caching``, llama.cpp's prompt cache) has already done the prefill when the first real question, which starts with the same bytes, arrives. See :mod:`llm.warmup`. The default raises :class:`NotImplementedError`: a backend with no real server behind it (a stub, a test double) has no cache to warm, and quietly generating a full answer instead would cost what the warm-up exists to save. :class:`~llm.providers.OpenAIBackend` overrides it. Parameters ---------- prefix: The static prompt prefix, exactly as the real requests start. timeout: Request timeout in seconds; the backend's own when ``None``. Returns ------- dict[str, Any] Whatever the backend learned, at least ``prompt_tokens`` (an ``int`` or ``None``). Never the prompt or response text. Raises ------ NotImplementedError For a backend with nothing to warm. Examples -------- >>> class Stub(LLMBackend): ... def generate(self, prompt: str) -> str: ... return "" >>> Stub().warm_prefix("p") Traceback (most recent call last): ... NotImplementedError: Stub has no prefix cache to warm """ raise NotImplementedError(f"{type(self).__name__} has no prefix cache to warm") def generate_structured( self, segments: "PromptSegments", schema: dict[str, Any] ) -> tuple[dict[str, Any], dict[str, Any]]: """Return an object matching *schema*, not a raw string. This is the interface :mod:`llm.router` calls — "give me an object matching this schema", not "here is a prompt string" (Phase 2 task 5's second design point). Constrained decoding is provider-specific (OpenAI-compatible ``response_format: json_schema``, Ollama's ``format``, Anthropic's forced tool use), so a backend that supports it natively — e.g. :class:`~llm.providers.OpenAIBackend` — overrides this method. This default makes it fully opt-in for everything else: flatten *segments* into one prompt, append a plain-text JSON instruction, call :meth:`generate`, and best-effort parse the result with :func:`llm.providers.parse_json_response` — exactly the "every provider collapses to string parsing" fallback the design note warns is the wrong DEFAULT behaviour to build the whole system around, but which remains a legitimate degrade path for a provider that genuinely has no structured-output mode. Parameters ---------- segments: A :class:`~llm.router.PromptSegments` (``static_prefix``, ``session_context``, ``question``) — accepted as a duck-typed object exposing ``.flatten() -> str`` here to avoid a circular import with :mod:`llm.router`, which imports this module. schema: A JSON Schema describing the expected object shape. Returns ------- tuple[dict[str, Any], dict[str, Any]] ``(parsed_object, meta)``. Raises ------ ValueError If the response cannot be parsed as JSON matching *schema* (schema conformance itself is NOT validated here — only that valid JSON was extracted; callers needing strict validation should check the result themselves). """ from llm.providers import parse_json_response # noqa: PLC0415 - avoid import cycle instruction = ( "\n\nRespond with ONLY a single JSON object matching this schema " f"(no markdown fences, no explanation): {schema}" ) raw, meta = self.generate_with_meta(segments.flatten() + instruction) parsed = parse_json_response(raw) if not isinstance(parsed, dict): raise ValueError(f"Expected a JSON object, got {type(parsed).__name__}: {raw[:200]!r}") meta = {**meta, "structured_output": False} return parsed, meta @property def name(self) -> str: """Human-readable identifier shown in logs and REPL header.""" return type(self).__name__ @property def trusted(self) -> bool: """Whether this backend's endpoint may see schema/business-rule/row data. Defaults to ``True``. This is deliberately *not* the same default the endpoint-trust computation itself uses for a real network transport (see :func:`~llm.trust.default_trust_for_url`, which defaults an *unrecognised base_url* to untrusted) — this default is for backends that have no endpoint at all: a test double, a stub, or a bespoke :class:`LLMBackend` subclass someone writes for their own local model. None of those are the thing :meth:`~llm.router.LLMRouter._governance_check` exists to catch (a real hosted API that would otherwise receive data silently), so they should not need any ceremony (constructing an :class:`~llm.endpoints.EndpointConfig`, computing a trust flag) to stay usable. A backend that talks to a real, configurable endpoint — :class:`~llm.providers.OpenAIBackend` — overrides this to report its own resolved per-endpoint trust instead of inheriting this default. """ return True @property def endpoint(self) -> str | None: """The endpoint identifier shown in the ``llm`` status block, or ``None``. ``None`` for a backend with no real network endpoint (a test double, or a bespoke subclass that doesn't override this). A backend with a genuine transport — :class:`~llm.providers.OpenAIBackend` — overrides this to return its ``base_url``. """ return None