feat(ingestion): add bounded, benchmark-aligned embedding execution
Why: - Plan 001 Phase 4 needs batched, concurrency-bounded embedding wired into the inline upload path, with process-wide capacity/timeout/chunk-limit guards (ADR-0017). - The BM25 analyzer and dense-model config are ported from the `emet` evaluation lab, which benchmarked them against the real Farsi corpus (bm25-fa-norm-stop; nomic-embed-text-v2-moe at 768-dim; text-embedding-3-large at native 3072-dim), closing open items in ADR-0001/ADR-0005. Changes: - New: embedding ports, orchestration (embed_chunks), request-bounds helpers, and dense/sparse adapters (analyzers.py, bm25.py, openai_compatible.py). - upload.py now parses/chunks/embeds inline behind INGESTION_MAX_CONCURRENCY (503), INGESTION_TIMEOUT_SECONDS (504), and the chunk-count ceiling (413); every failure path still writes a terminal job row. - Lifespan builds and warms both dense embedders at startup (fail-soft) and creates the sparse embedder and concurrency semaphore. - httpx moves from dev to main dependencies (adapters use it directly). Impact: - Qdrant point upserts are still Phase 5 -- chunks_indexed stays 0. - New EMBEDDING_* env vars documented in .env.example; safe defaults. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
0
src/infrastructure/embedding/__init__.py
Normal file
0
src/infrastructure/embedding/__init__.py
Normal file
150
src/infrastructure/embedding/analyzers.py
Normal file
150
src/infrastructure/embedding/analyzers.py
Normal file
@@ -0,0 +1,150 @@
|
||||
"""The `fa_norm_stop` BM25 analyzer (ADR-0001, ADR-0005).
|
||||
|
||||
Ported from the `emet` evaluation lab
|
||||
(`src/chatbot_gh/adapters/sparse/analyzers.py`), which benchmarked four Farsi
|
||||
analyzer variants on the real corpus and found `fa_norm_stop` the best
|
||||
performer. This is a *measured* artifact: changing the normalization,
|
||||
tokenization, or stopword list invalidates that result, so improvements belong
|
||||
in a new emet benchmark run rather than in an edit here.
|
||||
|
||||
Deliberately independent of `src/application/ingestion/normalization.py`.
|
||||
Those solve different problems: `normalize_persian_text` shapes chunk content
|
||||
that gets cited back to the reader, so ADR-0018 has it preserve digits and
|
||||
punctuation as authored. This module shapes index terms nobody ever sees, so
|
||||
it folds digits and diacritics freely. Sharing one function between them would
|
||||
let a display-motivated tweak silently perturb the benchmarked sparse index.
|
||||
"""
|
||||
|
||||
import re
|
||||
import unicodedata
|
||||
|
||||
# Several Arabic letterforms are visually indistinguishable from Latin ones in
|
||||
# a monospace editor (alef from "l", heh from "o"), and literals render
|
||||
# right-to-left, visually reordering the source line. `normalization.py` writes
|
||||
# them as codepoints for that reason; this module follows the same convention.
|
||||
_ZWNJ = 0x200C
|
||||
_ARABIC_YEH = 0x064A
|
||||
_ARABIC_KAF = 0x0643
|
||||
_TEH_MARBUTA = 0x0629
|
||||
_HAMZA_ON_WAW = 0x0624
|
||||
_ALEF_HAMZA_BELOW = 0x0625
|
||||
_ALEF_HAMZA_ABOVE = 0x0623
|
||||
|
||||
_PERSIAN_YEH = 0x06CC
|
||||
_PERSIAN_KEHEH = 0x06A9
|
||||
_HEH = 0x0647
|
||||
_WAW = 0x0648
|
||||
_ALEF = 0x0627
|
||||
_SPACE = 0x0020
|
||||
|
||||
# Persian (U+06F0-U+06F9) and Arabic-Indic (U+0660-U+0669) digits both fold to
|
||||
# ASCII, so the same number matches however it was authored.
|
||||
_EASTERN_DIGITS = str.maketrans("۰۱۲۳۴۵۶۷۸۹٠١٢٣٤٥٦٧٨٩", "01234567890123456789")
|
||||
|
||||
# ZWNJ becomes a space (splitting compounds into separate terms) and the
|
||||
# Arabic letterforms fold to their Persian equivalents. Every entry is a
|
||||
# single codepoint mapping to a single codepoint over disjoint sources, so
|
||||
# applying them in one pass is equivalent to emet's chained `str.replace`
|
||||
# calls -- provided NFC runs first, since NFC is what composes the hamza
|
||||
# forms this table then folds.
|
||||
_FOLDING: dict[int, int] = {
|
||||
_ZWNJ: _SPACE,
|
||||
_ARABIC_YEH: _PERSIAN_YEH,
|
||||
_ARABIC_KAF: _PERSIAN_KEHEH,
|
||||
_TEH_MARBUTA: _HEH,
|
||||
_HAMZA_ON_WAW: _WAW,
|
||||
_ALEF_HAMZA_BELOW: _ALEF,
|
||||
_ALEF_HAMZA_ABOVE: _ALEF,
|
||||
}
|
||||
_FOLDING.update(_EASTERN_DIGITS)
|
||||
|
||||
# Common Persian/Arabic stopwords (function words + FAQ noise), plus the
|
||||
# English function words that appear in a mixed-script corpus. Kept small and
|
||||
# explicit -- deliberately not a full hazm list. `_HEH_ALEF` is the plural
|
||||
# suffix "ha"; written as a codepoint pair because both of its letters are
|
||||
# Latin-confusable, which is exactly the case Ruff's RUF001 flags.
|
||||
_HEH_ALEF = chr(_HEH) + chr(_ALEF)
|
||||
|
||||
_PERSIAN_STOPWORDS: frozenset[str] = frozenset(
|
||||
{
|
||||
"و",
|
||||
"در",
|
||||
"به",
|
||||
"از",
|
||||
"که",
|
||||
"این",
|
||||
"را",
|
||||
"با",
|
||||
"برای",
|
||||
"آن",
|
||||
"یک",
|
||||
"است",
|
||||
"شد",
|
||||
"شده",
|
||||
"می",
|
||||
"های",
|
||||
_HEH_ALEF,
|
||||
"یا",
|
||||
"تا",
|
||||
"بر",
|
||||
"اگر",
|
||||
"هم",
|
||||
"نیز",
|
||||
"ولی",
|
||||
"اما",
|
||||
"چه",
|
||||
"چون",
|
||||
"روی",
|
||||
"پس",
|
||||
"پیش",
|
||||
"هر",
|
||||
"هیچ",
|
||||
"بود",
|
||||
"باشد",
|
||||
"هست",
|
||||
"نیست",
|
||||
"کند",
|
||||
"کرد",
|
||||
"کردن",
|
||||
"شود",
|
||||
"the",
|
||||
"a",
|
||||
"an",
|
||||
"of",
|
||||
"to",
|
||||
"and",
|
||||
"in",
|
||||
"on",
|
||||
"for",
|
||||
"is",
|
||||
"are",
|
||||
}
|
||||
)
|
||||
|
||||
# Word characters minus underscore. Note this KEEPS digits: an insurance
|
||||
# corpus is full of policy numbers, dates, and amounts, and those are exactly
|
||||
# the tokens a lexical index should be able to match on.
|
||||
_TOKEN_RE = re.compile(r"[^\W_]+", re.UNICODE)
|
||||
|
||||
FA_NORM_STOP = "fa_norm_stop"
|
||||
|
||||
|
||||
def _normalize_fa(text: str) -> str:
|
||||
return unicodedata.normalize("NFC", text).translate(_FOLDING)
|
||||
|
||||
|
||||
def _tokenize_raw(text: str) -> list[str]:
|
||||
return [m.group(0).lower() for m in _TOKEN_RE.finditer(text)]
|
||||
|
||||
|
||||
def analyze(text: str, analyzer: str = FA_NORM_STOP) -> list[str]:
|
||||
"""Tokenize `text` into sparse-index terms.
|
||||
|
||||
Only `fa_norm_stop` is implemented -- emet's other three variants
|
||||
(`raw`, `fa_norm`, `fa_norm_stem`) lost the benchmark and exist there as
|
||||
experiment arms, not as configurations this service should run.
|
||||
"""
|
||||
if analyzer != FA_NORM_STOP:
|
||||
raise ValueError(f"Unknown analyzer '{analyzer}'")
|
||||
tokens = _tokenize_raw(_normalize_fa(text))
|
||||
return [token for token in tokens if token not in _PERSIAN_STOPWORDS]
|
||||
92
src/infrastructure/embedding/bm25.py
Normal file
92
src/infrastructure/embedding/bm25.py
Normal file
@@ -0,0 +1,92 @@
|
||||
"""The `bm25-fa-norm-stop` sparse embedder (ADR-0001, ADR-0005).
|
||||
|
||||
Ported from the `emet` evaluation lab
|
||||
(`src/chatbot_gh/adapters/sparse/bm25_embedder.py`), the configuration that
|
||||
won its Farsi analyzer benchmark. Not Qdrant's hosted `Qdrant/bm25` FastEmbed
|
||||
model, whose documented language support omits Farsi (ADR-0005).
|
||||
|
||||
**The BM25 work is split across two systems.** This adapter applies the
|
||||
term-frequency saturation half client-side -- the `k` and `b` parameters,
|
||||
including document-length normalization. IDF is *not* computed here: Qdrant
|
||||
supplies it from collection-wide statistics when the sparse vector field is
|
||||
created with `modifier="idf"`.
|
||||
|
||||
That split is load-bearing. A Qdrant collection created without
|
||||
`modifier="idf"` will silently score these vectors as saturated term
|
||||
frequencies with no IDF weighting at all -- no error, just materially worse
|
||||
lexical retrieval. The collection bootstrap (plan 001 Phase 5) must set it.
|
||||
"""
|
||||
|
||||
from collections import Counter
|
||||
from collections.abc import Sequence
|
||||
from hashlib import blake2b
|
||||
|
||||
from src.application.ingestion.models import SparseVector
|
||||
from src.config import SparseEmbeddingSettings
|
||||
from src.infrastructure.embedding.analyzers import analyze
|
||||
|
||||
# Qdrant sparse indices must be non-negative and fit a signed 32-bit int.
|
||||
_INDEX_SPACE = 2**31 - 1
|
||||
|
||||
|
||||
def _token_index(token: str) -> int:
|
||||
"""Map a term to its sparse-vector index.
|
||||
|
||||
A hash rather than a vocabulary table, so the mapping needs no shared
|
||||
state and stays identical across processes, restarts, and — critically —
|
||||
between ingest-time and query-time encoding. `blake2b` rather than
|
||||
`hash()`, which is PYTHONHASHSEED-salted and therefore differs per
|
||||
process.
|
||||
"""
|
||||
digest = blake2b(token.encode("utf-8"), digest_size=8).digest()
|
||||
return int.from_bytes(digest, "big") % _INDEX_SPACE
|
||||
|
||||
|
||||
def text_to_sparse_vector(
|
||||
text: str, *, settings: SparseEmbeddingSettings, query: bool = False
|
||||
) -> SparseVector:
|
||||
"""Encode one text as a BM25-saturated sparse vector (IDF applied by Qdrant).
|
||||
|
||||
Document and query sides differ in exactly one term: documents carry the
|
||||
`b` length normalization, queries do not (standard BM25 practice — a
|
||||
query's own length should not discount its terms).
|
||||
"""
|
||||
tokens = analyze(text, settings.analyzer)
|
||||
if not tokens:
|
||||
return SparseVector(indices=[], values=[])
|
||||
|
||||
frequencies = Counter(tokens)
|
||||
doc_length = float(len(tokens))
|
||||
k = settings.k
|
||||
b = settings.b
|
||||
|
||||
indices: list[int] = []
|
||||
values: list[float] = []
|
||||
# Sorted so the emitted vector is deterministic for a given text, which
|
||||
# keeps re-ingestion byte-stable and makes the output testable.
|
||||
for token, freq in sorted(frequencies.items()):
|
||||
if query:
|
||||
weight = freq * (k + 1.0) / (freq + k)
|
||||
else:
|
||||
weight = freq * (k + 1.0) / (freq + k * (1.0 - b + b * doc_length / settings.avg_len))
|
||||
indices.append(_token_index(token))
|
||||
values.append(float(weight))
|
||||
|
||||
return SparseVector(indices=indices, values=values)
|
||||
|
||||
|
||||
class Bm25SparseEmbedder:
|
||||
"""A `SparseEmbedder` (see `src/application/ports/embedding.py`).
|
||||
|
||||
Pure CPU work with no network calls, so it is blocking: callers offload it
|
||||
via `anyio.to_thread.run_sync` with the ingestion `CapacityLimiter`
|
||||
(ADR-0017), never awaiting it directly on the event loop.
|
||||
"""
|
||||
|
||||
name = "sparse"
|
||||
|
||||
def __init__(self, settings: SparseEmbeddingSettings) -> None:
|
||||
self._settings = settings
|
||||
|
||||
def embed_batch(self, texts: Sequence[str], *, query: bool = False) -> list[SparseVector]:
|
||||
return [text_to_sparse_vector(text, settings=self._settings, query=query) for text in texts]
|
||||
81
src/infrastructure/embedding/openai_compatible.py
Normal file
81
src/infrastructure/embedding/openai_compatible.py
Normal file
@@ -0,0 +1,81 @@
|
||||
"""Dense embedding adapter for OpenAI-compatible `/embeddings` endpoints.
|
||||
|
||||
Backs both `dense_nomic` (self-hosted `nomic-embed-text-v2-moe` behind
|
||||
Ollama's OpenAI-compatible shim) and `dense_openai` (OpenAI's hosted API) —
|
||||
both speak the same request/response shape, so one adapter serves both named
|
||||
vectors with different config (ADR-0001).
|
||||
|
||||
Uses `httpx` directly rather than the `openai` SDK. The `emet` benchmark this
|
||||
configuration comes from uses the SDK, but it is a synchronous batch tool;
|
||||
ADR-0017 requires async, semaphore-bounded batches here, and the request shape
|
||||
is small enough that the SDK earns nothing.
|
||||
|
||||
`httpx.AsyncClient` is application-lifetime (ADR-0012): built once in the
|
||||
FastAPI lifespan and passed in, never constructed per call.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
# Ollama's default port, plus the alternate the benchmarked deployment uses.
|
||||
_OLLAMA_PORTS = frozenset({11434, 11435})
|
||||
|
||||
|
||||
def is_ollama_base_url(base_url: str) -> bool:
|
||||
"""Whether `base_url` looks like an Ollama OpenAI-compatible endpoint.
|
||||
|
||||
Ollama unloads an idle model, and reloading `nomic-embed-text-v2-moe`
|
||||
costs well over two minutes — longer than `INGESTION_TIMEOUT_SECONDS`, so
|
||||
a cold upload would 504. `keep_alive` is how the model is kept resident,
|
||||
and it is an Ollama extension, hence the sniffing.
|
||||
"""
|
||||
parsed = urlparse(base_url)
|
||||
host = parsed.hostname or ""
|
||||
return parsed.port in _OLLAMA_PORTS or "ollama" in host.lower()
|
||||
|
||||
|
||||
class OpenAICompatibleEmbedder:
|
||||
"""A `DenseEmbedder` (see `src/application/ports/embedding.py`) over one
|
||||
OpenAI-compatible `/embeddings` endpoint.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: httpx.AsyncClient,
|
||||
*,
|
||||
name: str,
|
||||
model: str,
|
||||
dimensions: int | None = None,
|
||||
document_prefix: str = "",
|
||||
keep_alive: str | None = None,
|
||||
) -> None:
|
||||
self.name = name
|
||||
self._client = client
|
||||
self._model = model
|
||||
self._dimensions = dimensions
|
||||
self._document_prefix = document_prefix
|
||||
self._keep_alive = keep_alive
|
||||
|
||||
async def embed_batch(self, texts: Sequence[str]) -> list[list[float]]:
|
||||
inputs = (
|
||||
[f"{self._document_prefix}{text}" for text in texts]
|
||||
if self._document_prefix
|
||||
else list(texts)
|
||||
)
|
||||
payload: dict[str, object] = {"model": self._model, "input": inputs}
|
||||
if self._dimensions is not None:
|
||||
payload["dimensions"] = self._dimensions
|
||||
if self._keep_alive is not None:
|
||||
payload["keep_alive"] = self._keep_alive
|
||||
|
||||
response = await self._client.post("/embeddings", json=payload)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
|
||||
# Sort by `index` rather than trusting response order: the contract
|
||||
# guarantees the field, not the ordering, and a silently permuted
|
||||
# batch would attach every vector to the wrong chunk.
|
||||
data = sorted(body["data"], key=lambda item: item["index"])
|
||||
return [item["embedding"] for item in data]
|
||||
Reference in New Issue
Block a user