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:
48
src/application/ingestion/bounds.py
Normal file
48
src/application/ingestion/bounds.py
Normal file
@@ -0,0 +1,48 @@
|
||||
"""Request bounds for inline ingestion (ADR-0017).
|
||||
|
||||
Three independent bounds, each mapping to its own status code: the chunk
|
||||
ceiling (`413`, checked before embedding starts), process-wide concurrency
|
||||
(`503` + `Retry-After`, rejected rather than queued), and the work-phase
|
||||
deadline (`504`, and the caller must still write a terminal job status).
|
||||
"""
|
||||
|
||||
from collections.abc import AsyncIterator, Sequence
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from anyio import Semaphore, WouldBlock
|
||||
|
||||
from src.application.ingestion.errors import ChunkLimitExceededError, IngestionAtCapacityError
|
||||
from src.application.ingestion.models import Chunk
|
||||
|
||||
|
||||
def enforce_chunk_limit(chunks: Sequence[Chunk], *, max_chunks: int) -> None:
|
||||
"""Raise `ChunkLimitExceededError` if `chunks` exceeds `max_chunks`.
|
||||
|
||||
Call this immediately after parsing/chunking and before any embedding
|
||||
call — the ceiling must be discovered up front, not mid-batch.
|
||||
"""
|
||||
if len(chunks) > max_chunks:
|
||||
raise ChunkLimitExceededError(
|
||||
f"document produced {len(chunks)} chunks, over the {max_chunks}-chunk limit"
|
||||
)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def acquire_ingestion_slot(limiter: Semaphore) -> AsyncIterator[None]:
|
||||
"""Hold one of `INGESTION_MAX_CONCURRENCY` process-wide slots for the block.
|
||||
|
||||
`limiter` is an `anyio.Semaphore` created once in the lifespan. Rejects
|
||||
immediately with `IngestionAtCapacityError` when the process is already at
|
||||
capacity, rather than queueing the request behind an unbounded wait
|
||||
(ADR-0017) — the semaphore's own async `acquire()` would do the latter.
|
||||
"""
|
||||
try:
|
||||
limiter.acquire_nowait()
|
||||
except WouldBlock:
|
||||
raise IngestionAtCapacityError(
|
||||
"ingestion is at capacity; retry after the configured backoff"
|
||||
) from None
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
limiter.release()
|
||||
112
src/application/ingestion/embedding.py
Normal file
112
src/application/ingestion/embedding.py
Normal file
@@ -0,0 +1,112 @@
|
||||
"""The one caller-facing entry point for embedding chunks (ADR-0001, ADR-0017).
|
||||
|
||||
`embed_chunks` is the only version of this step callers should reach for: it
|
||||
owns batching, the `embed_concurrency` semaphore bounding in-flight dense
|
||||
batches, and the `anyio.to_thread.run_sync` + `CapacityLimiter` offload for
|
||||
the blocking BM25 pipeline. Composing these correctly at every call site is
|
||||
exactly the obligation a deep module absorbs once (see CLAUDE.md's "prefer
|
||||
deep modules").
|
||||
|
||||
Per-provider text shaping — task prefixes, `keep_alive`, request payload —
|
||||
belongs to the adapters in `src/infrastructure/embedding/`, not here. This
|
||||
module knows only that an embedder turns texts into vectors.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import Sequence
|
||||
from functools import partial
|
||||
|
||||
from anyio import CapacityLimiter, to_thread
|
||||
|
||||
from src.application.ingestion.errors import EmbedderError
|
||||
from src.application.ingestion.models import Chunk, EmbeddedChunk
|
||||
from src.application.ports.embedding import DenseEmbedder, SparseEmbedder
|
||||
from src.config import IngestionSettings
|
||||
|
||||
|
||||
def _batches(texts: Sequence[str], size: int) -> list[Sequence[str]]:
|
||||
return [texts[i : i + size] for i in range(0, len(texts), size)]
|
||||
|
||||
|
||||
async def _embed_dense_bounded(
|
||||
embedder: DenseEmbedder,
|
||||
batch: Sequence[str],
|
||||
*,
|
||||
semaphore: asyncio.Semaphore,
|
||||
) -> list[list[float]]:
|
||||
async with semaphore:
|
||||
try:
|
||||
return await embedder.embed_batch(batch)
|
||||
except Exception as exc:
|
||||
raise EmbedderError(f"{embedder.name} embedding batch failed: {exc}") from exc
|
||||
|
||||
|
||||
async def _embed_dense_all(
|
||||
embedder: DenseEmbedder,
|
||||
texts: Sequence[str],
|
||||
*,
|
||||
batch_size: int,
|
||||
semaphore: asyncio.Semaphore,
|
||||
) -> list[list[float]]:
|
||||
batches = _batches(texts, batch_size)
|
||||
results = await asyncio.gather(
|
||||
*(_embed_dense_bounded(embedder, batch, semaphore=semaphore) for batch in batches)
|
||||
)
|
||||
return [vector for batch_result in results for vector in batch_result]
|
||||
|
||||
|
||||
def _embed_sparse_sync(embedder: SparseEmbedder, texts: Sequence[str]):
|
||||
try:
|
||||
return embedder.embed_batch(texts)
|
||||
except Exception as exc:
|
||||
raise EmbedderError(f"{embedder.name} embedding batch failed: {exc}") from exc
|
||||
|
||||
|
||||
async def embed_chunks(
|
||||
chunks: Sequence[Chunk],
|
||||
*,
|
||||
dense_embedders: Sequence[DenseEmbedder],
|
||||
sparse_embedder: SparseEmbedder,
|
||||
settings: IngestionSettings,
|
||||
thread_limiter: CapacityLimiter,
|
||||
) -> list[EmbeddedChunk]:
|
||||
"""Embed every chunk into all dense vectors plus the sparse vector.
|
||||
|
||||
Dense embedders run concurrently with each other; each one's batches are
|
||||
concurrent among themselves too, bounded by one `embed_concurrency`
|
||||
semaphore shared across all dense embedders (ADR-0017: the limit exists
|
||||
for both providers' rate limits and the self-hosted server's capacity —
|
||||
not a per-provider budget). The sparse (BM25) pass is blocking and runs
|
||||
once, off the event loop.
|
||||
|
||||
Raises `EmbedderError` (502) if any embedder call fails.
|
||||
"""
|
||||
if not chunks:
|
||||
return []
|
||||
|
||||
texts = [chunk.content for chunk in chunks]
|
||||
semaphore = asyncio.Semaphore(settings.embed_concurrency)
|
||||
|
||||
dense_task = asyncio.gather(
|
||||
*(
|
||||
_embed_dense_all(
|
||||
embedder, texts, batch_size=settings.embed_batch_size, semaphore=semaphore
|
||||
)
|
||||
for embedder in dense_embedders
|
||||
)
|
||||
)
|
||||
sparse_task = to_thread.run_sync(
|
||||
partial(_embed_sparse_sync, sparse_embedder, texts), limiter=thread_limiter
|
||||
)
|
||||
dense_results, sparse_vectors = await asyncio.gather(dense_task, sparse_task)
|
||||
|
||||
dense_by_name = {
|
||||
embedder.name: vectors
|
||||
for embedder, vectors in zip(dense_embedders, dense_results, strict=True)
|
||||
}
|
||||
|
||||
embedded: list[EmbeddedChunk] = []
|
||||
for index, chunk in enumerate(chunks):
|
||||
dense = {name: vectors[index] for name, vectors in dense_by_name.items()}
|
||||
embedded.append(EmbeddedChunk(chunk=chunk, dense=dense, sparse=sparse_vectors[index]))
|
||||
return embedded
|
||||
@@ -39,3 +39,26 @@ class ChunkTooLargeError(IngestionError):
|
||||
against is silent — `nomic-embed-text-v2-moe` truncates over-long input
|
||||
without raising (ADR-0004).
|
||||
"""
|
||||
|
||||
|
||||
class EmbedderError(IngestionError):
|
||||
"""A dense or sparse embedder call failed (transport error, non-2xx, or
|
||||
a malformed response).
|
||||
|
||||
Maps to `502` per ADR-0017.
|
||||
"""
|
||||
|
||||
|
||||
class IngestionAtCapacityError(IngestionError):
|
||||
"""`INGESTION_MAX_CONCURRENCY` in-process ingestions are already running.
|
||||
|
||||
Maps to `503` with `Retry-After`, not a queued wait (ADR-0017).
|
||||
"""
|
||||
|
||||
|
||||
class IngestionTimeoutError(IngestionError):
|
||||
"""The work phase (parse/embed/upsert) exceeded `INGESTION_TIMEOUT_SECONDS`.
|
||||
|
||||
Maps to `504`. The caller must still write a terminal `failed` job status
|
||||
before this propagates (ADR-0017).
|
||||
"""
|
||||
|
||||
@@ -64,3 +64,28 @@ class Chunk(BaseModel):
|
||||
next_chunk_id: uuid.UUID | None = None
|
||||
token_count: int
|
||||
character_count: int
|
||||
|
||||
|
||||
class SparseVector(BaseModel):
|
||||
"""A sparse (term-index -> weight) vector, Qdrant's `modifier="idf"` shape.
|
||||
|
||||
Kept free of the `qdrant_client` SDK (ADR-0015: ports carry no infra
|
||||
imports) — `src/infrastructure/qdrant/` converts this to the SDK's own
|
||||
`SparseVector` type at upsert time (Phase 5).
|
||||
"""
|
||||
|
||||
indices: list[int]
|
||||
values: list[float]
|
||||
|
||||
|
||||
class EmbeddedChunk(BaseModel):
|
||||
"""A chunk plus every vector it will be upserted with (ADR-0001).
|
||||
|
||||
`dense` is keyed by named-vector name (`dense_nomic`, `dense_openai`).
|
||||
`late_interaction` is deliberately absent — not computed at ingest
|
||||
(ADR-0017).
|
||||
"""
|
||||
|
||||
chunk: Chunk
|
||||
dense: dict[str, list[float]]
|
||||
sparse: SparseVector
|
||||
|
||||
Reference in New Issue
Block a user