"""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_`` 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_" 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