"""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): 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 logging from collections.abc import AsyncIterator from dataclasses import dataclass from typing import Literal, cast from openai import AsyncOpenAI from openai.types.chat import 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).""" @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 # 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]] ) -> AsyncIterator[StreamPiece]: """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). 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. """ try: # ``{role, content}`` dicts are exactly what the message params # accept; the cast keeps pyright honest about the SDK's union. stream = await self._client.chat.completions.create( model=self.settings.llm_chat_model, messages=cast("list[ChatCompletionMessageParam]", messages), temperature=0.4, max_tokens=self.settings.max_output_tokens, stream=True, ) async for chunk in stream: if not chunk.choices: continue delta = chunk.choices[0].delta 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) 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