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:
@@ -1,13 +1,20 @@
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from collections.abc import AsyncIterator, Callable, Sequence
|
||||
from contextlib import AbstractAsyncContextManager, asynccontextmanager
|
||||
|
||||
import httpx
|
||||
import structlog
|
||||
from anyio import CapacityLimiter, to_thread
|
||||
from anyio import CapacityLimiter, Semaphore, to_thread
|
||||
from fastapi import FastAPI
|
||||
|
||||
from src.application.ingestion import get_encoder
|
||||
from src.application.ports.embedding import DenseEmbedder
|
||||
from src.bootstrap.dependencies import AppResources
|
||||
from src.config import Settings
|
||||
from src.infrastructure.embedding.bm25 import Bm25SparseEmbedder
|
||||
from src.infrastructure.embedding.openai_compatible import (
|
||||
OpenAICompatibleEmbedder,
|
||||
is_ollama_base_url,
|
||||
)
|
||||
from src.infrastructure.minio.client import create_client as create_minio_client
|
||||
from src.infrastructure.minio.storage import MinioObjectStorage
|
||||
from src.infrastructure.observability.logging import configure_logging
|
||||
@@ -17,6 +24,34 @@ from src.infrastructure.qdrant.client import create_client as create_qdrant_clie
|
||||
logger = structlog.get_logger(__name__)
|
||||
|
||||
|
||||
def _auth_headers(api_key: str | None) -> dict[str, str]:
|
||||
"""Bearer header, or none at all when no key is configured.
|
||||
|
||||
Sending an empty `Bearer ` is worse than sending nothing: some gateways
|
||||
treat a malformed credential as an auth failure rather than as anonymous.
|
||||
"""
|
||||
return {"Authorization": f"Bearer {api_key}"} if api_key else {}
|
||||
|
||||
|
||||
async def _warm_dense_embedders(embedders: Sequence[DenseEmbedder]) -> None:
|
||||
"""Force each dense model to load before the first upload needs it.
|
||||
|
||||
Same rationale as the tiktoken warm-up above, but with the opposite
|
||||
failure policy. A self-hosted embedder that has unloaded the model takes
|
||||
minutes to serve its first request — longer than
|
||||
`INGESTION_TIMEOUT_SECONDS` — so paying that once at boot keeps it off a
|
||||
user's upload. Unlike the tokenizer this is best-effort: an embedder that
|
||||
is merely *down* must not stop the process from booting and reporting its
|
||||
own health, and `/readyz` is where that condition belongs.
|
||||
"""
|
||||
for embedder in embedders:
|
||||
try:
|
||||
await embedder.embed_batch(["warmup"])
|
||||
logger.info("lifespan.embedder.warmed", embedder=embedder.name)
|
||||
except Exception:
|
||||
logger.warning("lifespan.embedder.warm_failed", embedder=embedder.name, exc_info=True)
|
||||
|
||||
|
||||
def create_lifespan(
|
||||
settings: Settings | None = None,
|
||||
) -> Callable[[FastAPI], AbstractAsyncContextManager[None, bool | None]]:
|
||||
@@ -44,6 +79,48 @@ def create_lifespan(
|
||||
qdrant_client = create_qdrant_client(resolved_settings.qdrant)
|
||||
logger.info("lifespan.qdrant.client.created")
|
||||
|
||||
nomic_settings = resolved_settings.embedding.nomic
|
||||
nomic_http_client = httpx.AsyncClient(
|
||||
base_url=nomic_settings.base_url,
|
||||
timeout=nomic_settings.timeout_seconds,
|
||||
headers=_auth_headers(nomic_settings.api_key),
|
||||
)
|
||||
openai_settings = resolved_settings.embedding.openai
|
||||
openai_http_client = httpx.AsyncClient(
|
||||
base_url=openai_settings.base_url,
|
||||
timeout=openai_settings.timeout_seconds,
|
||||
headers=_auth_headers(openai_settings.api_key),
|
||||
)
|
||||
dense_embedders = (
|
||||
OpenAICompatibleEmbedder(
|
||||
nomic_http_client,
|
||||
name="dense_nomic",
|
||||
model=nomic_settings.model,
|
||||
document_prefix=nomic_settings.document_prefix,
|
||||
keep_alive=(
|
||||
nomic_settings.keep_alive
|
||||
if is_ollama_base_url(nomic_settings.base_url)
|
||||
else None
|
||||
),
|
||||
),
|
||||
OpenAICompatibleEmbedder(
|
||||
openai_http_client,
|
||||
name="dense_openai",
|
||||
model=openai_settings.model,
|
||||
dimensions=openai_settings.dimensions,
|
||||
document_prefix=openai_settings.document_prefix,
|
||||
),
|
||||
)
|
||||
sparse_embedder = Bm25SparseEmbedder(resolved_settings.embedding.sparse)
|
||||
logger.info("lifespan.embedders.created")
|
||||
|
||||
await _warm_dense_embedders(dense_embedders)
|
||||
|
||||
# Bounds how many ingestions run in this process at once (ADR-0017);
|
||||
# a distinct resource from ingestion_limiter, which bounds threads
|
||||
# spent on blocking work within a single ingestion.
|
||||
ingestion_concurrency_limiter = Semaphore(resolved_settings.ingestion.max_concurrency)
|
||||
|
||||
# Bounds threads spent on blocking ingestion work (parsing, chunking,
|
||||
# hashing, the sync minio SDK) so it cannot exhaust Starlette's own
|
||||
# thread pool (ADR-0017).
|
||||
@@ -60,6 +137,9 @@ def create_lifespan(
|
||||
qdrant_client=qdrant_client,
|
||||
object_storage=object_storage,
|
||||
ingestion_limiter=ingestion_limiter,
|
||||
dense_embedders=dense_embedders,
|
||||
sparse_embedder=sparse_embedder,
|
||||
ingestion_concurrency_limiter=ingestion_concurrency_limiter,
|
||||
)
|
||||
|
||||
try:
|
||||
@@ -75,4 +155,14 @@ def create_lifespan(
|
||||
except Exception:
|
||||
logger.exception("lifespan.qdrant.close.failed")
|
||||
|
||||
try:
|
||||
await nomic_http_client.aclose()
|
||||
except Exception:
|
||||
logger.exception("lifespan.embedding.nomic_client.close.failed")
|
||||
|
||||
try:
|
||||
await openai_http_client.aclose()
|
||||
except Exception:
|
||||
logger.exception("lifespan.embedding.openai_client.close.failed")
|
||||
|
||||
return lifespan
|
||||
|
||||
Reference in New Issue
Block a user