azk/chatbot_v3
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ingestion): add bounded, benchmark-aligned embedding execution

Browse Source 
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:
Ali Zarin Kolah
2026-08-19 17:13:32 +03:30
parent aa6d595424
commit 5c0a5938f8
33 changed files with 2455 additions and 536 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
tests/conftest.py
View File

@@ -12,12 +12,22 @@ from src.main import create_app
@pytest.fixture
def settings() -> Settings:
# Every external dependency points at a closed port so unit tests never
# reach real infrastructure. For the embedders this matters twice over:
# the real defaults are a colleague's Ollama box and OpenAI's paid API,
# and ADR-0016 forbids routine runs calling either. Connection-refused is
# immediate, so the lifespan's fail-soft warm-up costs nothing here --
# and these tests passing at all is what proves it is fail-soft.
return Settings(
postgres={"host": "127.0.0.1", "port": 1},
minio={"endpoint": "127.0.0.1:1"},
ingestion={"timeout_seconds": 1.0},
qdrant={"url": "http://127.0.0.1:1"},
app={"readiness_check_timeout_seconds": 0.5},
embedding={
"nomic": {"base_url": "http://127.0.0.1:1/v1", "timeout_seconds": 0.5},
"openai": {"base_url": "http://127.0.0.1:1/v1", "timeout_seconds": 0.5},
},
)

43
tests/fakes.py
View File

@@ -1,7 +1,11 @@
"""Hand-written fakes for narrow application-owned ports (ADR-0016)."""
import asyncio
from collections.abc import Sequence
from dataclasses import dataclass, field
from src.application.ingestion.models import SparseVector
@dataclass
class FakeObjectStorage:
@@ -15,3 +19,42 @@ class FakeObjectStorage:
self.fail_next = False
raise OSError("simulated object storage failure")
self.objects[key] = data
@dataclass
class FakeDenseEmbedder:
"""A scripted `DenseEmbedder`. Returns a fixed-dimension zero vector per
text by default; `fail_next` simulates one batch failure.
"""
name: str
dimensions: int = 4
calls: list[list[str]] = field(default_factory=list)
fail_next: bool = False
delay_seconds: float = 0.0
"""Simulates a slow provider call, e.g. to exercise timeout handling."""
async def embed_batch(self, texts: Sequence[str]) -> list[list[float]]:
self.calls.append(list(texts))
if self.delay_seconds:
await asyncio.sleep(self.delay_seconds)
if self.fail_next:
self.fail_next = False
raise RuntimeError("simulated embedder failure")
return [[0.0] * self.dimensions for _ in texts]
@dataclass
class FakeSparseEmbedder:
"""A scripted `SparseEmbedder`. Returns an empty sparse vector per text."""
name: str = "sparse"
calls: list[list[str]] = field(default_factory=list)
fail_next: bool = False
def embed_batch(self, texts: Sequence[str], *, query: bool = False) -> list[SparseVector]:
self.calls.append(list(texts))
if self.fail_next:
self.fail_next = False
raise RuntimeError("simulated embedder failure")
return [SparseVector(indices=[], values=[]) for _ in texts]

210
tests/integration/postgres/test_upload_service.py
View File

@@ -1,19 +1,23 @@
"""ADR-0017's two-transaction upload shape, exercised against real Postgres.
Object storage is faked (`FakeObjectStorage`) -- it's a port, not the thing
under test here. `MinioObjectStorage` itself is covered in
`tests/integration/minio/`.
Object storage and the embedders are faked (`tests/fakes.py`) -- they're
ports, not the thing under test here. `MinioObjectStorage` itself is covered
in `tests/integration/minio/`.
"""
import pytest
from anyio import CapacityLimiter
from anyio import CapacityLimiter, Semaphore
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from src.application.auth.context import AuthContext
from src.application.files.models import UploadResult
from src.application.files.upload import upload_source_file
from src.application.ingestion.errors import IngestionAtCapacityError, IngestionTimeoutError
from src.application.ports.object_storage import ObjectStorage
from src.config import ChunkingSettings, IngestionSettings
from src.infrastructure.postgres.models.ingestion_job import IngestionJob
from tests.fakes import FakeObjectStorage
from tests.fakes import FakeDenseEmbedder, FakeObjectStorage, FakeSparseEmbedder
from tests.support.factories import create_api_key, create_tenant
pytestmark = [
@@ -25,6 +29,34 @@ pytestmark = [
_CSV_BYTES = b"name,value\nfirst,1\n"
async def _upload(
*,
sessionmaker: async_sessionmaker[AsyncSession],
storage: ObjectStorage,
auth: AuthContext,
domain: str = "general",
filename: str = "report.csv",
data: bytes = _CSV_BYTES,
) -> UploadResult:
return await upload_source_file(
sessionmaker=sessionmaker,
storage=storage,
auth=auth,
domain=domain,
filename=filename,
data=data,
ingestion_settings=IngestionSettings(),
chunking_settings=ChunkingSettings(),
thread_limiter=CapacityLimiter(2),
concurrency_limiter=Semaphore(2),
dense_embedders=[
FakeDenseEmbedder(name="dense_nomic"),
FakeDenseEmbedder(name="dense_openai"),
],
sparse_embedder=FakeSparseEmbedder(),
)
async def _auth_for(db_session: AsyncSession) -> AuthContext:
tenant = await create_tenant(db_session)
api_key, _ = await create_api_key(db_session, tenant=tenant)
@@ -48,17 +80,7 @@ async def test_upload_source_file_commits_running_job_before_storage_write(
auth = await _auth_for(db_session)
storage = FakeObjectStorage()
result = await upload_source_file(
sessionmaker=db_sessionmaker,
storage=storage,
auth=auth,
domain="general",
filename="report.csv",
data=_CSV_BYTES,
max_upload_size_bytes=1_000_000,
chunking_strategy="fixed_size",
validation_limiter=CapacityLimiter(2),
)
result = await _upload(sessionmaker=db_sessionmaker, storage=storage, auth=auth)
assert result.status == "succeeded"
assert result.chunks_indexed == 0
@@ -77,28 +99,8 @@ async def test_upload_source_file_duplicate_hash_returns_existing_without_reinge
auth = await _auth_for(db_session)
storage = FakeObjectStorage()
first = await upload_source_file(
sessionmaker=db_sessionmaker,
storage=storage,
auth=auth,
domain="general",
filename="report.csv",
data=_CSV_BYTES,
max_upload_size_bytes=1_000_000,
chunking_strategy="fixed_size",
validation_limiter=CapacityLimiter(2),
)
second = await upload_source_file(
sessionmaker=db_sessionmaker,
storage=storage,
auth=auth,
domain="general",
filename="report.csv",
data=_CSV_BYTES,
max_upload_size_bytes=1_000_000,
chunking_strategy="fixed_size",
validation_limiter=CapacityLimiter(2),
)
first = await _upload(sessionmaker=db_sessionmaker, storage=storage, auth=auth)
second = await _upload(sessionmaker=db_sessionmaker, storage=storage, auth=auth)
assert second.file_id == first.file_id
assert second.ingestion_job_id == first.ingestion_job_id
@@ -117,17 +119,7 @@ async def test_upload_source_file_retries_after_failed_job(
failing_storage = FakeObjectStorage(fail_next=True)
with pytest.raises(OSError):
await upload_source_file(
sessionmaker=db_sessionmaker,
storage=failing_storage,
auth=auth,
domain="general",
filename="report.csv",
data=_CSV_BYTES,
max_upload_size_bytes=1_000_000,
chunking_strategy="fixed_size",
validation_limiter=CapacityLimiter(2),
)
await _upload(sessionmaker=db_sessionmaker, storage=failing_storage, auth=auth)
async with db_sessionmaker() as verify_session:
jobs = (
@@ -142,17 +134,7 @@ async def test_upload_source_file_retries_after_failed_job(
assert len(jobs) == 1
assert jobs[0].status == "failed"
retry = await upload_source_file(
sessionmaker=db_sessionmaker,
storage=FakeObjectStorage(),
auth=auth,
domain="general",
filename="report.csv",
data=_CSV_BYTES,
max_upload_size_bytes=1_000_000,
chunking_strategy="fixed_size",
validation_limiter=CapacityLimiter(2),
)
retry = await _upload(sessionmaker=db_sessionmaker, storage=FakeObjectStorage(), auth=auth)
assert retry.status == "succeeded"
assert retry.is_new_attempt
@@ -168,27 +150,91 @@ async def test_upload_source_file_is_tenant_isolated_for_identical_content(
auth_a = await _auth_for(db_session)
auth_b = await _auth_for(db_session)
result_a = await upload_source_file(
sessionmaker=db_sessionmaker,
storage=FakeObjectStorage(),
auth=auth_a,
domain="general",
filename="report.csv",
data=_CSV_BYTES,
max_upload_size_bytes=1_000_000,
chunking_strategy="fixed_size",
validation_limiter=CapacityLimiter(2),
)
result_b = await upload_source_file(
sessionmaker=db_sessionmaker,
storage=FakeObjectStorage(),
auth=auth_b,
domain="general",
filename="report.csv",
data=_CSV_BYTES,
max_upload_size_bytes=1_000_000,
chunking_strategy="fixed_size",
validation_limiter=CapacityLimiter(2),
)
result_a = await _upload(sessionmaker=db_sessionmaker, storage=FakeObjectStorage(), auth=auth_a)
result_b = await _upload(sessionmaker=db_sessionmaker, storage=FakeObjectStorage(), auth=auth_b)
assert result_a.file_id != result_b.file_id
async def test_upload_source_file_timeout_writes_failed_job_and_raises(
db_session: AsyncSession, db_sessionmaker: async_sessionmaker[AsyncSession]
) -> None:
"""A work phase that outruns `INGESTION_TIMEOUT_SECONDS` must still leave
a terminal `failed` job -- never one stuck in `running` (ADR-0017, plan
001 Phase 4).
"""
auth = await _auth_for(db_session)
slow_embedder = FakeDenseEmbedder(name="dense_nomic", delay_seconds=10)
with pytest.raises(IngestionTimeoutError):
await upload_source_file(
sessionmaker=db_sessionmaker,
storage=FakeObjectStorage(),
auth=auth,
domain="general",
filename="report.csv",
data=_CSV_BYTES,
ingestion_settings=IngestionSettings(timeout_seconds=0.05),
chunking_settings=ChunkingSettings(),
thread_limiter=CapacityLimiter(2),
concurrency_limiter=Semaphore(2),
dense_embedders=[slow_embedder, FakeDenseEmbedder(name="dense_openai")],
sparse_embedder=FakeSparseEmbedder(),
)
async with db_sessionmaker() as verify_session:
jobs = (
(
await verify_session.execute(
select(IngestionJob).where(IngestionJob.tenant_id == auth.tenant_id)
)
)
.scalars()
.all()
)
assert len(jobs) == 1
assert jobs[0].status == "failed"
assert jobs[0].error_code == "timeout"
async def test_upload_source_file_at_capacity_rejects_before_any_job_row(
db_session: AsyncSession, db_sessionmaker: async_sessionmaker[AsyncSession]
) -> None:
"""`INGESTION_MAX_CONCURRENCY` rejects outright rather than queuing
(ADR-0017) -- and does so before txn A, so a rejected request leaves no
job row behind at all.
"""
auth = await _auth_for(db_session)
concurrency_limiter = Semaphore(1)
concurrency_limiter.acquire_nowait() # simulate the one slot already in use
with pytest.raises(IngestionAtCapacityError):
await upload_source_file(
sessionmaker=db_sessionmaker,
storage=FakeObjectStorage(),
auth=auth,
domain="general",
filename="report.csv",
data=_CSV_BYTES,
ingestion_settings=IngestionSettings(),
chunking_settings=ChunkingSettings(),
thread_limiter=CapacityLimiter(2),
concurrency_limiter=concurrency_limiter,
dense_embedders=[
FakeDenseEmbedder(name="dense_nomic"),
FakeDenseEmbedder(name="dense_openai"),
],
sparse_embedder=FakeSparseEmbedder(),
)
async with db_sessionmaker() as verify_session:
jobs = (
(
await verify_session.execute(
select(IngestionJob).where(IngestionJob.tenant_id == auth.tenant_id)
)
)
.scalars()
.all()
)
assert jobs == []

78
tests/unit/application/test_bounds.py Normal file
View File

@@ -0,0 +1,78 @@
"""Request bounds for inline ingestion: chunk ceiling and the process-wide
concurrency gate (ADR-0017, plan 001 Phase 4).
"""
import uuid
import pytest
from anyio import Semaphore
from src.application.ingestion.bounds import acquire_ingestion_slot, enforce_chunk_limit
from src.application.ingestion.errors import ChunkLimitExceededError, IngestionAtCapacityError
from src.application.ingestion.models import Chunk, ContentType
pytestmark = pytest.mark.unit
def _chunk(index: int) -> Chunk:
return Chunk(
chunk_id=uuid.uuid4(),
chunk_index=index,
order_id=float(index + 1),
content="x",
content_type=ContentType.PARAGRAPH,
token_count=1,
character_count=1,
)
def test_enforce_chunk_limit_within_bound_does_not_raise() -> None:
enforce_chunk_limit([_chunk(0), _chunk(1)], max_chunks=2)
def test_enforce_chunk_limit_over_bound_raises() -> None:
with pytest.raises(ChunkLimitExceededError):
enforce_chunk_limit([_chunk(0), _chunk(1), _chunk(2)], max_chunks=2)
@pytest.mark.asyncio
async def test_acquire_ingestion_slot_allows_up_to_the_limit() -> None:
limiter = Semaphore(2)
async with acquire_ingestion_slot(limiter), acquire_ingestion_slot(limiter):
pass # two concurrent holders within a limit of two: no rejection
@pytest.mark.asyncio
async def test_acquire_ingestion_slot_rejects_beyond_the_limit() -> None:
limiter = Semaphore(1)
async with acquire_ingestion_slot(limiter):
with pytest.raises(IngestionAtCapacityError):
async with acquire_ingestion_slot(limiter):
pass
@pytest.mark.asyncio
async def test_acquire_ingestion_slot_releases_on_exit() -> None:
limiter = Semaphore(1)
async with acquire_ingestion_slot(limiter):
pass
# The slot from the first `async with` must be released by the time it
# exits, or every subsequent request would see permanent capacity loss.
async with acquire_ingestion_slot(limiter):
pass
@pytest.mark.asyncio
async def test_acquire_ingestion_slot_releases_after_body_raises() -> None:
limiter = Semaphore(1)
with pytest.raises(ValueError, match="boom"):
async with acquire_ingestion_slot(limiter):
raise ValueError("boom")
async with acquire_ingestion_slot(limiter):
pass

211
tests/unit/application/test_embedding.py Normal file
View File

@@ -0,0 +1,211 @@
"""Batching, concurrency bounding, and failure translation for `embed_chunks`
(ADR-0001, ADR-0017, plan 001 Phase 4).
"""
import asyncio
import uuid
from collections.abc import Sequence
from dataclasses import dataclass, field
import pytest
from anyio import CapacityLimiter
from src.application.ingestion.embedding import embed_chunks
from src.application.ingestion.errors import EmbedderError
from src.application.ingestion.models import Chunk, ContentType, SparseVector
from src.config import IngestionSettings
pytestmark = [pytest.mark.unit, pytest.mark.asyncio]
def _chunk(index: int, content: str = "hello world") -> Chunk:
return Chunk(
chunk_id=uuid.uuid4(),
chunk_index=index,
order_id=float(index + 1),
content=content,
content_type=ContentType.PARAGRAPH,
token_count=2,
character_count=len(content),
)
@dataclass
class _TrackingDenseEmbedder:
"""Records each batch's texts and the peak number of batches in flight
at once, to prove concurrency is bounded, not serial.
"""
name: str
dimensions: int = 3
batches: list[list[str]] = field(default_factory=list)
in_flight: int = 0
peak_in_flight: int = 0
async def embed_batch(self, texts: Sequence[str]) -> list[list[float]]:
self.in_flight += 1
self.peak_in_flight = max(self.peak_in_flight, self.in_flight)
self.batches.append(list(texts))
await asyncio.sleep(0) # yield so overlapping calls can interleave
self.in_flight -= 1
return [[0.0] * self.dimensions for _ in texts]
@dataclass
class _FailingDenseEmbedder:
name: str
async def embed_batch(self, texts: Sequence[str]) -> list[list[float]]:
raise RuntimeError("boom")
@dataclass
class _StubSparseEmbedder:
name: str = "sparse"
calls: list[list[str]] = field(default_factory=list)
def embed_batch(self, texts: Sequence[str], *, query: bool = False) -> list[SparseVector]:
self.calls.append(list(texts))
return [SparseVector(indices=[i], values=[1.0]) for i in range(len(texts))]
@dataclass
class _FailingSparseEmbedder:
name: str = "sparse"
def embed_batch(self, texts: Sequence[str], *, query: bool = False) -> list[SparseVector]:
raise RuntimeError("boom")
@pytest.fixture
def settings() -> IngestionSettings:
return IngestionSettings(embed_batch_size=2, embed_concurrency=2)
@pytest.fixture
def thread_limiter() -> CapacityLimiter:
return CapacityLimiter(4)
async def test_embed_chunks_empty_input_returns_empty(
settings: IngestionSettings, thread_limiter: CapacityLimiter
) -> None:
result = await embed_chunks(
[],
dense_embedders=[_TrackingDenseEmbedder(name="dense_nomic")],
sparse_embedder=_StubSparseEmbedder(),
settings=settings,
thread_limiter=thread_limiter,
)
assert result == []
async def test_embed_chunks_batches_before_parallelizing(
settings: IngestionSettings, thread_limiter: CapacityLimiter
) -> None:
chunks = [_chunk(i) for i in range(5)]
embedder = _TrackingDenseEmbedder(name="dense_nomic")
await embed_chunks(
chunks,
dense_embedders=[embedder],
sparse_embedder=_StubSparseEmbedder(),
settings=settings,
thread_limiter=thread_limiter,
)
# embed_batch_size=2 over 5 chunks -> 3 batches (2, 2, 1), never one call
# per chunk.
assert len(embedder.batches) == 3
assert [len(batch) for batch in embedder.batches] == [2, 2, 1]
async def test_embed_chunks_bounds_concurrency_by_embed_concurrency(
thread_limiter: CapacityLimiter,
) -> None:
settings = IngestionSettings(embed_batch_size=1, embed_concurrency=2)
chunks = [_chunk(i) for i in range(6)]
embedder = _TrackingDenseEmbedder(name="dense_nomic")
await embed_chunks(
chunks,
dense_embedders=[embedder],
sparse_embedder=_StubSparseEmbedder(),
settings=settings,
thread_limiter=thread_limiter,
)
assert embedder.peak_in_flight <= settings.embed_concurrency
assert embedder.peak_in_flight > 1 # proves it isn't serial either
async def test_embed_chunks_passes_chunk_text_through_unmodified(
settings: IngestionSettings, thread_limiter: CapacityLimiter
) -> None:
"""Text shaping (task prefixes, `keep_alive`) is the adapter's job, not
this module's -- see `src/infrastructure/embedding/openai_compatible.py`.
Orchestration here must stay provider-agnostic.
"""
chunks = [_chunk(0, content="salam")]
nomic = _TrackingDenseEmbedder(name="dense_nomic")
openai = _TrackingDenseEmbedder(name="dense_openai")
await embed_chunks(
chunks,
dense_embedders=[nomic, openai],
sparse_embedder=_StubSparseEmbedder(),
settings=settings,
thread_limiter=thread_limiter,
)
assert nomic.batches[0] == ["salam"]
assert openai.batches[0] == ["salam"]
async def test_embed_chunks_assembles_dense_and_sparse_per_chunk_in_order(
settings: IngestionSettings, thread_limiter: CapacityLimiter
) -> None:
chunks = [_chunk(0, "a"), _chunk(1, "b")]
nomic = _TrackingDenseEmbedder(name="dense_nomic", dimensions=3)
result = await embed_chunks(
chunks,
dense_embedders=[nomic],
sparse_embedder=_StubSparseEmbedder(),
settings=settings,
thread_limiter=thread_limiter,
)
assert [item.chunk.chunk_id for item in result] == [c.chunk_id for c in chunks]
assert all(len(item.dense["dense_nomic"]) == 3 for item in result)
assert all("sparse" not in item.dense for item in result)
async def test_embed_chunks_dense_failure_raises_embedder_error(
settings: IngestionSettings, thread_limiter: CapacityLimiter
) -> None:
chunks = [_chunk(0)]
with pytest.raises(EmbedderError):
await embed_chunks(
chunks,
dense_embedders=[_FailingDenseEmbedder(name="dense_nomic")],
sparse_embedder=_StubSparseEmbedder(),
settings=settings,
thread_limiter=thread_limiter,
)
async def test_embed_chunks_sparse_failure_raises_embedder_error(
settings: IngestionSettings, thread_limiter: CapacityLimiter
) -> None:
chunks = [_chunk(0)]
with pytest.raises(EmbedderError):
await embed_chunks(
chunks,
dense_embedders=[_TrackingDenseEmbedder(name="dense_nomic")],
sparse_embedder=_FailingSparseEmbedder(),
settings=settings,
thread_limiter=thread_limiter,
)

31
tests/unit/bootstrap/test_lifespan.py
View File

@@ -2,7 +2,9 @@ import pytest
from fastapi import FastAPI
from httpx import AsyncClient
from src.bootstrap.lifespan import _warm_dense_embedders
from src.config import Settings
from tests.fakes import FakeDenseEmbedder
pytestmark = [pytest.mark.unit, pytest.mark.asyncio]
@@ -17,3 +19,32 @@ async def test_lifespan_binds_resources_to_app_state(app: FastAPI, client: Async
resources = app.state.resources
assert isinstance(resources.settings, Settings)
assert resources.minio_client is not None
assert len(resources.dense_embedders) == 2
assert {e.name for e in resources.dense_embedders} == {"dense_nomic", "dense_openai"}
assert resources.sparse_embedder.name == "sparse"
assert resources.ingestion_concurrency_limiter is not None
async def test_warm_dense_embedders_calls_every_embedder() -> None:
"""Pays the model-load cost at boot instead of on a user's first upload:
a cold nomic load outruns INGESTION_TIMEOUT_SECONDS entirely.
"""
embedders = [FakeDenseEmbedder(name="dense_nomic"), FakeDenseEmbedder(name="dense_openai")]
await _warm_dense_embedders(embedders)
assert all(e.calls for e in embedders)
async def test_warm_dense_embedders_survives_an_unreachable_embedder() -> None:
"""A down embedder must not stop the process booting -- otherwise the
service cannot come up to report its own health. `/readyz` owns that
signal, not startup.
"""
failing = FakeDenseEmbedder(name="dense_nomic", fail_next=True)
healthy = FakeDenseEmbedder(name="dense_openai")
await _warm_dense_embedders([failing, healthy])
# The failure is swallowed *and* does not abort the remaining warm-ups.
assert healthy.calls

0
tests/unit/infrastructure/__init__.py Normal file
View File

0
tests/unit/infrastructure/embedding/__init__.py Normal file
View File

170
tests/unit/infrastructure/embedding/test_bm25.py Normal file
View 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

170
tests/unit/infrastructure/embedding/test_openai_compatible.py Normal file
View File

@@ -0,0 +1,170 @@
"""`OpenAICompatibleEmbedder` against a mocked `/embeddings` endpoint.
Backs both `dense_nomic` and `dense_openai` (ADR-0001) -- `httpx.MockTransport`
stands in for the real self-hosted/OpenAI server so this stays a unit test
with no network dependency.
"""
import json
from collections.abc import Callable
import httpx
import pytest
from src.infrastructure.embedding.openai_compatible import (
OpenAICompatibleEmbedder,
is_ollama_base_url,
)
pytestmark = pytest.mark.unit
def _client(handler: Callable[[httpx.Request], httpx.Response]) -> httpx.AsyncClient:
return httpx.AsyncClient(transport=httpx.MockTransport(handler), base_url="http://embedder")
@pytest.mark.asyncio
async def test_embed_batch_returns_vectors_in_input_order() -> None:
def handler(request: httpx.Request) -> httpx.Response:
# Respond out of order to prove the adapter re-sorts by `index`.
return httpx.Response(
200,
json={
"data": [
{"index": 1, "embedding": [0.2, 0.2]},
{"index": 0, "embedding": [0.1, 0.1]},
],
"model": "test-model",
},
)
embedder = OpenAICompatibleEmbedder(_client(handler), name="dense_nomic", model="test-model")
vectors = await embedder.embed_batch(["first", "second"])
assert vectors == [[0.1, 0.1], [0.2, 0.2]]
@pytest.mark.asyncio
async def test_embed_batch_sends_model_and_input() -> None:
captured: dict[str, object] = {}
def handler(request: httpx.Request) -> httpx.Response:
captured.update(json.loads(request.read()))
return httpx.Response(
200, json={"data": [{"index": 0, "embedding": [0.0]}], "model": "test-model"}
)
embedder = OpenAICompatibleEmbedder(_client(handler), name="dense_openai", model="test-model")
await embedder.embed_batch(["only text"])
assert captured["model"] == "test-model"
assert captured["input"] == ["only text"]
assert "dimensions" not in captured
@pytest.mark.asyncio
async def test_embed_batch_sends_dimensions_when_configured() -> None:
captured: dict[str, object] = {}
def handler(request: httpx.Request) -> httpx.Response:
captured.update(json.loads(request.read()))
return httpx.Response(
200, json={"data": [{"index": 0, "embedding": [0.0] * 256}], "model": "test-model"}
)
embedder = OpenAICompatibleEmbedder(
_client(handler), name="dense_openai", model="test-model", dimensions=256
)
await embedder.embed_batch(["only text"])
assert captured["dimensions"] == 256
@pytest.mark.asyncio
async def test_embed_batch_raises_on_non_2xx_response() -> None:
def handler(request: httpx.Request) -> httpx.Response:
return httpx.Response(500, json={"error": "boom"})
embedder = OpenAICompatibleEmbedder(_client(handler), name="dense_nomic", model="test-model")
with pytest.raises(httpx.HTTPStatusError):
await embedder.embed_batch(["text"])
@pytest.mark.asyncio
async def test_embed_batch_sends_no_prefix_by_default() -> None:
"""emet parity: the benchmarked run used no task prefix (ADR-0004)."""
captured: dict[str, object] = {}
def handler(request: httpx.Request) -> httpx.Response:
captured.update(json.loads(request.read()))
return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.0]}], "model": "m"})
embedder = OpenAICompatibleEmbedder(_client(handler), name="dense_nomic", model="m")
await embedder.embed_batch(["سلام"])
assert captured["input"] == ["سلام"]
@pytest.mark.asyncio
async def test_embed_batch_applies_document_prefix_when_configured() -> None:
captured: dict[str, object] = {}
def handler(request: httpx.Request) -> httpx.Response:
captured.update(json.loads(request.read()))
return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.0]}], "model": "m"})
embedder = OpenAICompatibleEmbedder(
_client(handler), name="dense_nomic", model="m", document_prefix="search_document: "
)
await embedder.embed_batch(["سلام"])
assert captured["input"] == ["search_document: سلام"]
@pytest.mark.asyncio
async def test_embed_batch_sends_keep_alive_when_configured() -> None:
"""Keeps an Ollama-hosted model resident; a cold load outruns the
ingestion timeout entirely.
"""
captured: dict[str, object] = {}
def handler(request: httpx.Request) -> httpx.Response:
captured.update(json.loads(request.read()))
return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.0]}], "model": "m"})
embedder = OpenAICompatibleEmbedder(
_client(handler), name="dense_nomic", model="m", keep_alive="30m"
)
await embedder.embed_batch(["text"])
assert captured["keep_alive"] == "30m"
@pytest.mark.asyncio
async def test_embed_batch_omits_keep_alive_when_not_configured() -> None:
"""OpenAI would reject an unknown field, so it must not be sent there."""
captured: dict[str, object] = {}
def handler(request: httpx.Request) -> httpx.Response:
captured.update(json.loads(request.read()))
return httpx.Response(200, json={"data": [{"index": 0, "embedding": [0.0]}], "model": "m"})
embedder = OpenAICompatibleEmbedder(_client(handler), name="dense_openai", model="m")
await embedder.embed_batch(["text"])
assert "keep_alive" not in captured
@pytest.mark.parametrize(
("base_url", "expected"),
[
("http://192.168.10.10:11435/v1", True),
("http://127.0.0.1:11434/v1", True),
("http://ollama.internal/v1", True),
("https://api.openai.com/v1", False),
("http://127.0.0.1:8081/v1", False),
],
)
def test_is_ollama_base_url(base_url: str, expected: bool) -> None:
assert is_ollama_base_url(base_url) is expected