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:
@@ -17,10 +17,22 @@ from src.application.auth.errors import (
|
||||
TenantInactiveError,
|
||||
)
|
||||
from src.application.files.errors import FileTooLargeError, InvalidUploadError
|
||||
from src.application.ingestion.errors import DocumentParseError, UnsupportedSourceTypeError
|
||||
from src.application.ingestion.errors import (
|
||||
ChunkLimitExceededError,
|
||||
DocumentParseError,
|
||||
EmbedderError,
|
||||
IngestionAtCapacityError,
|
||||
IngestionTimeoutError,
|
||||
UnsupportedSourceTypeError,
|
||||
)
|
||||
|
||||
logger = structlog.get_logger(__name__)
|
||||
|
||||
# A fixed backoff hint, not a computed retry budget: ADR-0017 rejects a
|
||||
# request outright at capacity rather than queueing it, so there is no
|
||||
# in-process estimate of when a slot will free up to report instead.
|
||||
_CAPACITY_RETRY_AFTER_SECONDS = 1
|
||||
|
||||
# (exception type, status code, stable error code)
|
||||
_MAPPING: tuple[tuple[type[Exception], int, str], ...] = (
|
||||
(InvalidApiKeyError, status.HTTP_401_UNAUTHORIZED, "invalid_api_key"),
|
||||
@@ -30,6 +42,9 @@ _MAPPING: tuple[tuple[type[Exception], int, str], ...] = (
|
||||
(DocumentParseError, status.HTTP_400_BAD_REQUEST, "validation_error"),
|
||||
(UnsupportedSourceTypeError, status.HTTP_415_UNSUPPORTED_MEDIA_TYPE, "unsupported_media_type"),
|
||||
(FileTooLargeError, status.HTTP_413_CONTENT_TOO_LARGE, "payload_too_large"),
|
||||
(ChunkLimitExceededError, status.HTTP_413_CONTENT_TOO_LARGE, "payload_too_large"),
|
||||
(EmbedderError, status.HTTP_502_BAD_GATEWAY, "embedder_error"),
|
||||
(IngestionTimeoutError, status.HTTP_504_GATEWAY_TIMEOUT, "ingestion_timeout"),
|
||||
)
|
||||
|
||||
|
||||
@@ -66,6 +81,14 @@ def register_exception_handlers(app: FastAPI) -> None:
|
||||
|
||||
app.add_exception_handler(exc_type, _handler)
|
||||
|
||||
@app.exception_handler(IngestionAtCapacityError)
|
||||
def _capacity_handler(request: Request, exc: IngestionAtCapacityError) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
|
||||
content=_envelope("ingestion_at_capacity", str(exc), _request_id(request)),
|
||||
headers={"Retry-After": str(_CAPACITY_RETRY_AFTER_SECONDS)},
|
||||
)
|
||||
|
||||
@app.exception_handler(RequestValidationError)
|
||||
def _validation_handler(request: Request, exc: RequestValidationError) -> JSONResponse:
|
||||
return JSONResponse(
|
||||
|
||||
@@ -5,9 +5,10 @@ touch MinIO/Qdrant, or otherwise carry ingestion business logic (ADR-0015).
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from typing import Annotated
|
||||
|
||||
from anyio import CapacityLimiter
|
||||
from anyio import CapacityLimiter, Semaphore
|
||||
from fastapi import APIRouter, Depends, Form, HTTPException, Response, UploadFile, status
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
|
||||
@@ -16,12 +17,16 @@ from src.api.schemas.files import FileStatusResponse, FileUploadResponse
|
||||
from src.application.auth.context import AuthContext
|
||||
from src.application.files.status import get_file_status
|
||||
from src.application.files.upload import upload_source_file
|
||||
from src.application.ports.embedding import DenseEmbedder, SparseEmbedder
|
||||
from src.application.ports.object_storage import ObjectStorage
|
||||
from src.bootstrap.dependencies import (
|
||||
get_dense_embedders,
|
||||
get_ingestion_concurrency_limiter,
|
||||
get_ingestion_limiter,
|
||||
get_object_storage,
|
||||
get_sessionmaker,
|
||||
get_settings,
|
||||
get_sparse_embedder,
|
||||
)
|
||||
from src.config import Settings
|
||||
|
||||
@@ -32,6 +37,9 @@ _SessionmakerDep = Annotated[async_sessionmaker[AsyncSession], Depends(get_sessi
|
||||
_ObjectStorageDep = Annotated[ObjectStorage, Depends(get_object_storage)]
|
||||
_SettingsDep = Annotated[Settings, Depends(get_settings)]
|
||||
_IngestionLimiterDep = Annotated[CapacityLimiter, Depends(get_ingestion_limiter)]
|
||||
_ConcurrencyLimiterDep = Annotated[Semaphore, Depends(get_ingestion_concurrency_limiter)]
|
||||
_DenseEmbeddersDep = Annotated[Sequence[DenseEmbedder], Depends(get_dense_embedders)]
|
||||
_SparseEmbedderDep = Annotated[SparseEmbedder, Depends(get_sparse_embedder)]
|
||||
|
||||
|
||||
@router.post("", status_code=status.HTTP_201_CREATED)
|
||||
@@ -44,6 +52,9 @@ async def upload_file(
|
||||
storage: _ObjectStorageDep,
|
||||
settings: _SettingsDep,
|
||||
limiter: _IngestionLimiterDep,
|
||||
concurrency_limiter: _ConcurrencyLimiterDep,
|
||||
dense_embedders: _DenseEmbeddersDep,
|
||||
sparse_embedder: _SparseEmbedderDep,
|
||||
) -> FileUploadResponse:
|
||||
data = await file.read()
|
||||
result = await upload_source_file(
|
||||
@@ -53,9 +64,12 @@ async def upload_file(
|
||||
domain=domain,
|
||||
filename=file.filename or "",
|
||||
data=data,
|
||||
max_upload_size_bytes=settings.ingestion.max_upload_size_bytes,
|
||||
chunking_strategy=settings.chunking.strategy,
|
||||
validation_limiter=limiter,
|
||||
ingestion_settings=settings.ingestion,
|
||||
chunking_settings=settings.chunking,
|
||||
thread_limiter=limiter,
|
||||
concurrency_limiter=concurrency_limiter,
|
||||
dense_embedders=dense_embedders,
|
||||
sparse_embedder=sparse_embedder,
|
||||
)
|
||||
if not result.is_new_attempt:
|
||||
response.status_code = status.HTTP_200_OK
|
||||
|
||||
@@ -5,21 +5,27 @@ one request-scoped session, because the request is two units of work
|
||||
(ADR-0012, ADR-0017):
|
||||
|
||||
txn A (short): source_files [+ ingestion_jobs(status='running')], commit
|
||||
no txn: store bytes in MinIO
|
||||
no txn: store bytes in MinIO, parse/chunk (threads),
|
||||
embed dense+sparse (bounded/batched)
|
||||
txn B (short): ingestion_jobs -> succeeded/failed, append event, commit
|
||||
|
||||
No Postgres session is open during the MinIO write. A storage failure between
|
||||
No Postgres session is open during phase 2. A failure at any point between
|
||||
txn A and txn B still leaves a durable, inspectable `failed` job — never a
|
||||
job stuck in `running`.
|
||||
job stuck in `running`. The whole request additionally holds one of
|
||||
`INGESTION_MAX_CONCURRENCY` process-wide slots (`503` when exhausted) and
|
||||
phase 2 is bounded by `INGESTION_TIMEOUT_SECONDS` (`504`) (ADR-0017, plan 001
|
||||
Phase 4).
|
||||
|
||||
Parsing/chunking/Qdrant indexing are Phase 4/5 work, not implemented here:
|
||||
this phase stores bytes only, so a successful job reports `chunks_indexed=0`.
|
||||
Qdrant point upserts are Phase 5 work, not implemented here: this phase
|
||||
parses, chunks, and embeds, so a successful job still reports
|
||||
`chunks_indexed=0` — nothing is searchable yet.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
|
||||
import structlog
|
||||
from anyio import CapacityLimiter, to_thread
|
||||
from anyio import CapacityLimiter, Semaphore, fail_after, to_thread
|
||||
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
|
||||
|
||||
from src.application.auth.context import AuthContext
|
||||
@@ -27,7 +33,22 @@ from src.application.files.errors import InvalidUploadError
|
||||
from src.application.files.models import UploadResult
|
||||
from src.application.files.storage_keys import source_file_object_key
|
||||
from src.application.files.validation import validate_and_hash_upload
|
||||
from src.application.ingestion import (
|
||||
ChunkTooLargeError,
|
||||
DocumentParseError,
|
||||
UnsupportedSourceTypeError,
|
||||
parse_and_chunk_document,
|
||||
)
|
||||
from src.application.ingestion.bounds import acquire_ingestion_slot, enforce_chunk_limit
|
||||
from src.application.ingestion.embedding import embed_chunks
|
||||
from src.application.ingestion.errors import (
|
||||
ChunkLimitExceededError,
|
||||
EmbedderError,
|
||||
IngestionTimeoutError,
|
||||
)
|
||||
from src.application.ports.embedding import DenseEmbedder, SparseEmbedder
|
||||
from src.application.ports.object_storage import ObjectStorage
|
||||
from src.config import ChunkingSettings, IngestionSettings
|
||||
from src.infrastructure.postgres.repositories import ingestion_jobs as jobs_repo
|
||||
from src.infrastructure.postgres.repositories import source_files as source_files_repo
|
||||
|
||||
@@ -71,9 +92,12 @@ async def upload_source_file(
|
||||
domain: str,
|
||||
filename: str,
|
||||
data: bytes,
|
||||
max_upload_size_bytes: int,
|
||||
chunking_strategy: str,
|
||||
validation_limiter: CapacityLimiter,
|
||||
ingestion_settings: IngestionSettings,
|
||||
chunking_settings: ChunkingSettings,
|
||||
thread_limiter: CapacityLimiter,
|
||||
concurrency_limiter: Semaphore,
|
||||
dense_embedders: Sequence[DenseEmbedder],
|
||||
sparse_embedder: SparseEmbedder,
|
||||
) -> UploadResult:
|
||||
domain = domain.strip()
|
||||
if not domain:
|
||||
@@ -81,114 +105,186 @@ async def upload_source_file(
|
||||
|
||||
validated = await to_thread.run_sync(
|
||||
lambda: validate_and_hash_upload(
|
||||
filename=filename, data=data, max_size_bytes=max_upload_size_bytes
|
||||
filename=filename, data=data, max_size_bytes=ingestion_settings.max_upload_size_bytes
|
||||
),
|
||||
limiter=validation_limiter,
|
||||
limiter=thread_limiter,
|
||||
)
|
||||
|
||||
async with sessionmaker() as session:
|
||||
existing = await source_files_repo.find_active_by_content_hash(
|
||||
session,
|
||||
tenant_id=auth.tenant_id,
|
||||
domain=domain,
|
||||
content_sha256=validated.content_sha256,
|
||||
)
|
||||
|
||||
if existing is not None:
|
||||
latest_job = await jobs_repo.get_latest_for_source_file(
|
||||
session, tenant_id=auth.tenant_id, source_file_id=existing.id
|
||||
)
|
||||
if latest_job is not None and latest_job.status == "succeeded":
|
||||
logger.info(
|
||||
"files.upload.duplicate",
|
||||
tenant_id=str(auth.tenant_id),
|
||||
file_id=str(existing.id),
|
||||
)
|
||||
return UploadResult(
|
||||
file_id=existing.id,
|
||||
ingestion_job_id=latest_job.id,
|
||||
status=latest_job.status,
|
||||
chunks_indexed=latest_job.points_created,
|
||||
is_new_attempt=False,
|
||||
)
|
||||
source_file_id = existing.id
|
||||
object_key = existing.storage_uri or source_file_object_key(
|
||||
auth.tenant_id, source_file_id
|
||||
)
|
||||
else:
|
||||
source_file_id = uuid.uuid4()
|
||||
object_key = source_file_object_key(auth.tenant_id, source_file_id)
|
||||
source_files_repo.create(
|
||||
async with acquire_ingestion_slot(concurrency_limiter):
|
||||
async with sessionmaker() as session:
|
||||
existing = await source_files_repo.find_active_by_content_hash(
|
||||
session,
|
||||
source_file_id=source_file_id,
|
||||
tenant_id=auth.tenant_id,
|
||||
domain=domain,
|
||||
source_filename=filename,
|
||||
source_type=validated.source_type,
|
||||
content_sha256=validated.content_sha256,
|
||||
byte_size=len(data),
|
||||
storage_uri=object_key,
|
||||
created_by_api_key_id=auth.api_key_id,
|
||||
)
|
||||
# `ingestion_jobs.source_file_id` FKs to this row; flush so the
|
||||
# insert below sees it, since the two mapped classes carry no
|
||||
# ORM relationship for the unit of work to order by itself.
|
||||
await session.flush()
|
||||
|
||||
job = jobs_repo.create_running(
|
||||
session,
|
||||
tenant_id=auth.tenant_id,
|
||||
source_file_id=source_file_id,
|
||||
requested_by_api_key_id=auth.api_key_id,
|
||||
chunking_strategy=chunking_strategy,
|
||||
)
|
||||
jobs_repo.append_event(
|
||||
session,
|
||||
tenant_id=auth.tenant_id,
|
||||
ingestion_job_id=job.id,
|
||||
level="info",
|
||||
stage="received",
|
||||
message="upload accepted, storing object",
|
||||
)
|
||||
await session.commit()
|
||||
ingestion_job_id = job.id
|
||||
if existing is not None:
|
||||
latest_job = await jobs_repo.get_latest_for_source_file(
|
||||
session, tenant_id=auth.tenant_id, source_file_id=existing.id
|
||||
)
|
||||
if latest_job is not None and latest_job.status == "succeeded":
|
||||
logger.info(
|
||||
"files.upload.duplicate",
|
||||
tenant_id=str(auth.tenant_id),
|
||||
file_id=str(existing.id),
|
||||
)
|
||||
return UploadResult(
|
||||
file_id=existing.id,
|
||||
ingestion_job_id=latest_job.id,
|
||||
status=latest_job.status,
|
||||
chunks_indexed=latest_job.points_created,
|
||||
is_new_attempt=False,
|
||||
)
|
||||
source_file_id = existing.id
|
||||
object_key = existing.storage_uri or source_file_object_key(
|
||||
auth.tenant_id, source_file_id
|
||||
)
|
||||
else:
|
||||
source_file_id = uuid.uuid4()
|
||||
object_key = source_file_object_key(auth.tenant_id, source_file_id)
|
||||
source_files_repo.create(
|
||||
session,
|
||||
source_file_id=source_file_id,
|
||||
tenant_id=auth.tenant_id,
|
||||
domain=domain,
|
||||
source_filename=filename,
|
||||
source_type=validated.source_type,
|
||||
content_sha256=validated.content_sha256,
|
||||
byte_size=len(data),
|
||||
storage_uri=object_key,
|
||||
created_by_api_key_id=auth.api_key_id,
|
||||
)
|
||||
# `ingestion_jobs.source_file_id` FKs to this row; flush so the
|
||||
# insert below sees it, since the two mapped classes carry no
|
||||
# ORM relationship for the unit of work to order by itself.
|
||||
await session.flush()
|
||||
|
||||
# Phase 2: no Postgres session open across this work (ADR-0017).
|
||||
try:
|
||||
await storage.put_object(key=object_key, data=data, content_type=validated.content_type)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"files.upload.storage_failed",
|
||||
tenant_id=str(auth.tenant_id),
|
||||
file_id=str(source_file_id),
|
||||
ingestion_job_id=str(ingestion_job_id),
|
||||
)
|
||||
await _mark_job_failed(
|
||||
sessionmaker,
|
||||
tenant_id=auth.tenant_id,
|
||||
ingestion_job_id=ingestion_job_id,
|
||||
error_code="storage_upload_failed",
|
||||
error_message=f"failed to store object: {exc}",
|
||||
)
|
||||
raise
|
||||
job = jobs_repo.create_running(
|
||||
session,
|
||||
tenant_id=auth.tenant_id,
|
||||
source_file_id=source_file_id,
|
||||
requested_by_api_key_id=auth.api_key_id,
|
||||
chunking_strategy=chunking_settings.strategy,
|
||||
)
|
||||
jobs_repo.append_event(
|
||||
session,
|
||||
tenant_id=auth.tenant_id,
|
||||
ingestion_job_id=job.id,
|
||||
level="info",
|
||||
stage="received",
|
||||
message="upload accepted, storing object",
|
||||
)
|
||||
await session.commit()
|
||||
ingestion_job_id = job.id
|
||||
|
||||
async with sessionmaker() as session:
|
||||
await jobs_repo.mark_terminal(
|
||||
session,
|
||||
tenant_id=auth.tenant_id,
|
||||
ingestion_job_id=ingestion_job_id,
|
||||
status="succeeded",
|
||||
points_created=0,
|
||||
)
|
||||
jobs_repo.append_event(
|
||||
session,
|
||||
tenant_id=auth.tenant_id,
|
||||
ingestion_job_id=ingestion_job_id,
|
||||
level="info",
|
||||
stage="completed",
|
||||
message="object stored; parsing/embedding/indexing not yet implemented",
|
||||
)
|
||||
await session.commit()
|
||||
# Phase 2: no Postgres session open across this work (ADR-0017),
|
||||
# bounded end-to-end by INGESTION_TIMEOUT_SECONDS.
|
||||
try:
|
||||
with fail_after(ingestion_settings.timeout_seconds):
|
||||
try:
|
||||
await storage.put_object(
|
||||
key=object_key, data=data, content_type=validated.content_type
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"files.upload.storage_failed",
|
||||
tenant_id=str(auth.tenant_id),
|
||||
file_id=str(source_file_id),
|
||||
ingestion_job_id=str(ingestion_job_id),
|
||||
)
|
||||
await _mark_job_failed(
|
||||
sessionmaker,
|
||||
tenant_id=auth.tenant_id,
|
||||
ingestion_job_id=ingestion_job_id,
|
||||
error_code="storage_upload_failed",
|
||||
error_message=f"failed to store object: {exc}",
|
||||
)
|
||||
raise
|
||||
|
||||
try:
|
||||
chunks = await parse_and_chunk_document(
|
||||
data,
|
||||
source_type=validated.source_type,
|
||||
file_id=source_file_id,
|
||||
settings=chunking_settings,
|
||||
limiter=thread_limiter,
|
||||
)
|
||||
except (DocumentParseError, UnsupportedSourceTypeError, ChunkTooLargeError) as exc:
|
||||
await _mark_job_failed(
|
||||
sessionmaker,
|
||||
tenant_id=auth.tenant_id,
|
||||
ingestion_job_id=ingestion_job_id,
|
||||
error_code="parse_failed",
|
||||
error_message=str(exc),
|
||||
)
|
||||
raise
|
||||
|
||||
try:
|
||||
enforce_chunk_limit(chunks, max_chunks=ingestion_settings.max_chunks_per_file)
|
||||
except ChunkLimitExceededError as exc:
|
||||
await _mark_job_failed(
|
||||
sessionmaker,
|
||||
tenant_id=auth.tenant_id,
|
||||
ingestion_job_id=ingestion_job_id,
|
||||
error_code="chunk_limit_exceeded",
|
||||
error_message=str(exc),
|
||||
)
|
||||
raise
|
||||
|
||||
try:
|
||||
embedded = await embed_chunks(
|
||||
chunks,
|
||||
dense_embedders=dense_embedders,
|
||||
sparse_embedder=sparse_embedder,
|
||||
settings=ingestion_settings,
|
||||
thread_limiter=thread_limiter,
|
||||
)
|
||||
except EmbedderError as exc:
|
||||
await _mark_job_failed(
|
||||
sessionmaker,
|
||||
tenant_id=auth.tenant_id,
|
||||
ingestion_job_id=ingestion_job_id,
|
||||
error_code="embedding_failed",
|
||||
error_message=str(exc),
|
||||
)
|
||||
raise
|
||||
except TimeoutError:
|
||||
logger.warning(
|
||||
"files.upload.timeout",
|
||||
tenant_id=str(auth.tenant_id),
|
||||
file_id=str(source_file_id),
|
||||
ingestion_job_id=str(ingestion_job_id),
|
||||
)
|
||||
await _mark_job_failed(
|
||||
sessionmaker,
|
||||
tenant_id=auth.tenant_id,
|
||||
ingestion_job_id=ingestion_job_id,
|
||||
error_code="timeout",
|
||||
error_message=f"ingestion exceeded {ingestion_settings.timeout_seconds}s",
|
||||
)
|
||||
raise IngestionTimeoutError(
|
||||
f"ingestion exceeded {ingestion_settings.timeout_seconds}s"
|
||||
) from None
|
||||
|
||||
async with sessionmaker() as session:
|
||||
await jobs_repo.mark_terminal(
|
||||
session,
|
||||
tenant_id=auth.tenant_id,
|
||||
ingestion_job_id=ingestion_job_id,
|
||||
status="succeeded",
|
||||
points_created=0,
|
||||
)
|
||||
jobs_repo.append_event(
|
||||
session,
|
||||
tenant_id=auth.tenant_id,
|
||||
ingestion_job_id=ingestion_job_id,
|
||||
level="info",
|
||||
stage="completed",
|
||||
message="chunks parsed and embedded; Qdrant indexing not yet implemented",
|
||||
details={"chunks_parsed": len(chunks), "chunks_embedded": len(embedded)},
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
logger.info(
|
||||
"files.upload.succeeded",
|
||||
|
||||
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
|
||||
|
||||
49
src/application/ports/embedding.py
Normal file
49
src/application/ports/embedding.py
Normal file
@@ -0,0 +1,49 @@
|
||||
"""Embedding ports (ADR-0001, ADR-0017).
|
||||
|
||||
`src/infrastructure/embedding/` holds the production adapters; tests use
|
||||
scripted fakes (ADR-0016). Application code depends on these Protocols, not
|
||||
on `httpx`/provider SDKs directly.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Protocol
|
||||
|
||||
from src.application.ingestion.models import SparseVector
|
||||
|
||||
|
||||
class DenseEmbedder(Protocol):
|
||||
"""One named dense vector's embedding client (`dense_nomic`/`dense_openai`).
|
||||
|
||||
`embed_batch` is a single batched network call — callers own concurrency
|
||||
bounding (ADR-0017's `embed_concurrency` semaphore), not this Protocol.
|
||||
"""
|
||||
|
||||
name: str
|
||||
|
||||
async def embed_batch(self, texts: Sequence[str]) -> list[list[float]]:
|
||||
"""Return one vector per input text, same order. Raises `EmbedderError`
|
||||
(see `src/application/ingestion/errors.py`) on transport/response
|
||||
failure.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
class SparseEmbedder(Protocol):
|
||||
"""The `sparse` (BM25) vector's embedding client.
|
||||
|
||||
Blocking/CPU-bound (ADR-0017): callers offload it via
|
||||
`anyio.to_thread.run_sync` with the ingestion `CapacityLimiter`, not call
|
||||
it directly from an `async def`.
|
||||
"""
|
||||
|
||||
name: str
|
||||
|
||||
def embed_batch(self, texts: Sequence[str], *, query: bool = False) -> list[SparseVector]:
|
||||
"""Return one sparse vector per input text, same order.
|
||||
|
||||
`query=True` selects the query-side weighting, which omits document
|
||||
length normalization. Ingestion always passes `False`; the flag exists
|
||||
so retrieval (ADR-0003) encodes queries through this same port rather
|
||||
than growing a second, silently divergent implementation.
|
||||
"""
|
||||
...
|
||||
@@ -1,12 +1,13 @@
|
||||
from collections.abc import AsyncIterator
|
||||
from collections.abc import AsyncIterator, Sequence
|
||||
from dataclasses import dataclass
|
||||
|
||||
from anyio import CapacityLimiter
|
||||
from anyio import CapacityLimiter, Semaphore
|
||||
from fastapi import Request
|
||||
from minio import Minio
|
||||
from qdrant_client import AsyncQdrantClient
|
||||
from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession, async_sessionmaker
|
||||
|
||||
from src.application.ports.embedding import DenseEmbedder, SparseEmbedder
|
||||
from src.application.ports.object_storage import ObjectStorage
|
||||
from src.config import Settings
|
||||
|
||||
@@ -20,6 +21,9 @@ class AppResources:
|
||||
qdrant_client: AsyncQdrantClient
|
||||
object_storage: ObjectStorage
|
||||
ingestion_limiter: CapacityLimiter
|
||||
dense_embedders: Sequence[DenseEmbedder]
|
||||
sparse_embedder: SparseEmbedder
|
||||
ingestion_concurrency_limiter: Semaphore
|
||||
|
||||
|
||||
def _resources(request: Request) -> AppResources:
|
||||
@@ -46,6 +50,18 @@ def get_ingestion_limiter(request: Request) -> CapacityLimiter:
|
||||
return _resources(request).ingestion_limiter
|
||||
|
||||
|
||||
def get_dense_embedders(request: Request) -> Sequence[DenseEmbedder]:
|
||||
return _resources(request).dense_embedders
|
||||
|
||||
|
||||
def get_sparse_embedder(request: Request) -> SparseEmbedder:
|
||||
return _resources(request).sparse_embedder
|
||||
|
||||
|
||||
def get_ingestion_concurrency_limiter(request: Request) -> Semaphore:
|
||||
return _resources(request).ingestion_concurrency_limiter
|
||||
|
||||
|
||||
def get_sessionmaker(request: Request) -> async_sessionmaker[AsyncSession]:
|
||||
"""The session *factory*, not a request-scoped session.
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -87,6 +87,79 @@ class QdrantSettings(BaseSettings):
|
||||
api_key: str | None = None
|
||||
|
||||
|
||||
class NomicEmbeddingSettings(BaseSettings):
|
||||
"""Self-hosted `nomic-embed-text-v2-moe` server (ADR-0001, 768-dim).
|
||||
|
||||
Talks the OpenAI-compatible `/embeddings` endpoint shape. The default
|
||||
points at the Ollama host the `emet` benchmark used, whose OpenAI-compat
|
||||
shim accepts any non-empty API key.
|
||||
|
||||
`document_prefix` is empty by design. ADR-0004 cites the model card's
|
||||
`search_document: ` requirement, but emet's winning run used no prefix
|
||||
(Ollama's template is a bare passthrough and injects none), and the
|
||||
prefix is not cosmetic -- it moves the vector substantially. Turning it
|
||||
on here obliges the query side to send `search_query: ` too (ADR-0003),
|
||||
so it stays off until an emet run measures the pair together.
|
||||
"""
|
||||
|
||||
model_config = SettingsConfigDict(env_prefix="EMBEDDING_NOMIC_", extra="ignore")
|
||||
|
||||
base_url: str = "http://192.168.10.10:11435/v1"
|
||||
model: str = "nomic-embed-text-v2-moe"
|
||||
api_key: str = "sk-not-set"
|
||||
document_prefix: str = ""
|
||||
keep_alive: str = "30m"
|
||||
timeout_seconds: float = 30.0
|
||||
|
||||
|
||||
class OpenaiEmbeddingSettings(BaseSettings):
|
||||
"""OpenAI's hosted embedding API (ADR-0001's second dense signal).
|
||||
|
||||
`dimensions` is unset, so `text-embedding-3-large` returns its native
|
||||
3072 dimensions -- the configuration emet benchmarked. Setting it would
|
||||
truncate via Matryoshka and is a re-embedding migration, not a config
|
||||
tweak.
|
||||
"""
|
||||
|
||||
model_config = SettingsConfigDict(env_prefix="EMBEDDING_OPENAI_", extra="ignore")
|
||||
|
||||
base_url: str = "https://api.openai.com/v1"
|
||||
model: str = "text-embedding-3-large"
|
||||
api_key: str | None = None
|
||||
dimensions: int | None = None
|
||||
document_prefix: str = ""
|
||||
timeout_seconds: float = 30.0
|
||||
|
||||
|
||||
class SparseEmbeddingSettings(BaseSettings):
|
||||
"""BM25 sparse-vector parameters (ADR-0001, ADR-0005).
|
||||
|
||||
`analyzer`, `k`, and `b` are the configuration emet benchmarked as
|
||||
`bm25-fa-norm-stop`; changing them invalidates that result. `k`/`b`
|
||||
saturation is applied client-side here, while IDF comes from Qdrant's
|
||||
`modifier="idf"` sparse-vector config at query time.
|
||||
|
||||
`avg_len` is emet's placeholder average document length in analyzer
|
||||
tokens, exposed as a setting so it can be recalibrated from real corpus
|
||||
statistics without a code change.
|
||||
"""
|
||||
|
||||
model_config = SettingsConfigDict(env_prefix="EMBEDDING_SPARSE_", extra="ignore")
|
||||
|
||||
analyzer: str = "fa_norm_stop"
|
||||
k: float = 1.2
|
||||
b: float = 0.75
|
||||
avg_len: float = 256.0
|
||||
|
||||
|
||||
class EmbeddingSettings(BaseSettings):
|
||||
model_config = SettingsConfigDict(extra="ignore")
|
||||
|
||||
nomic: NomicEmbeddingSettings = Field(default_factory=NomicEmbeddingSettings)
|
||||
openai: OpenaiEmbeddingSettings = Field(default_factory=OpenaiEmbeddingSettings)
|
||||
sparse: SparseEmbeddingSettings = Field(default_factory=SparseEmbeddingSettings)
|
||||
|
||||
|
||||
class AppLimitSettings(BaseSettings):
|
||||
model_config = SettingsConfigDict(env_prefix="APP_", extra="ignore")
|
||||
|
||||
@@ -109,5 +182,6 @@ class Settings(BaseSettings):
|
||||
ingestion: IngestionSettings = Field(default_factory=IngestionSettings)
|
||||
chunking: ChunkingSettings = Field(default_factory=ChunkingSettings)
|
||||
qdrant: QdrantSettings = Field(default_factory=QdrantSettings)
|
||||
embedding: EmbeddingSettings = Field(default_factory=EmbeddingSettings)
|
||||
app: AppLimitSettings = Field(default_factory=AppLimitSettings)
|
||||
logging: LoggingSettings = Field(default_factory=LoggingSettings)
|
||||
|
||||
0
src/infrastructure/embedding/__init__.py
Normal file
0
src/infrastructure/embedding/__init__.py
Normal file
150
src/infrastructure/embedding/analyzers.py
Normal file
150
src/infrastructure/embedding/analyzers.py
Normal file
@@ -0,0 +1,150 @@
|
||||
"""The `fa_norm_stop` BM25 analyzer (ADR-0001, ADR-0005).
|
||||
|
||||
Ported from the `emet` evaluation lab
|
||||
(`src/chatbot_gh/adapters/sparse/analyzers.py`), which benchmarked four Farsi
|
||||
analyzer variants on the real corpus and found `fa_norm_stop` the best
|
||||
performer. This is a *measured* artifact: changing the normalization,
|
||||
tokenization, or stopword list invalidates that result, so improvements belong
|
||||
in a new emet benchmark run rather than in an edit here.
|
||||
|
||||
Deliberately independent of `src/application/ingestion/normalization.py`.
|
||||
Those solve different problems: `normalize_persian_text` shapes chunk content
|
||||
that gets cited back to the reader, so ADR-0018 has it preserve digits and
|
||||
punctuation as authored. This module shapes index terms nobody ever sees, so
|
||||
it folds digits and diacritics freely. Sharing one function between them would
|
||||
let a display-motivated tweak silently perturb the benchmarked sparse index.
|
||||
"""
|
||||
|
||||
import re
|
||||
import unicodedata
|
||||
|
||||
# Several Arabic letterforms are visually indistinguishable from Latin ones in
|
||||
# a monospace editor (alef from "l", heh from "o"), and literals render
|
||||
# right-to-left, visually reordering the source line. `normalization.py` writes
|
||||
# them as codepoints for that reason; this module follows the same convention.
|
||||
_ZWNJ = 0x200C
|
||||
_ARABIC_YEH = 0x064A
|
||||
_ARABIC_KAF = 0x0643
|
||||
_TEH_MARBUTA = 0x0629
|
||||
_HAMZA_ON_WAW = 0x0624
|
||||
_ALEF_HAMZA_BELOW = 0x0625
|
||||
_ALEF_HAMZA_ABOVE = 0x0623
|
||||
|
||||
_PERSIAN_YEH = 0x06CC
|
||||
_PERSIAN_KEHEH = 0x06A9
|
||||
_HEH = 0x0647
|
||||
_WAW = 0x0648
|
||||
_ALEF = 0x0627
|
||||
_SPACE = 0x0020
|
||||
|
||||
# Persian (U+06F0-U+06F9) and Arabic-Indic (U+0660-U+0669) digits both fold to
|
||||
# ASCII, so the same number matches however it was authored.
|
||||
_EASTERN_DIGITS = str.maketrans("۰۱۲۳۴۵۶۷۸۹٠١٢٣٤٥٦٧٨٩", "01234567890123456789")
|
||||
|
||||
# ZWNJ becomes a space (splitting compounds into separate terms) and the
|
||||
# Arabic letterforms fold to their Persian equivalents. Every entry is a
|
||||
# single codepoint mapping to a single codepoint over disjoint sources, so
|
||||
# applying them in one pass is equivalent to emet's chained `str.replace`
|
||||
# calls -- provided NFC runs first, since NFC is what composes the hamza
|
||||
# forms this table then folds.
|
||||
_FOLDING: dict[int, int] = {
|
||||
_ZWNJ: _SPACE,
|
||||
_ARABIC_YEH: _PERSIAN_YEH,
|
||||
_ARABIC_KAF: _PERSIAN_KEHEH,
|
||||
_TEH_MARBUTA: _HEH,
|
||||
_HAMZA_ON_WAW: _WAW,
|
||||
_ALEF_HAMZA_BELOW: _ALEF,
|
||||
_ALEF_HAMZA_ABOVE: _ALEF,
|
||||
}
|
||||
_FOLDING.update(_EASTERN_DIGITS)
|
||||
|
||||
# Common Persian/Arabic stopwords (function words + FAQ noise), plus the
|
||||
# English function words that appear in a mixed-script corpus. Kept small and
|
||||
# explicit -- deliberately not a full hazm list. `_HEH_ALEF` is the plural
|
||||
# suffix "ha"; written as a codepoint pair because both of its letters are
|
||||
# Latin-confusable, which is exactly the case Ruff's RUF001 flags.
|
||||
_HEH_ALEF = chr(_HEH) + chr(_ALEF)
|
||||
|
||||
_PERSIAN_STOPWORDS: frozenset[str] = frozenset(
|
||||
{
|
||||
"و",
|
||||
"در",
|
||||
"به",
|
||||
"از",
|
||||
"که",
|
||||
"این",
|
||||
"را",
|
||||
"با",
|
||||
"برای",
|
||||
"آن",
|
||||
"یک",
|
||||
"است",
|
||||
"شد",
|
||||
"شده",
|
||||
"می",
|
||||
"های",
|
||||
_HEH_ALEF,
|
||||
"یا",
|
||||
"تا",
|
||||
"بر",
|
||||
"اگر",
|
||||
"هم",
|
||||
"نیز",
|
||||
"ولی",
|
||||
"اما",
|
||||
"چه",
|
||||
"چون",
|
||||
"روی",
|
||||
"پس",
|
||||
"پیش",
|
||||
"هر",
|
||||
"هیچ",
|
||||
"بود",
|
||||
"باشد",
|
||||
"هست",
|
||||
"نیست",
|
||||
"کند",
|
||||
"کرد",
|
||||
"کردن",
|
||||
"شود",
|
||||
"the",
|
||||
"a",
|
||||
"an",
|
||||
"of",
|
||||
"to",
|
||||
"and",
|
||||
"in",
|
||||
"on",
|
||||
"for",
|
||||
"is",
|
||||
"are",
|
||||
}
|
||||
)
|
||||
|
||||
# Word characters minus underscore. Note this KEEPS digits: an insurance
|
||||
# corpus is full of policy numbers, dates, and amounts, and those are exactly
|
||||
# the tokens a lexical index should be able to match on.
|
||||
_TOKEN_RE = re.compile(r"[^\W_]+", re.UNICODE)
|
||||
|
||||
FA_NORM_STOP = "fa_norm_stop"
|
||||
|
||||
|
||||
def _normalize_fa(text: str) -> str:
|
||||
return unicodedata.normalize("NFC", text).translate(_FOLDING)
|
||||
|
||||
|
||||
def _tokenize_raw(text: str) -> list[str]:
|
||||
return [m.group(0).lower() for m in _TOKEN_RE.finditer(text)]
|
||||
|
||||
|
||||
def analyze(text: str, analyzer: str = FA_NORM_STOP) -> list[str]:
|
||||
"""Tokenize `text` into sparse-index terms.
|
||||
|
||||
Only `fa_norm_stop` is implemented -- emet's other three variants
|
||||
(`raw`, `fa_norm`, `fa_norm_stem`) lost the benchmark and exist there as
|
||||
experiment arms, not as configurations this service should run.
|
||||
"""
|
||||
if analyzer != FA_NORM_STOP:
|
||||
raise ValueError(f"Unknown analyzer '{analyzer}'")
|
||||
tokens = _tokenize_raw(_normalize_fa(text))
|
||||
return [token for token in tokens if token not in _PERSIAN_STOPWORDS]
|
||||
92
src/infrastructure/embedding/bm25.py
Normal file
92
src/infrastructure/embedding/bm25.py
Normal file
@@ -0,0 +1,92 @@
|
||||
"""The `bm25-fa-norm-stop` sparse embedder (ADR-0001, ADR-0005).
|
||||
|
||||
Ported from the `emet` evaluation lab
|
||||
(`src/chatbot_gh/adapters/sparse/bm25_embedder.py`), the configuration that
|
||||
won its Farsi analyzer benchmark. Not Qdrant's hosted `Qdrant/bm25` FastEmbed
|
||||
model, whose documented language support omits Farsi (ADR-0005).
|
||||
|
||||
**The BM25 work is split across two systems.** This adapter applies the
|
||||
term-frequency saturation half client-side -- the `k` and `b` parameters,
|
||||
including document-length normalization. IDF is *not* computed here: Qdrant
|
||||
supplies it from collection-wide statistics when the sparse vector field is
|
||||
created with `modifier="idf"`.
|
||||
|
||||
That split is load-bearing. A Qdrant collection created without
|
||||
`modifier="idf"` will silently score these vectors as saturated term
|
||||
frequencies with no IDF weighting at all -- no error, just materially worse
|
||||
lexical retrieval. The collection bootstrap (plan 001 Phase 5) must set it.
|
||||
"""
|
||||
|
||||
from collections import Counter
|
||||
from collections.abc import Sequence
|
||||
from hashlib import blake2b
|
||||
|
||||
from src.application.ingestion.models import SparseVector
|
||||
from src.config import SparseEmbeddingSettings
|
||||
from src.infrastructure.embedding.analyzers import analyze
|
||||
|
||||
# Qdrant sparse indices must be non-negative and fit a signed 32-bit int.
|
||||
_INDEX_SPACE = 2**31 - 1
|
||||
|
||||
|
||||
def _token_index(token: str) -> int:
|
||||
"""Map a term to its sparse-vector index.
|
||||
|
||||
A hash rather than a vocabulary table, so the mapping needs no shared
|
||||
state and stays identical across processes, restarts, and — critically —
|
||||
between ingest-time and query-time encoding. `blake2b` rather than
|
||||
`hash()`, which is PYTHONHASHSEED-salted and therefore differs per
|
||||
process.
|
||||
"""
|
||||
digest = blake2b(token.encode("utf-8"), digest_size=8).digest()
|
||||
return int.from_bytes(digest, "big") % _INDEX_SPACE
|
||||
|
||||
|
||||
def text_to_sparse_vector(
|
||||
text: str, *, settings: SparseEmbeddingSettings, query: bool = False
|
||||
) -> SparseVector:
|
||||
"""Encode one text as a BM25-saturated sparse vector (IDF applied by Qdrant).
|
||||
|
||||
Document and query sides differ in exactly one term: documents carry the
|
||||
`b` length normalization, queries do not (standard BM25 practice — a
|
||||
query's own length should not discount its terms).
|
||||
"""
|
||||
tokens = analyze(text, settings.analyzer)
|
||||
if not tokens:
|
||||
return SparseVector(indices=[], values=[])
|
||||
|
||||
frequencies = Counter(tokens)
|
||||
doc_length = float(len(tokens))
|
||||
k = settings.k
|
||||
b = settings.b
|
||||
|
||||
indices: list[int] = []
|
||||
values: list[float] = []
|
||||
# Sorted so the emitted vector is deterministic for a given text, which
|
||||
# keeps re-ingestion byte-stable and makes the output testable.
|
||||
for token, freq in sorted(frequencies.items()):
|
||||
if query:
|
||||
weight = freq * (k + 1.0) / (freq + k)
|
||||
else:
|
||||
weight = freq * (k + 1.0) / (freq + k * (1.0 - b + b * doc_length / settings.avg_len))
|
||||
indices.append(_token_index(token))
|
||||
values.append(float(weight))
|
||||
|
||||
return SparseVector(indices=indices, values=values)
|
||||
|
||||
|
||||
class Bm25SparseEmbedder:
|
||||
"""A `SparseEmbedder` (see `src/application/ports/embedding.py`).
|
||||
|
||||
Pure CPU work with no network calls, so it is blocking: callers offload it
|
||||
via `anyio.to_thread.run_sync` with the ingestion `CapacityLimiter`
|
||||
(ADR-0017), never awaiting it directly on the event loop.
|
||||
"""
|
||||
|
||||
name = "sparse"
|
||||
|
||||
def __init__(self, settings: SparseEmbeddingSettings) -> None:
|
||||
self._settings = settings
|
||||
|
||||
def embed_batch(self, texts: Sequence[str], *, query: bool = False) -> list[SparseVector]:
|
||||
return [text_to_sparse_vector(text, settings=self._settings, query=query) for text in texts]
|
||||
81
src/infrastructure/embedding/openai_compatible.py
Normal file
81
src/infrastructure/embedding/openai_compatible.py
Normal file
@@ -0,0 +1,81 @@
|
||||
"""Dense embedding adapter for OpenAI-compatible `/embeddings` endpoints.
|
||||
|
||||
Backs both `dense_nomic` (self-hosted `nomic-embed-text-v2-moe` behind
|
||||
Ollama's OpenAI-compatible shim) and `dense_openai` (OpenAI's hosted API) —
|
||||
both speak the same request/response shape, so one adapter serves both named
|
||||
vectors with different config (ADR-0001).
|
||||
|
||||
Uses `httpx` directly rather than the `openai` SDK. The `emet` benchmark this
|
||||
configuration comes from uses the SDK, but it is a synchronous batch tool;
|
||||
ADR-0017 requires async, semaphore-bounded batches here, and the request shape
|
||||
is small enough that the SDK earns nothing.
|
||||
|
||||
`httpx.AsyncClient` is application-lifetime (ADR-0012): built once in the
|
||||
FastAPI lifespan and passed in, never constructed per call.
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
# Ollama's default port, plus the alternate the benchmarked deployment uses.
|
||||
_OLLAMA_PORTS = frozenset({11434, 11435})
|
||||
|
||||
|
||||
def is_ollama_base_url(base_url: str) -> bool:
|
||||
"""Whether `base_url` looks like an Ollama OpenAI-compatible endpoint.
|
||||
|
||||
Ollama unloads an idle model, and reloading `nomic-embed-text-v2-moe`
|
||||
costs well over two minutes — longer than `INGESTION_TIMEOUT_SECONDS`, so
|
||||
a cold upload would 504. `keep_alive` is how the model is kept resident,
|
||||
and it is an Ollama extension, hence the sniffing.
|
||||
"""
|
||||
parsed = urlparse(base_url)
|
||||
host = parsed.hostname or ""
|
||||
return parsed.port in _OLLAMA_PORTS or "ollama" in host.lower()
|
||||
|
||||
|
||||
class OpenAICompatibleEmbedder:
|
||||
"""A `DenseEmbedder` (see `src/application/ports/embedding.py`) over one
|
||||
OpenAI-compatible `/embeddings` endpoint.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: httpx.AsyncClient,
|
||||
*,
|
||||
name: str,
|
||||
model: str,
|
||||
dimensions: int | None = None,
|
||||
document_prefix: str = "",
|
||||
keep_alive: str | None = None,
|
||||
) -> None:
|
||||
self.name = name
|
||||
self._client = client
|
||||
self._model = model
|
||||
self._dimensions = dimensions
|
||||
self._document_prefix = document_prefix
|
||||
self._keep_alive = keep_alive
|
||||
|
||||
async def embed_batch(self, texts: Sequence[str]) -> list[list[float]]:
|
||||
inputs = (
|
||||
[f"{self._document_prefix}{text}" for text in texts]
|
||||
if self._document_prefix
|
||||
else list(texts)
|
||||
)
|
||||
payload: dict[str, object] = {"model": self._model, "input": inputs}
|
||||
if self._dimensions is not None:
|
||||
payload["dimensions"] = self._dimensions
|
||||
if self._keep_alive is not None:
|
||||
payload["keep_alive"] = self._keep_alive
|
||||
|
||||
response = await self._client.post("/embeddings", json=payload)
|
||||
response.raise_for_status()
|
||||
body = response.json()
|
||||
|
||||
# Sort by `index` rather than trusting response order: the contract
|
||||
# guarantees the field, not the ordering, and a silently permuted
|
||||
# batch would attach every vector to the wrong chunk.
|
||||
data = sorted(body["data"], key=lambda item: item["index"])
|
||||
return [item["embedding"] for item in data]
|
||||
Reference in New Issue
Block a user