"""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 TYPE_CHECKING, Any, Literal, cast from openai import AsyncOpenAI, AsyncStream from openai.types.chat import ChatCompletionChunk, ChatCompletionMessageParam from app.config import Settings, get_settings from app.core.errors import sanitize_error if TYPE_CHECKING: # Phase 71: the filter type is only needed for typing (the module # stays import-graph-clean; callers pass their own instances). from app.rag.scaffolding import ScaffoldingFilter 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, ``ls`` / ``read`` / ``grep``, phase 70), 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 # "ls" | "read" | "grep" (whatever AGENT_TOOLS names, phase 70) arguments: dict[str, Any] @dataclass(frozen=True) class ToolResultPiece: """One executed tool call whose result was truncated (phase 95). A15 extension (owner permission 2026-09-10, ``TODO.md`` L5 — recorded in the phase 95 overview ``00_phase.md``; PLAN.md is being redone by the owner): the ``read`` tool caps its result at ``settings.read_max_chars`` (``BOR_READ_MAX_CHARS``, default 128 000 chars ≈ 32k tokens). When a document is longer than the cap, the agent loop appends the shared ``[…truncated…]`` marker + the grep-pointer notice to the result the model sees AND yields one of these pieces so the API layer can surface the truncation to the user (an SSE ``tool_result`` frame, task 02). It is the ONLY piece the agent loop yields that does not come from the model stream — it is derived from the executed call. ``argument`` is the combined ``source/path`` the model passed (the same value the matching ``tool`` frame carries), ``chars_shown`` is the cap (``settings.read_max_chars``) and ``chars_total`` is the document's true length — so the UI can render "(truncated — showing N of M chars)". ``truncated`` is always ``True`` on a yielded piece (a non-truncated read yields nothing). Frozen like its siblings: an immutable wire value. """ name: str # the tool that was executed (always "read" today) argument: str | None # the model's argument (the combined source/path) truncated: bool # always True on a yielded piece chars_shown: int # the cap actually kept (settings.read_max_chars) chars_total: int # the document's true length @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 an unscoped ``ls``). """ 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=self.settings.llm_timeout, ) 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( # Phase 84 (SEC-13): a base URL configured with # embedded ``user:pass@`` credentials must not reach # the error string — sanitized at construction. f"embeddings request to {sanitize_error(self.settings.llm_base_url)} " f"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, Any]], 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( # Phase 84 (SEC-13): sanitize the base URL (see embed). f"chat completion from {sanitize_error(self.settings.llm_base_url)} " f"failed: {e}" ) from e if not resp.choices: raise LLMError( f"chat completion from {sanitize_error(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 {sanitize_error(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, Any]], tools: list[dict[str, Any]] | None = None, scaffolding: ScaffoldingFilter | None = None, ) -> AsyncGenerator[StreamPiece | ToolCallPiece, None]: """Stream assistant pieces from the chat model (PLAN A5/A15, phase 17). Messages are passed to the request body VERBATIM: string-only ``{role, content}`` dicts are byte-identical on the wire to the pre-phase-74 requests, and an assistant message may additionally carry ``reasoning_content`` (the client's prior thinking, phase 74 — the same wire field the model uses for its OWN reasoning on the response side; the ``openai`` SDK passes message dicts through untouched, so no transport change). ``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 (after the stream's chunks are exhausted — aipi ends the stream at ``finish_reason="tool_calls"``, so this is the wire's emission point). 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``). Scaffolding guardrail (phase 71): the caller may pass a ``ScaffoldingFilter`` — one per request, caller-owned (this method never creates or resets one). When present, only **content** is filtered: ``delta.content`` is fed through the filter and only the clean text is yielded (an empty clean result yields **no** piece — no empty ``delta`` frames); thinking pieces are never filtered (the scratchpad stays raw, phase 17). At stream end the filter's held tail is flushed to a content piece **before** any tool-call materialization (content- before-tools wire convention). ``None`` (the default) keeps today's byte-identical raw path for callers that opt out. 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] = {} 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: if scaffolding is not None: # Phase 71: content only — an empty clean result # yields nothing (no empty delta frames). cleaned = scaffolding.feed(content) if cleaned: yield StreamPiece("content", cleaned) else: yield StreamPiece("content", content) # Phase 71: flush the filter's held tail at stream end, BEFORE # any tool-call materialization — flushed-tail content precedes # ToolCallPieces (content-before-tools wire convention). if scaffolding is not None: tail = scaffolding.flush() if tail: yield StreamPiece("content", tail) if calls: for piece in _materialize_tool_calls(calls): yield piece except LLMError: raise except Exception as e: # noqa: BLE001 — wrap transport-level failures # Phase 84 (SEC-13): sanitize the base URL (see embed). raise LLMError( f"chat stream from {sanitize_error(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, Any]], *, tools: list[dict[str, Any]] | None = None, retries: int = 0, delay: float = 0.0, scaffolding: ScaffoldingFilter | None = None, ) -> 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*. Scaffolding guardrail (phase 71): *scaffolding* is passed through to every attempt's ``chat_stream``. The SAME caller-owned filter object across the retry attempts of one logical request is safe by construction: a restarted attempt only happens while no piece was emitted, i.e. the filter was never fed (its pending buffer is still empty). 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, scaffolding=scaffolding) 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