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

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 == []