Files
brain-of-reese/app/rag/llm.py
T

567 lines
24 KiB
Python

"""Async OpenAI-compatible client for the self-hosted aipi endpoint (PLAN A5).
Provides the embeddings surface (importer, retrieval), one-shot chat
completions (phase 30: the ``lite`` model summarizes non-markdown
documents at import time), and chat streaming (PLAN A15) for the RAG
pipeline. Chat streaming yields typed :class:`StreamPiece` values
(phase 17) and — when the caller passes a ``tools`` list —
:class:`ToolCallPiece` values (phase 37): aipi's ``turbo`` model streams
its reasoning as ``delta.reasoning_content`` chunks (deepseek/litellm
wire convention, verified live 2026-08-23) **before** the answer's
``delta.content`` chunks, and reasoning counts against ``max_tokens``
(an answer can in principle be empty).
Fail-loud rule (PLAN A6): the ``chunks.embedding`` column is fixed at 768
dimensions when the table is created, so a model that returns a different
dimension must abort the import with an actionable error — never store
vectors that pgvector rejects.
"""
from __future__ import annotations
import asyncio
import json
import logging
from collections.abc import AsyncGenerator
from dataclasses import dataclass
from typing import Any, Literal, cast
from openai import AsyncOpenAI, AsyncStream
from openai.types.chat import ChatCompletionChunk, ChatCompletionMessageParam
from app.config import Settings, get_settings
logger = logging.getLogger("app.llm")
class EmbeddingError(RuntimeError):
"""The embeddings endpoint failed (network, HTTP, or malformed reply)."""
class EmbeddingDimensionError(EmbeddingError):
"""Embedding dimension != BOR_EMBEDDING_DIM — import must fail loudly."""
class LLMError(RuntimeError):
"""The chat-completions endpoint failed (network, HTTP, or mid-stream)."""
class ModelUnavailableError(LLMError):
"""One of the models a sync needs is unreachable (phase 41 probe).
Raised by :func:`check_models` when the pre-sync probe finds the
embedding or summary model down; the message names the model so the
admin can fix the right thing.
"""
@dataclass(frozen=True)
class StreamPiece:
"""One piece of a streamed chat turn (phase 17, PLAN §4 extension).
``kind`` is ``"content"`` for answer text (an SSE ``delta`` frame)
or ``"thinking"`` for the model's reasoning (an SSE ``thinking``
frame). Frozen: pieces are immutable wire values, not accumulators.
"""
kind: Literal["content", "thinking"]
text: str
@dataclass(frozen=True)
class ToolCallPiece:
"""One model-requested tool call accumulated from stream deltas (phase 37).
``id`` is the model's tool_call id (synthesized as ``call_<index>``
when the wire never carried one), ``name`` is the function name
(whatever the caller's ``tools`` list names — for the agent loop,
``list_documents`` / ``read_document``), and ``arguments`` is the
parsed JSON object (``{}`` when the model sent none).
"""
id: str # the model's tool_call id; synthesized "call_<index>" when absent
name: str # "list_documents" | "read_document" (whatever AGENT_TOOLS names)
arguments: dict[str, Any]
@dataclass(frozen=True)
class RetryPiece:
"""One LLM request retry that is about to start (phase 67, locked A2).
``attempt`` is the 1-based number of the attempt that is about to be
tried — the one AFTER the attempt that just failed (a first-attempt
failure carries ``attempt=2``, so the API's SSE ``retry`` frame reads
"retrying (2 of N)" — the same convention the endpoint's embedding
retry loop uses, phase 67 task 02); ``max_attempts`` is the total
attempt budget (``llm_retries + 1``). One piece per wait: the API
layer turns it into an SSE ``retry`` frame, and it always precedes
the pre-retry sleep so the frame reaches the client before the wait
starts.
"""
attempt: int
max_attempts: int
@dataclass
class _ToolCallSlot:
"""Mutable accumulator for one streamed tool call (phase 37, private).
``id`` and ``function.name`` arrive on the first partial for an index;
``function.arguments`` arrives in fragments to concatenate (OpenAI wire
convention, verified live against aipi 2026-08-26).
"""
id: str | None = None
name: str = ""
arguments: str = ""
def _materialize_tool_calls(
slots: dict[int, _ToolCallSlot],
) -> list[ToolCallPiece]:
"""Turn accumulated slots into ordered :class:`ToolCallPiece` values.
Malformed ``arguments`` JSON raises :class:`LLMError` — a silently
dropped tool call would corrupt the agent loop (fail-loud house
style). Empty/``null`` arguments become ``{}`` (a no-parameter call
such as ``list_documents``).
"""
pieces: list[ToolCallPiece] = []
for index in sorted(slots):
slot = slots[index]
raw = slot.arguments.strip()
label = slot.name or f"index {index}"
if raw:
try:
parsed: Any = json.loads(raw)
except json.JSONDecodeError as e:
raise LLMError(
f"model sent malformed tool-call arguments for '{label}': "
f"{raw[:200]!r} ({e})"
) from e
else:
parsed = None
if parsed is None:
arguments: dict[str, Any] = {}
elif isinstance(parsed, dict):
arguments = cast("dict[str, Any]", parsed)
else:
raise LLMError(
f"model sent non-object tool-call arguments for '{label}': "
f"{raw[:200]!r}"
)
pieces.append(
ToolCallPiece(
id=slot.id or f"call_{index}", name=slot.name, arguments=arguments
)
)
return pieces
# aipi's local embedding model rejects requests over ~1024 input tokens
# ("input is too large to process"). Batch by estimated tokens, with a
# safety margin under that cap — code-dense text can tokenize at ~3
# chars/token, so stay well below the ceiling. A retry that halves an
# over-large batch (:meth:`LLMClient._embed_batch`) is the backstop.
_TOKENS_PER_CHAR = 0.25
_REQUEST_TOKEN_BUDGET = 700
class _TooLarge(RuntimeError):
"""Internal: the endpoint rejected the request's input size."""
class LLMClient:
"""Thin async wrapper over the aipi OpenAI-compatible API."""
def __init__(self, settings: Settings | None = None) -> None:
self.settings = settings or get_settings()
#: Number of embedding HTTP requests made so far (importer logging).
self.embed_batches: int = 0
self._client = AsyncOpenAI(
base_url=self.settings.llm_base_url,
api_key=self.settings.effective_api_key,
timeout=120.0,
)
async def _post_embeddings(self, texts: list[str]) -> list[list[float]]:
"""One POST /embeddings with a minimal OpenAI-compatible payload.
The ``openai`` SDK (1.x and 2.x) injects ``encoding_format`` into
every embeddings request (defaulting to ``"base64"``), and aipi's
litellm ``openai_like`` model group rejects that parameter in any
form — so we reuse the openai client's own httpx transport (same
base URL, TLS, and connection pooling) and send a clean payload.
The endpoint's default is floats, which is what pgvector needs.
"""
http = self._client._client # pyright: ignore[reportAttributeAccessIssue]
resp = await http.post(
"embeddings",
json={"model": self.settings.llm_embed_model, "input": texts},
headers={"Authorization": f"Bearer {self.settings.effective_api_key}"},
)
if resp.status_code >= 400:
if "too large" in resp.text:
raise _TooLarge(resp.text[:300])
raise EmbeddingError(
f"embeddings endpoint returned HTTP {resp.status_code}: {resp.text[:300]}"
)
payload = resp.json()
rows = payload.get("data") if isinstance(payload, dict) else None
if not isinstance(rows, list) or len(rows) != len(texts):
n = len(rows) if isinstance(rows, list) else "?"
raise EmbeddingError(
f"embeddings endpoint returned {n} vectors for {len(texts)} inputs — "
"refusing to guess which is which"
)
ordered = sorted(rows, key=lambda r: r["index"])
return [[float(x) for x in row["embedding"]] for row in ordered]
async def _embed_batch(self, chunk: list[str]) -> list[list[float]]:
"""Embed *chunk*, halving the request if the endpoint says the input
is too large (tokenizer estimates can be wrong for dense content).
A single text that still fails is a hard, actionable error."""
try:
return await self._post_embeddings(chunk)
except _TooLarge:
if len(chunk) == 1:
raise EmbeddingError(
f"a single {len(chunk[0])}-char chunk exceeded the endpoint's "
"per-request input token cap — lower BOR_CHUNK_TARGET_CHARS "
"and re-import"
) from None
mid = len(chunk) // 2
left = await self._embed_batch(chunk[:mid])
right = await self._embed_batch(chunk[mid:])
return [*left, *right]
def _check_dims(self, vecs: list[list[float]]) -> None:
expected = self.settings.embedding_dim
dims = sorted({len(v) for v in vecs})
if dims != [expected]:
raise EmbeddingDimensionError(
f"embedding dimension mismatch: model '{self.settings.llm_embed_model}' "
f"returned dims {dims} but BOR_EMBEDDING_DIM={expected}. The chunks "
"table stores a fixed dimension — run `uv run python -m scripts.llm_probe`,"
" update BOR_EMBEDDING_DIM, and recreate the schema (README → "
"Troubleshooting: 'Embedding dimension mismatch')."
)
def _batch_texts(self, texts: list[str]) -> list[list[str]]:
"""Group *texts* into requests under the endpoint's token cap.
``BOR_EMBED_BATCH_SIZE`` stays a hard cap on *texts* per request;
the token budget usually binds first for 2000-char chunks.
"""
budget_chars = int(_REQUEST_TOKEN_BUDGET / _TOKENS_PER_CHAR)
max_texts = max(1, self.settings.embed_batch_size)
batches: list[list[str]] = []
cur: list[str] = []
cur_chars = 0
for text in texts:
if cur and (len(cur) >= max_texts or cur_chars + len(text) > budget_chars):
batches.append(cur)
cur, cur_chars = [], 0
cur.append(text)
cur_chars += len(text)
if cur:
batches.append(cur)
return batches
async def embed(self, texts: list[str]) -> list[list[float]]:
"""Embed *texts* in token-budgeted batches (order always kept)."""
if not texts:
return []
out: list[list[float]] = []
for chunk in self._batch_texts(texts):
try:
vecs = await self._embed_batch(chunk)
except EmbeddingError:
raise
except Exception as e: # noqa: BLE001 — wrap transport-level failures
raise EmbeddingError(
f"embeddings request to {self.settings.llm_base_url} failed: {e}"
) from e
self.embed_batches += 1
self._check_dims(vecs)
out.extend(vecs)
return out
async def embed_one(self, text: str) -> list[float]:
"""Convenience: embed a single text (retrieval path)."""
(vec,) = await self.embed([text])
return vec
async def chat(
self, messages: list[dict[str, str]], model: str | None = None
) -> str:
"""One-shot (non-streaming) completion (A5 extended, phase 30).
Short, low-temperature request (``temperature=0.2``, 2048-token
cap — summaries and outlines are small, so a fixed budget is
enough) against ``BOR_LLM_SUMMARY_MODEL`` (default ``lite``)
unless *model* names another. Used by the document summarizer
(phase 30) and the KB overview generator (phase 31).
Any transport/HTTP/malformed failure, a choiceless reply, or an
empty/missing ``content`` field raises :class:`LLMError` — a
silent empty summary must never be stored.
"""
try:
resp = await self._client.chat.completions.create(
model=model or self.settings.llm_summary_model,
messages=cast("list[ChatCompletionMessageParam]", messages),
temperature=0.2,
max_tokens=2048,
stream=False,
)
except LLMError:
raise
except Exception as e: # noqa: BLE001 — wrap transport-level failures
raise LLMError(
f"chat completion from {self.settings.llm_base_url} failed: {e}"
) from e
if not resp.choices:
raise LLMError(
f"chat completion from {self.settings.llm_base_url} "
"returned no choices"
)
content = resp.choices[0].message.content
if content is None or not content.strip():
raise LLMError(
f"chat completion from {self.settings.llm_base_url} returned "
"empty content — refusing to store a silent summary"
)
return content.strip()
async def chat_stream(
self,
messages: list[dict[str, str]],
tools: list[dict[str, Any]] | None = None,
) -> AsyncGenerator[StreamPiece | ToolCallPiece, None]:
"""Stream assistant pieces from the chat model (PLAN A5/A15, phase 17).
``stream=True`` against the OpenAI-compatible endpoint, yielding
typed :class:`StreamPiece` values. Wire convention (verified live
against aipi's ``turbo`` on 2026-08-23): the model's reasoning
arrives as ``delta.reasoning_content`` chunks (deepseek/litellm
convention) **before** the first ``delta.content`` chunk, so in
practice thinking pieces precede content pieces. The ``openai``
SDK keeps unknown delta fields in ``model_extra``, so ``getattr``
is the right accessor — no raw-HTTP parsing is needed. A chunk
carrying both fields yields the thinking piece **first**.
Reasoning counts against ``max_tokens``: an answer can in principle
be empty (thinking with no content) — the UI handles that.
Answers are allowed up to ``BOR_MAX_OUTPUT_TOKENS`` (default
32 768) output tokens — the old hard 700-token cap cut long
answers off mid-sentence (owner report 2026-08-22).
Tool calls (phase 37): when *tools* (an OpenAI ``tools`` list) is
not None it is passed through as ``tools=…``; when None the key is
**not** included, so the request is byte-identical to pre-phase-37
and no tool pieces can be produced. A tool-calling model replies
with ``delta.tool_calls`` partials — keyed by ``index``, with
``id`` and ``function.name`` on the first partial and
``function.arguments`` in fragments — which are accumulated into
one :class:`ToolCallPiece` per call, yielded in index order at
stream end (or immediately once a chunk carries
``finish_reason="tool_calls"``). Malformed ``arguments`` JSON
raises :class:`LLMError`. Wire convention verified live against
aipi's ``turbo`` on 2026-08-26 via
``uv run python -m scripts.llm_probe --tools`` (phase 37, task 01:
``probe: turbo tool_calls=supported 2026-08-26``).
Any failure (network, HTTP, malformed stream) surfaces as
:class:`LLMError` so the API layer can turn it into an SSE
``error`` event instead of a hung request.
Teardown (phase 48, 2026-08-29, ``TODO.md`` L3): once
``create()`` succeeded, the endpoint stream's lifetime is
explicit — it is closed on **every** exit: normal exhaustion
(a quiet no-op on the already-ended SDK stream, so the
completed path stays byte-identical), a wrapped mid-stream
failure, and consumer abandon (stop/cancel — ``GeneratorExit``;
awaiting in the ``finally`` is safe because it does not yield).
The SDK's ``close()`` awaits the underlying httpx response's
``aclose()``, so the local model stops generating as soon as
the SSE consumer goes away.
"""
kwargs: dict[str, Any] = {
# ``{role, content}`` dicts are exactly what the message params
# accept; the cast keeps pyright honest about the SDK's union.
"model": self.settings.llm_chat_model,
"messages": cast("list[ChatCompletionMessageParam]", messages),
"temperature": 0.4,
"max_tokens": self.settings.max_output_tokens,
"stream": True,
}
if tools is not None:
kwargs["tools"] = tools
stream: AsyncStream[ChatCompletionChunk] | None = None
try:
stream = cast(
"AsyncStream[ChatCompletionChunk]",
await self._client.chat.completions.create(**kwargs),
)
calls: dict[int, _ToolCallSlot] = {}
emitted = False
async for chunk in stream:
if not chunk.choices:
continue
choice = chunk.choices[0]
delta = choice.delta
# Tool-call partials (phase 37) accumulate across chunks,
# keyed by index; a missing index (not seen on aipi) falls
# back to the next synthetic slot.
for tc in getattr(delta, "tool_calls", None) or []:
idx = getattr(tc, "index", None)
key = idx if isinstance(idx, int) else (max(calls) + 1 if calls else 0)
slot = calls.setdefault(key, _ToolCallSlot())
tc_id = getattr(tc, "id", None)
if tc_id and slot.id is None:
slot.id = tc_id
fn = getattr(tc, "function", None)
if fn is not None:
if fn.name:
slot.name += fn.name
if fn.arguments:
slot.arguments += fn.arguments
reasoning = getattr(delta, "reasoning_content", None)
if not reasoning:
# Future-proofing: the same wire convention under a
# shorter field name.
reasoning = getattr(delta, "reasoning", None)
if reasoning:
yield StreamPiece("thinking", reasoning)
content = delta.content
if content:
yield StreamPiece("content", content)
if (
calls
and not emitted
and getattr(choice, "finish_reason", None) == "tool_calls"
):
for piece in _materialize_tool_calls(calls):
yield piece
emitted = True
if calls and not emitted:
for piece in _materialize_tool_calls(calls):
yield piece
except LLMError:
raise
except Exception as e: # noqa: BLE001 — wrap transport-level failures
raise LLMError(f"chat stream from {self.settings.llm_base_url} failed: {e}") from e
finally:
# Phase 48: deterministic teardown — whenever ``create()``
# succeeded, close the endpoint's stream on every subsequent
# exit (normal exhaustion, wrapped failures, and consumer
# abandon). A failure of ``create()`` itself never sets
# ``stream``, so it stays the plain wrap above.
if stream is not None:
await stream.close()
async def chat_stream_retried(
llm: LLMClient,
messages: list[dict[str, str]],
*,
tools: list[dict[str, Any]] | None = None,
retries: int = 0,
delay: float = 0.0,
) -> AsyncGenerator[StreamPiece | ToolCallPiece | RetryPiece, None]:
"""Stream a chat turn, retrying a dead endpoint (phase 67).
Wraps :meth:`LLMClient.chat_stream` with the retry-before-first-piece
rule (owner-locked A2): a request is restarted **only** while no output
piece (thinking/tool/delta) has been yielded for it. Once pieces have
flowed, an :class:`LLMError` is re-raised unchanged — a partial answer
is never redone, and the API layer's terminal ``error`` frame applies.
This primitive is the ONLY place that rule lives (the chat endpoint
and the agent loop both build on it).
Up to *retries* restarts after the initial attempt (``retries + 1``
attempts total; ``retries=0`` is exactly one attempt with no
:class:`RetryPiece` — the pre-phase-67 kill-switch path). Each restart
is preceded by one :class:`RetryPiece` (``attempt`` = the 1-based
number of the attempt about to be tried — the failed attempt + 1 —
and ``max_attempts`` = ``retries + 1``) and a flat
``asyncio.sleep(delay)`` — the TODO-locked fixed interval, no
backoff. The ``RetryPiece`` always precedes its sleep: the API frame
must reach the client before the wait starts.
The request is restarted byte-identical: ``chat_stream`` is stateless,
so every attempt is opened with the SAME *messages*/*tools*.
Teardown (phase 48, extended): every attempt's stream is explicitly
closed in a ``finally`` — normal exhaustion, a terminal
:class:`LLMError`, and a consumer abandon (``GeneratorExit`` mid-attempt
or during the pre-retry sleep) all run it, so an abandoned turn never
leaves the endpoint's stream open.
"""
max_attempts = retries + 1
for attempt in range(1, max_attempts + 1):
emitted = False
stream = llm.chat_stream(messages, tools=tools)
try:
async for piece in stream:
emitted = True
yield piece
except LLMError as e:
if emitted:
# Locked A2: tokens already flowed — the failure is
# terminal, never redo a partial answer.
raise
if attempt >= max_attempts:
# Retries exhausted — the API layer turns this into the
# terminal error frame.
raise
logger.warning(
"llm stream failed before the first piece (attempt %d/%d) — "
"retrying in %.1fs: %s",
attempt,
max_attempts,
delay,
e,
)
yield RetryPiece(attempt + 1, max_attempts)
await asyncio.sleep(delay)
else:
# A fully consumed attempt — the turn is done; the loop must
# NOT open another stream for an attempt that never failed.
return
finally:
# Phase 48 teardown for THIS attempt's stream (a no-op once the
# stream already ended; the real close on a consumer abandon).
await stream.aclose()
async def check_models(llm: LLMClient) -> None:
"""Verify the models a sync needs (embed + summary) before any
expensive work; raise ModelUnavailableError naming the model.
The probe is deliberately tiny — one short embedding
(``sync model check``) and one 1-token-scale completion (``ping``)
— so a dead endpoint is discovered cheaper than a single git
clone. The sync sanitizer downstream (``app.api.sync
._sanitize_error``) still masks any credentials embedded in the
wrapped error text, so the raw endpoint URL in the original
exception is safe to include.
"""
embed_model = llm.settings.llm_embed_model
try:
await llm.embed_one("sync model check")
except Exception as e: # noqa: BLE001 — wrap EmbeddingError + transport failures
raise ModelUnavailableError(
f"The embedding model ('{embed_model}') is not available — "
f"check the model endpoint and retry. ({e})"
) from e
summary_model = llm.settings.llm_summary_model
try:
await llm.chat([{"role": "user", "content": "ping"}])
except Exception as e: # noqa: BLE001 — wrap LLMError + transport failures
raise ModelUnavailableError(
f"The summary model ('{summary_model}') is not available — "
f"check the model endpoint and retry. ({e})"
) from e