feat(rag): index markdown KB — chunker, embed client, delta importer, Sources page
Phase 02 (story: import documents):
- fence-aware markdown chunker (heading sections, 200-char overlap,
heading anchor on every chunk, 1200-char hard cap, fence blocks
kept atomic and split under the cap)
- LLMClient over aipi (LiteLLM) reusing the openai client's httpx
transport to send a clean {model, input} payload — the openai SDK
injects encoding_format, which aipi's openai_like group rejects;
token-budget batching + halving retry for the endpoint's
~1024-token per-request input cap
- two-phase per-file upsert importer: sha256 delta (unchanged skip),
atomic commit, A9 exclusion walk, per-source prune, per-file error
tolerance (rollback + log + continue, non-zero CLI exit), adaptive
re-chunk at half target for URL-dense files the endpoint rejects
- scripts/import_docs CLI (repeatable --source, --prune, --limit,
defaults ~/Homelab + ~/Deployments)
- GET /api/docs with per-doc chunk counts; Sources page wired to the
real endpoint (stat cards, full-width a11y table, designed empty
state, DOM-built rows — no innerHTML)
- tests: 63 passed (chunker/llm/importer units, docs API + importer
integration), story E2E 3/3 (real endpoints, in-thread import);
app/ coverage 98%
- real KB imported: 672 docs / 8969 chunks in ~3m, idempotent
re-run (672 unchanged, 0 batches)
- harness: .agent/validate.sh now gates through uv (pytest +
coverage >90% + ruff + pyright) instead of system python3
This commit is contained in:
@@ -0,0 +1 @@
|
||||
"""RAG pipeline package (chunker, LLM client, importer, retrieval)."""
|
||||
@@ -0,0 +1,205 @@
|
||||
"""Markdown-aware chunker (PLAN §5 chunking policy).
|
||||
|
||||
Pure functions, no I/O — fully unit-testable.
|
||||
|
||||
Policy
|
||||
------
|
||||
* **Sections** are split on ATX headings of level ≥ 2 (``## ``/``### ``/…).
|
||||
* A section that fits in ``target_chars`` becomes a single chunk.
|
||||
* A longer section is sub-split at paragraph boundaries (blank lines outside
|
||||
code fences); each chunk after the first starts with the trailing
|
||||
``overlap_chars`` of the previous chunk so context survives the cut.
|
||||
* Every chunk keeps its nearest preceding heading line (the section anchor),
|
||||
so a retrieval hit is always readable in context.
|
||||
* **Code fences** (``` / ~~~) are atomic: a chunk boundary never falls
|
||||
inside one, and lines inside a fence are never mistaken for headings or
|
||||
paragraph breaks. One exception: a fence *larger than* :data:`HARD_MAX_CHARS`
|
||||
is split by line, because aipi's local embedding model rejects requests
|
||||
over ~1024 input tokens and a single 5000-char code block would blow
|
||||
past that on its own.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
#: ATX heading of level ≥ 2 — the section splitter (PLAN §5).
|
||||
_HEADING_RE = re.compile(r"^(#{2,6})\s+\S")
|
||||
#: First H1 — the document title (used by :func:`extract_title`).
|
||||
_H1_RE = re.compile(r"^#\s+\S")
|
||||
#: Opening/closing code fence (three or more backticks or tildes).
|
||||
_FENCE_RE = re.compile(r"^(`{3,}|~{3,})")
|
||||
#: Absolute per-chunk size cap (chars). Measured worst case in the real KB:
|
||||
#: punctuation-heavy machine output (``lspci`` dumps in fences) tokenizes at
|
||||
#: ~1.4 chars/token, so 1200 chars stays under the endpoint's ~1024-token
|
||||
#: per-request input cap even as a solo-chunk request.
|
||||
HARD_MAX_CHARS = 1200
|
||||
|
||||
|
||||
def extract_title(content: str, fallback: str = "") -> str:
|
||||
"""First markdown H1, else *fallback* (the importer passes the file stem)."""
|
||||
for line in content.splitlines():
|
||||
m = _H1_RE.match(line)
|
||||
if m:
|
||||
return line.lstrip("#").strip() or fallback
|
||||
return fallback
|
||||
|
||||
|
||||
def _iter_fence_state(lines: Sequence[str]) -> list[bool]:
|
||||
"""Per-line flags: ``True`` on a fence delimiter or inside a fence."""
|
||||
flags: list[bool] = []
|
||||
marker: str | None = None
|
||||
for line in lines:
|
||||
stripped = line.lstrip()
|
||||
if marker is None:
|
||||
m = _FENCE_RE.match(stripped)
|
||||
if m:
|
||||
marker = m.group(1)[:3]
|
||||
flags.append(True)
|
||||
else:
|
||||
flags.append(False)
|
||||
else:
|
||||
flags.append(True)
|
||||
if stripped.startswith(marker):
|
||||
marker = None
|
||||
return flags
|
||||
|
||||
|
||||
def _section_ranges(lines: Sequence[str], flags: Sequence[bool]) -> list[tuple[int, int]]:
|
||||
"""(start, end) line-index ranges of every heading-delimited section."""
|
||||
starts = [0]
|
||||
for i, line in enumerate(lines):
|
||||
if i > 0 and not flags[i] and _HEADING_RE.match(line):
|
||||
starts.append(i)
|
||||
return list(zip(starts, [*starts[1:], len(lines)], strict=True))
|
||||
|
||||
|
||||
def _paragraph_blocks(lines: Sequence[str], flags: Sequence[bool]) -> list[str]:
|
||||
"""Blocks of lines separated by blank lines *outside* fences.
|
||||
|
||||
A fenced code block (possibly containing blank lines) is one block.
|
||||
"""
|
||||
blocks: list[str] = []
|
||||
cur: list[str] = []
|
||||
for line, in_fence in zip(lines, flags, strict=True):
|
||||
if in_fence or line.strip():
|
||||
cur.append(line)
|
||||
elif cur:
|
||||
blocks.append("\n".join(cur))
|
||||
cur = []
|
||||
if cur:
|
||||
blocks.append("\n".join(cur))
|
||||
return blocks
|
||||
|
||||
|
||||
def _split_oversized(block: str, hard: int = HARD_MAX_CHARS) -> list[str]:
|
||||
"""Split one oversized block (e.g. a huge fenced code block) into line
|
||||
groups under *hard* chars. A single line longer than *hard* is chopped
|
||||
at char boundaries — better than exceeding the endpoint's token cap."""
|
||||
parts: list[str] = []
|
||||
cur: list[str] = []
|
||||
cur_len = 0
|
||||
for line in block.splitlines():
|
||||
if len(line) > hard:
|
||||
if cur:
|
||||
parts.append("\n".join(cur))
|
||||
cur, cur_len = [], 0
|
||||
parts.extend(line[i : i + hard] for i in range(0, len(line), hard))
|
||||
continue
|
||||
if cur and cur_len + 1 + len(line) > hard:
|
||||
parts.append("\n".join(cur))
|
||||
cur, cur_len = [], 0
|
||||
cur.append(line)
|
||||
cur_len += 1 + len(line)
|
||||
if cur:
|
||||
parts.append("\n".join(cur))
|
||||
return parts
|
||||
|
||||
|
||||
def _pack_blocks(blocks: Sequence[str], target: int, overlap: int) -> list[str]:
|
||||
"""Greedy paragraph packing; consecutive chunks share ``overlap`` chars."""
|
||||
# Oversized blocks are split to the *budget*, so a solo-block chunk
|
||||
# (anchor + block) cannot exceed the target — let alone the hard cap.
|
||||
split_at = min(HARD_MAX_CHARS, max(1, target))
|
||||
expanded: list[str] = []
|
||||
for block in blocks:
|
||||
if len(block) > split_at:
|
||||
expanded.extend(_split_oversized(block, split_at))
|
||||
else:
|
||||
expanded.append(block)
|
||||
chunks: list[str] = []
|
||||
cur_parts: list[str] = []
|
||||
cur_len = 0
|
||||
for block in expanded:
|
||||
if cur_len == 0:
|
||||
cur_parts.append(block)
|
||||
cur_len = len(block)
|
||||
elif cur_len + 1 + len(block) <= target:
|
||||
cur_parts.append(block)
|
||||
cur_len += 1 + len(block)
|
||||
else:
|
||||
chunks.append("\n".join(cur_parts))
|
||||
# The tail is only carried over when tail + block still fits the
|
||||
# budget — otherwise the chunk would exceed the size cap.
|
||||
tail = (
|
||||
chunks[-1][-overlap:]
|
||||
if (overlap > 0 and overlap + 1 + len(block) <= target)
|
||||
else ""
|
||||
)
|
||||
cur_parts = [tail, block] if tail else [block]
|
||||
cur_len = len(tail) + 1 + len(block) if tail else len(block)
|
||||
if cur_parts:
|
||||
chunks.append("\n".join(cur_parts))
|
||||
return chunks
|
||||
|
||||
|
||||
def _chunk_section(
|
||||
lines: Sequence[str], flags: Sequence[bool], target: int, overlap: int
|
||||
) -> list[str]:
|
||||
"""Chunk one section (its heading line included, when it has one)."""
|
||||
first = next((i for i, line in enumerate(lines) if line.strip()), None)
|
||||
if first is None:
|
||||
return []
|
||||
anchor: str | None = None
|
||||
if not flags[first] and _HEADING_RE.match(lines[first]):
|
||||
anchor = lines[first]
|
||||
body: Sequence[str] = lines[first + 1 :]
|
||||
body_flags: Sequence[bool] = flags[first + 1 :]
|
||||
else:
|
||||
body = lines
|
||||
body_flags = flags
|
||||
|
||||
full = (anchor + "\n" + "\n".join(body) if anchor is not None else "\n".join(body)).strip()
|
||||
if not full:
|
||||
return []
|
||||
if len(full) <= target:
|
||||
return [full]
|
||||
|
||||
blocks = _paragraph_blocks(body, body_flags)
|
||||
if not blocks:
|
||||
return [full]
|
||||
budget = target - (len(anchor) + 2 if anchor is not None else 0)
|
||||
packed = _pack_blocks(blocks, max(1, budget), overlap)
|
||||
if anchor is None:
|
||||
return packed
|
||||
return [(anchor + "\n\n" + part).strip() for part in packed]
|
||||
|
||||
|
||||
def chunk_markdown(
|
||||
content: str, target_chars: int = 2000, overlap_chars: int = 200
|
||||
) -> list[str]:
|
||||
"""Split markdown into retrieval chunks (see module docstring for policy)."""
|
||||
if target_chars <= 0:
|
||||
raise ValueError("target_chars must be > 0")
|
||||
if overlap_chars < 0:
|
||||
raise ValueError("overlap_chars must be >= 0")
|
||||
# The endpoint's token cap is absolute — a larger target is unsafe.
|
||||
target = min(target_chars, HARD_MAX_CHARS)
|
||||
overlap = min(overlap_chars, target - 1)
|
||||
|
||||
lines = content.splitlines()
|
||||
flags = _iter_fence_state(lines)
|
||||
chunks: list[str] = []
|
||||
for start, end in _section_ranges(lines, flags):
|
||||
chunks.extend(_chunk_section(lines[start:end], flags[start:end], target, overlap))
|
||||
return [c for c in chunks if c.strip()]
|
||||
@@ -0,0 +1,247 @@
|
||||
"""Knowledge-base importer (PLAN §5 / §9 / §11).
|
||||
|
||||
Walks ``*.md`` files (A9 exclusion list), diffs by sha256 against
|
||||
``documents.content_hash`` and, for every new or changed file, runs the
|
||||
two-phase upsert:
|
||||
|
||||
1. upsert the document row and replace its chunk rows (embeddings NULL)
|
||||
2. embed the new chunks in batches and attach the vectors
|
||||
3. commit — one transaction per file, so a failed embedding leaves the
|
||||
database untouched and the file is simply retried on the next run
|
||||
|
||||
``prune=True`` deletes documents (of the imported sources only) whose files
|
||||
no longer exist. Per-file logging uses the verbs
|
||||
``added | updated | unchanged | pruned`` plus a summary line (PLAN §9).
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, datetime
|
||||
from pathlib import Path
|
||||
from typing import Protocol
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from app.config import Settings
|
||||
from app.db import SessionLocal
|
||||
from app.models import Chunk, Document
|
||||
from app.rag.chunker import chunk_markdown, extract_title
|
||||
from app.rag.llm import EmbeddingError
|
||||
|
||||
logger = logging.getLogger("app.importer")
|
||||
|
||||
#: Non-content directories never imported (PLAN anchor A9).
|
||||
EXCLUDED_DIRS: frozenset[str] = frozenset(
|
||||
{".venv", "node_modules", ".git", "__pycache__", ".pytest_cache", "dist", "build"}
|
||||
)
|
||||
|
||||
|
||||
class Embedder(Protocol):
|
||||
"""Everything the importer needs from the LLM client (duck-typed for tests)."""
|
||||
|
||||
settings: Settings
|
||||
embed_batches: int
|
||||
|
||||
async def embed(self, texts: list[str]) -> list[list[float]]: ...
|
||||
|
||||
|
||||
@dataclass
|
||||
class ImportSummary:
|
||||
"""Counts for one import run (also printed by the CLI)."""
|
||||
|
||||
files: int = 0
|
||||
added: int = 0
|
||||
updated: int = 0
|
||||
unchanged: int = 0
|
||||
pruned: int = 0
|
||||
errors: int = 0
|
||||
chunks: int = 0
|
||||
embed_batches: int = 0
|
||||
|
||||
def log(self) -> None:
|
||||
logger.info(
|
||||
"import: summary files=%d added=%d updated=%d unchanged=%d pruned=%d "
|
||||
"errors=%d chunks=%d embed_batches=%d",
|
||||
self.files,
|
||||
self.added,
|
||||
self.updated,
|
||||
self.unchanged,
|
||||
self.pruned,
|
||||
self.errors,
|
||||
self.chunks,
|
||||
self.embed_batches,
|
||||
)
|
||||
|
||||
|
||||
def iter_markdown_files(root: Path, excluded: frozenset[str] = EXCLUDED_DIRS) -> list[Path]:
|
||||
"""All ``*.md`` files under *root* (sorted), skipping excluded dirs (A9)."""
|
||||
if not root.is_dir():
|
||||
return []
|
||||
files: list[Path] = []
|
||||
for path in sorted(root.rglob("*.md")):
|
||||
rel = path.relative_to(root)
|
||||
if any(part in excluded for part in rel.parts[:-1]):
|
||||
continue
|
||||
files.append(path)
|
||||
return files
|
||||
|
||||
|
||||
async def import_sources(
|
||||
sources: list[Path],
|
||||
llm: Embedder,
|
||||
*,
|
||||
prune: bool = False,
|
||||
limit: int | None = None,
|
||||
session: Session | None = None,
|
||||
) -> ImportSummary:
|
||||
"""Import every ``*.md`` under *sources* (see module docstring).
|
||||
|
||||
``session`` may be supplied (tests); a private one is opened and closed
|
||||
otherwise. ``limit`` caps the number of files processed (debug only) and
|
||||
disables pruning, since an incomplete walk must not drive deletions.
|
||||
"""
|
||||
if limit is not None and limit <= 0:
|
||||
raise ValueError("limit must be >= 1")
|
||||
summary = ImportSummary()
|
||||
owns_session = session is None
|
||||
if session is None:
|
||||
session = SessionLocal()
|
||||
seen: set[tuple[str, str]] = set()
|
||||
source_names: set[str] = set()
|
||||
try:
|
||||
for root in sources:
|
||||
if not root.is_dir():
|
||||
logger.warning("import: source dir not found, skipping: %s", root)
|
||||
continue
|
||||
if limit is not None and summary.files >= limit:
|
||||
break
|
||||
source = root.name
|
||||
source_names.add(source)
|
||||
for path in iter_markdown_files(root):
|
||||
if limit is not None and summary.files >= limit:
|
||||
break
|
||||
rel = path.relative_to(root).as_posix()
|
||||
seen.add((source, rel))
|
||||
summary.files += 1
|
||||
try:
|
||||
await _index_file(
|
||||
session, source=source, rel=rel, full_path=path, llm=llm,
|
||||
summary=summary,
|
||||
)
|
||||
except EmbeddingError as e:
|
||||
# A pathological file (e.g. content the embedding endpoint
|
||||
# refuses) must not abort the whole KB: roll back its
|
||||
# uncommitted rows, log loudly, and keep going. The next
|
||||
# run retries it.
|
||||
session.rollback()
|
||||
summary.errors += 1
|
||||
logger.error("import: error source=%s path=%s — %s", source, rel, e)
|
||||
if prune:
|
||||
if limit is not None:
|
||||
logger.warning("import: --prune ignored because --limit was given")
|
||||
else:
|
||||
summary.pruned = _prune(session, source_names, seen)
|
||||
summary.embed_batches = llm.embed_batches
|
||||
summary.log()
|
||||
return summary
|
||||
finally:
|
||||
if owns_session:
|
||||
session.close()
|
||||
|
||||
|
||||
async def _index_file(
|
||||
session: Session,
|
||||
*,
|
||||
source: str,
|
||||
rel: str,
|
||||
full_path: Path,
|
||||
llm: Embedder,
|
||||
summary: ImportSummary,
|
||||
) -> None:
|
||||
"""Upsert one file: doc row + chunk rows + embeddings, one transaction."""
|
||||
settings = llm.settings
|
||||
content = full_path.read_text(encoding="utf-8", errors="replace")
|
||||
digest = hashlib.sha256(content.encode("utf-8")).hexdigest()
|
||||
doc = session.scalar(select(Document).where(Document.source == source, Document.path == rel))
|
||||
if doc is not None and doc.content_hash == digest:
|
||||
summary.unchanged += 1
|
||||
logger.info("import: unchanged source=%s path=%s", source, rel)
|
||||
return
|
||||
|
||||
verb = "updated" if doc is not None else "added"
|
||||
title = extract_title(content, fallback=full_path.stem)
|
||||
if doc is None:
|
||||
doc = Document(
|
||||
source=source,
|
||||
path=rel,
|
||||
full_path=str(full_path),
|
||||
title=title,
|
||||
content=content,
|
||||
content_hash=digest,
|
||||
indexed_at=datetime.now(UTC),
|
||||
)
|
||||
session.add(doc)
|
||||
else:
|
||||
doc.full_path = str(full_path)
|
||||
doc.title = title
|
||||
doc.content = content
|
||||
doc.content_hash = digest
|
||||
doc.indexed_at = datetime.now(UTC)
|
||||
|
||||
session.flush() # guarantees doc.id even for brand-new rows
|
||||
|
||||
# Phase 1+2 — chunk, replace the chunk rows (embeddings NULL), embed,
|
||||
# and commit the whole file atomically. If the endpoint rejects a chunk
|
||||
# as over its input token cap (URL-dense paragraphs tokenize at ~1.1
|
||||
# chars/token), halve this file's chunk target and retry — the global
|
||||
# policy stays intact for the rest of the KB.
|
||||
target = max(400, settings.chunk_target_chars)
|
||||
while True:
|
||||
chunks_text = chunk_markdown(content, target, settings.chunk_overlap_chars)
|
||||
doc.chunks = [
|
||||
Chunk(document_id=doc.id, position=i, content=c) for i, c in enumerate(chunks_text)
|
||||
]
|
||||
session.flush() # delete-orphan cascade drops the previous rows
|
||||
if not doc.chunks:
|
||||
break
|
||||
try:
|
||||
vectors = await llm.embed([c.content for c in doc.chunks])
|
||||
for row, vec in zip(doc.chunks, vectors, strict=True):
|
||||
row.embedding = vec
|
||||
break
|
||||
except EmbeddingError as e:
|
||||
if "token cap" not in str(e) or target <= 400:
|
||||
raise
|
||||
logger.info(
|
||||
"import: re-chunking at %d chars after endpoint token cap: %s",
|
||||
target // 2,
|
||||
rel,
|
||||
)
|
||||
target //= 2
|
||||
session.commit()
|
||||
|
||||
if verb == "added":
|
||||
summary.added += 1
|
||||
else:
|
||||
summary.updated += 1
|
||||
summary.chunks += len(chunks_text)
|
||||
logger.info("import: %s source=%s path=%s chunks=%d", verb, source, rel, len(chunks_text))
|
||||
|
||||
|
||||
def _prune(session: Session, source_names: set[str], seen: set[tuple[str, str]]) -> int:
|
||||
"""Delete documents of *source_names* whose file is no longer in *seen*."""
|
||||
if not source_names:
|
||||
return 0
|
||||
pruned = 0
|
||||
docs = session.scalars(select(Document).where(Document.source.in_(source_names))).all()
|
||||
for doc in docs:
|
||||
if (doc.source, doc.path) not in seen:
|
||||
session.delete(doc)
|
||||
pruned += 1
|
||||
logger.info("import: pruned source=%s path=%s", doc.source, doc.path)
|
||||
if pruned:
|
||||
session.commit()
|
||||
return pruned
|
||||
+162
@@ -0,0 +1,162 @@
|
||||
"""Async OpenAI-compatible client for the self-hosted aipi endpoint (PLAN A5).
|
||||
|
||||
Phase 02 adds the embeddings surface (the importer — and, from phase 03,
|
||||
retrieval — need it). Chat streaming lands in phase 03 on this same client.
|
||||
|
||||
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 openai import AsyncOpenAI
|
||||
|
||||
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."""
|
||||
|
||||
|
||||
# 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, phase 03)."""
|
||||
(vec,) = await self.embed([text])
|
||||
return vec
|
||||
Reference in New Issue
Block a user