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:
170
tests/unit/infrastructure/embedding/test_bm25.py
Normal file
170
tests/unit/infrastructure/embedding/test_bm25.py
Normal file
@@ -0,0 +1,170 @@
|
||||
"""The `bm25-fa-norm-stop` sparse embedder (ADR-0001, ADR-0005).
|
||||
|
||||
These assert the analyzer/weighting behaviour benchmarked in the `emet`
|
||||
evaluation lab. A change that makes one of these fail is a change that
|
||||
invalidates that benchmark, not just a failing test.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from src.config import SparseEmbeddingSettings
|
||||
from src.infrastructure.embedding.analyzers import analyze
|
||||
from src.infrastructure.embedding.bm25 import Bm25SparseEmbedder, _token_index
|
||||
|
||||
pytestmark = pytest.mark.unit
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def settings() -> SparseEmbeddingSettings:
|
||||
return SparseEmbeddingSettings()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def embedder(settings: SparseEmbeddingSettings) -> Bm25SparseEmbedder:
|
||||
return Bm25SparseEmbedder(settings)
|
||||
|
||||
|
||||
# --- analyzer ------------------------------------------------------------
|
||||
|
||||
|
||||
def test_analyze_keeps_digits_as_tokens() -> None:
|
||||
"""Policy numbers, dates, and amounts are exactly what a lexical index
|
||||
should match on. An earlier implementation dropped every digit.
|
||||
"""
|
||||
tokens = analyze("بیمهنامه شماره ۱۲۳۴۵ صادر شد")
|
||||
assert "12345" in tokens
|
||||
|
||||
|
||||
def test_analyze_folds_eastern_digits_to_ascii() -> None:
|
||||
"""The same number must match however it was authored."""
|
||||
assert analyze("۹۹۸۸۷۷") == analyze("٩٩٨٨٧٧") == analyze("998877")
|
||||
|
||||
|
||||
def test_analyze_splits_on_zwnj() -> None:
|
||||
"""ZWNJ joins compounds visually but they are separate index terms."""
|
||||
assert analyze("آتشسوزی") == ["آتش", "سوزی"]
|
||||
|
||||
|
||||
def test_analyze_folds_arabic_letterforms_to_persian() -> None:
|
||||
# Arabic kaf (U+0643) vs. Persian keheh (U+06A9): the same word typed on
|
||||
# two different keyboards must produce the same term.
|
||||
assert analyze("كتاب") == analyze("کتاب")
|
||||
|
||||
|
||||
def test_analyze_drops_persian_stopwords() -> None:
|
||||
assert analyze("این کتاب و آن مداد") == analyze("کتاب مداد")
|
||||
|
||||
|
||||
def test_analyze_drops_english_stopwords() -> None:
|
||||
"""The corpus is mixed-script, so the list carries English too."""
|
||||
assert analyze("the policy is valid") == ["policy", "valid"]
|
||||
|
||||
|
||||
def test_analyze_lowercases_latin() -> None:
|
||||
assert analyze("POLICY Number") == ["policy", "number"]
|
||||
|
||||
|
||||
def test_analyze_rejects_unknown_analyzer() -> None:
|
||||
with pytest.raises(ValueError, match="Unknown analyzer"):
|
||||
analyze("متن", "fa_norm_stem")
|
||||
|
||||
|
||||
# --- token indexing ------------------------------------------------------
|
||||
|
||||
|
||||
def test_token_index_is_stable_across_calls() -> None:
|
||||
assert _token_index("کتاب") == _token_index("کتاب")
|
||||
|
||||
|
||||
def test_token_index_is_pinned_to_known_values() -> None:
|
||||
"""A golden test. These indices are baked into every stored sparse vector,
|
||||
so changing the hash silently orphans the whole index -- a re-ingestion,
|
||||
not a deploy. Ingest-time and query-time encoding must agree forever.
|
||||
"""
|
||||
assert _token_index("کتاب") == 1701064151
|
||||
assert _token_index("policy") == 741331709
|
||||
assert _token_index("12345") == 1232178634
|
||||
|
||||
|
||||
def test_token_index_fits_signed_int32() -> None:
|
||||
for token in ("کتاب", "policy", "12345", "بیمه", "x" * 200):
|
||||
assert 0 <= _token_index(token) < 2**31 - 1
|
||||
|
||||
|
||||
def test_token_index_distinguishes_different_tokens() -> None:
|
||||
assert _token_index("کتاب") != _token_index("مداد")
|
||||
|
||||
|
||||
# --- vector construction -------------------------------------------------
|
||||
|
||||
|
||||
def test_embed_batch_returns_one_vector_per_text(embedder: Bm25SparseEmbedder) -> None:
|
||||
assert len(embedder.embed_batch(["سلام دنیا", "یک تست دیگر"])) == 2
|
||||
|
||||
|
||||
def test_embed_batch_empty_text_returns_empty_vector(embedder: Bm25SparseEmbedder) -> None:
|
||||
(vector,) = embedder.embed_batch([""])
|
||||
assert vector.indices == []
|
||||
assert vector.values == []
|
||||
|
||||
|
||||
def test_embed_batch_all_stopwords_returns_empty_vector(embedder: Bm25SparseEmbedder) -> None:
|
||||
(vector,) = embedder.embed_batch(["و در به از که"])
|
||||
assert vector.indices == []
|
||||
|
||||
|
||||
def test_embed_batch_emits_tokens_in_sorted_order(embedder: Bm25SparseEmbedder) -> None:
|
||||
"""Deterministic output keeps re-ingestion byte-stable."""
|
||||
text = "مداد کتاب دفتر"
|
||||
expected = [_token_index(token) for token in sorted(analyze(text))]
|
||||
(vector,) = embedder.embed_batch([text])
|
||||
assert vector.indices == expected
|
||||
|
||||
|
||||
def test_embed_batch_applies_bm25_saturation_not_raw_counts(
|
||||
embedder: Bm25SparseEmbedder,
|
||||
) -> None:
|
||||
"""Weight must be sublinear in term frequency: tripling a term must not
|
||||
triple its weight, which is the whole point of the `k` parameter.
|
||||
"""
|
||||
(once,) = embedder.embed_batch(["کتاب"])
|
||||
(thrice,) = embedder.embed_batch(["کتاب کتاب کتاب"])
|
||||
assert thrice.values[0] > once.values[0]
|
||||
assert thrice.values[0] < 3 * once.values[0]
|
||||
|
||||
|
||||
def test_embed_batch_query_side_omits_length_normalization(
|
||||
embedder: Bm25SparseEmbedder, settings: SparseEmbeddingSettings
|
||||
) -> None:
|
||||
text = "کتاب مداد دفتر خودکار"
|
||||
(document,) = embedder.embed_batch([text], query=False)
|
||||
(query,) = embedder.embed_batch([text], query=True)
|
||||
|
||||
assert document.indices == query.indices # same terms, same hashing
|
||||
assert document.values != query.values
|
||||
|
||||
k = settings.k
|
||||
assert query.values[0] == pytest.approx(1.0 * (k + 1.0) / (1.0 + k))
|
||||
|
||||
|
||||
def test_embed_batch_short_document_outweighs_long_one(
|
||||
embedder: Bm25SparseEmbedder,
|
||||
) -> None:
|
||||
"""The `b` term discounts a term appearing in a longer document."""
|
||||
(short,) = embedder.embed_batch(["کتاب"])
|
||||
(long,) = embedder.embed_batch(["کتاب " + " ".join(f"واژه{i}" for i in range(200))])
|
||||
|
||||
short_weight = short.values[short.indices.index(_token_index("کتاب"))]
|
||||
long_weight = long.values[long.indices.index(_token_index("کتاب"))]
|
||||
assert short_weight > long_weight
|
||||
|
||||
|
||||
def test_embed_batch_respects_configured_parameters() -> None:
|
||||
"""k/b/avg_len come from settings, so they can be retuned without a code
|
||||
change -- and so a retune is visibly a config decision.
|
||||
"""
|
||||
default = Bm25SparseEmbedder(SparseEmbeddingSettings())
|
||||
tuned = Bm25SparseEmbedder(SparseEmbeddingSettings(k=2.5, b=0.2, avg_len=64.0))
|
||||
text = "کتاب کتاب مداد"
|
||||
|
||||
assert default.embed_batch([text])[0].values != tuned.embed_batch([text])[0].values
|
||||
Reference in New Issue
Block a user