Ali Zarinkolah
5c0a5938f8
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>
241 lines
8.5 KiB
Python
241 lines
8.5 KiB
Python
"""ADR-0017's two-transaction upload shape, exercised against real Postgres.
|
|
|
|
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, 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 FakeDenseEmbedder, FakeObjectStorage, FakeSparseEmbedder
|
|
from tests.support.factories import create_api_key, create_tenant
|
|
|
|
pytestmark = [
|
|
pytest.mark.integration,
|
|
pytest.mark.postgres,
|
|
pytest.mark.asyncio(loop_scope="session"),
|
|
]
|
|
|
|
_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)
|
|
await db_session.commit()
|
|
return AuthContext(
|
|
tenant_id=tenant.id,
|
|
tenant_slug=tenant.slug,
|
|
api_key_id=api_key.id,
|
|
scopes=frozenset({"files:write"}),
|
|
actor_type="backend",
|
|
)
|
|
|
|
|
|
async def test_upload_source_file_commits_running_job_before_storage_write(
|
|
db_session: AsyncSession, db_sessionmaker: async_sessionmaker[AsyncSession]
|
|
) -> None:
|
|
"""After phase 1 commits, a *separate* session must already see the
|
|
running job -- proving txn A committed and released before phase 2 work,
|
|
per ADR-0017.
|
|
"""
|
|
auth = await _auth_for(db_session)
|
|
storage = FakeObjectStorage()
|
|
|
|
result = await _upload(sessionmaker=db_sessionmaker, storage=storage, auth=auth)
|
|
|
|
assert result.status == "succeeded"
|
|
assert result.chunks_indexed == 0
|
|
assert result.is_new_attempt
|
|
|
|
async with db_sessionmaker() as verify_session:
|
|
job = await verify_session.get(IngestionJob, result.ingestion_job_id)
|
|
assert job is not None
|
|
assert job.status == "succeeded"
|
|
assert storage.objects # bytes were actually written
|
|
|
|
|
|
async def test_upload_source_file_duplicate_hash_returns_existing_without_reingesting(
|
|
db_session: AsyncSession, db_sessionmaker: async_sessionmaker[AsyncSession]
|
|
) -> None:
|
|
auth = await _auth_for(db_session)
|
|
storage = FakeObjectStorage()
|
|
|
|
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
|
|
assert not second.is_new_attempt
|
|
assert len(storage.objects) == 1 # no second write
|
|
|
|
|
|
async def test_upload_source_file_retries_after_failed_job(
|
|
db_session: AsyncSession, db_sessionmaker: async_sessionmaker[AsyncSession]
|
|
) -> None:
|
|
"""A duplicate upload whose last attempt failed must get a fresh job,
|
|
not be treated as already-succeeded (ADR-0017: safe to run more than
|
|
once).
|
|
"""
|
|
auth = await _auth_for(db_session)
|
|
failing_storage = FakeObjectStorage(fail_next=True)
|
|
|
|
with pytest.raises(OSError):
|
|
await _upload(sessionmaker=db_sessionmaker, storage=failing_storage, auth=auth)
|
|
|
|
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"
|
|
|
|
retry = await _upload(sessionmaker=db_sessionmaker, storage=FakeObjectStorage(), auth=auth)
|
|
|
|
assert retry.status == "succeeded"
|
|
assert retry.is_new_attempt
|
|
assert retry.file_id == jobs[0].source_file_id
|
|
|
|
|
|
async def test_upload_source_file_is_tenant_isolated_for_identical_content(
|
|
db_session: AsyncSession, db_sessionmaker: async_sessionmaker[AsyncSession]
|
|
) -> None:
|
|
"""Two different tenants uploading byte-identical files must not collide
|
|
on the (tenant_id, domain, content_sha256) idempotency key.
|
|
"""
|
|
auth_a = await _auth_for(db_session)
|
|
auth_b = await _auth_for(db_session)
|
|
|
|
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 == []
|