287 lines
12 KiB
Python
287 lines
12 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): 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
|