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:
211
tests/unit/application/test_embedding.py
Normal file
211
tests/unit/application/test_embedding.py
Normal 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,
|
||||
)
|
||||
Reference in New Issue
Block a user