diff --git a/AGENTS.md b/AGENTS.md index fa13a28f..4dc21c38 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -589,8 +589,8 @@ are resolved to chunks, then assembled like classic. Map-nav is archived at `hydration.result_assembly.assemble_retrieval_results()`: 1. Filters by `exclude_document_ids` and `exclude_sections` -2. Filters by `allowed_chunk_types` (data_type parameter) -3. Hydrates `connect_to` targets (related table chunks inlined into text) +2. Hydrates `connect_to` targets +3. Filters by `allowed_chunk_types` while retaining body chunks connected to a requested image/table 4. Cleans asset path references from content 5. Public projection builds `source`: `{document_id, source_file_name, section_path}` plus `page_nums` for `chunk_type=page` when present diff --git a/apps/api/alembic/versions/0b1c2d3e4f5a_add_demo_materialization_claim_state.py b/apps/api/alembic/versions/0b1c2d3e4f5a_add_demo_materialization_claim_state.py new file mode 100644 index 00000000..81f676d0 --- /dev/null +++ b/apps/api/alembic/versions/0b1c2d3e4f5a_add_demo_materialization_claim_state.py @@ -0,0 +1,52 @@ +"""Add demo materialization claim state.""" + +from __future__ import annotations + +from collections.abc import Sequence + +from alembic import op +import sqlalchemy as sa + + +revision: str = "0b1c2d3e4f5a" +down_revision: str | Sequence[str] | None = "e4f5a6b7c8d9" +branch_labels: Sequence[str] | None = None +depends_on: Sequence[str] | None = None + + +def upgrade() -> None: + op.add_column( + "demo_materializations", + sa.Column("status", sa.String(length=32), nullable=True), + ) + op.add_column( + "demo_materializations", + sa.Column("claimed_at", sa.DateTime(), nullable=True), + ) + op.execute( + "UPDATE demo_materializations SET status = 'ready' WHERE status IS NULL" + ) + op.alter_column( + "demo_materializations", + "status", + existing_type=sa.String(length=32), + nullable=False, + server_default="ready", + ) + op.alter_column( + "demo_materializations", + "document_id", + existing_type=sa.String(length=36), + nullable=True, + ) + + +def downgrade() -> None: + op.alter_column( + "demo_materializations", + "document_id", + existing_type=sa.String(length=36), + nullable=False, + ) + op.drop_column("demo_materializations", "claimed_at") + op.drop_column("demo_materializations", "status") diff --git a/apps/api/app/services/demo/source_materializer.py b/apps/api/app/services/demo/source_materializer.py index df9e2f0c..1e903f53 100644 --- a/apps/api/app/services/demo/source_materializer.py +++ b/apps/api/app/services/demo/source_materializer.py @@ -4,23 +4,30 @@ import shutil import tempfile +import time +from collections.abc import Iterable from dataclasses import dataclass -from datetime import datetime, timezone +from datetime import datetime, timedelta, timezone from hashlib import blake2b from pathlib import Path from uuid import uuid4 +import logfire + from app.services.demo.source_catalog import DemoSourceCatalog, DemoSourceDefinition -from sqlalchemy import func, select +from sqlalchemy.exc import IntegrityError +from sqlalchemy import delete, func, select from sqlalchemy.ext.asyncio import AsyncSession -from shared.core.exceptions.domain_exceptions import ValidationException +from shared.core.exceptions.domain_exceptions import ConflictException, ValidationException from shared.models.database.demo_materialization import DemoMaterialization -from shared.models.database.document import Document from shared.models.database.job import Job from shared.models.database.job_result import JobResult from shared.services.retrieval.cache_service import invalidate_retrieval_cache_namespaces from shared.services.retrieval.publication_service import RetrievalPublicationService +from shared.services.retrieval.publication_models import DocumentPublicationScope +from shared.services.redis import RedisPublicationSemaphore, RedisServiceFactory +from shared.core.config import settings from shared.services.storage.result_storage import get_result_storage @@ -48,6 +55,7 @@ def __init__( ) -> None: self._catalog = catalog self._publication_service = publication_service or RetrievalPublicationService() + self._redis_service = RedisServiceFactory.get_service() async def materialize_sources( self, @@ -73,17 +81,46 @@ async def materialize_sources( self._catalog.require_source(demo_source_id) for demo_source_id in selected_demo_source_ids ] - results: list[MaterializedDemoSource] = [] - for source in selected_sources: - result = await self._materialize_source( + try: + claims = await self._claim_sources( db, user_id=user_id, namespace=namespace, - source=source, + sources=selected_sources, ) - results.append(result) - + except IntegrityError as error: + await db.rollback() + raise ConflictException( + user_message="This demo source is currently being materialized.", + reason="ABORTED", + resource="Demo materialization", + ) from error await db.commit() + results: list[MaterializedDemoSource] = [] + for source in selected_sources: + try: + semaphore = RedisPublicationSemaphore( + self._redis_service, + concurrency=settings.MATERIALIZATION_DB_PUBLICATION_CONCURRENCY, + lease_seconds=settings.MATERIALIZATION_DB_PUBLICATION_LEASE_SECONDS, + acquire_timeout_seconds=settings.MATERIALIZATION_DB_PUBLICATION_ACQUIRE_TIMEOUT_SECONDS, + ) + result = await self._materialize_source( + db, + user_id=user_id, + namespace=namespace, + source=source, + claim=claims[source.demo_source_id], + semaphore=semaphore, + ) + results.append(result) + except Exception: + await db.rollback() + await self._release_claims( + db, + claims=claims.values(), + ) + raise await invalidate_retrieval_cache_namespaces( user_id=user_id, namespaces=[namespace], @@ -97,37 +134,23 @@ async def _materialize_source( user_id: str, namespace: str, source: DemoSourceDefinition, + claim: DemoMaterialization, + semaphore: RedisPublicationSemaphore, ) -> MaterializedDemoSource: - await _lock_materialization_scope( - db, - user_id=user_id, - namespace=namespace, - demo_source_id=source.demo_source_id, - ) - existing = await self._get_existing_materialization( - db, - user_id=user_id, - namespace=namespace, - demo_source_id=source.demo_source_id, - ) - if existing is not None and await self._is_active_document( - db, - document_id=existing.document_id, - ): - return _materialized_source_payload( - source=source, - document_id=existing.document_id, - status="existing", - ) - document_id = f"doc_{uuid4().hex[:12]}" job_id = f"job_demo_{uuid4().hex[:12]}" job_result_id = str(uuid4()) timestamp = _utc_now() + stage_started_at = time.perf_counter() result_bundle = _upload_demo_result_bundle( job_id=job_id, source_directory=self._catalog.source_directory(source), ) + logfire.info( + "Demo materialization source bundle upload completed", + demo_source_id=source.demo_source_id, + duration_seconds=time.perf_counter() - stage_started_at, + ) db.add( Job( @@ -167,77 +190,168 @@ async def _materialize_source( updated_at=timestamp, ) ) - await db.flush() chunks = self._catalog.publication_chunks(source) - await db.run_sync( - lambda sync_db: self._publication_service.publish_document_state( - sync_db, - job_id=job_id, - job_result_id=job_result_id, - chunks=[dict(chunk) for chunk in chunks], - ) + stage_started_at = time.perf_counter() + wait_seconds = await semaphore.acquire() + logfire.info( + "Demo materialization publication semaphore acquired", + demo_source_id=source.demo_source_id, + wait_seconds=wait_seconds, ) - await db.run_sync( - lambda sync_db: self._publication_service.publish_document_graph( - sync_db, - job_id=job_id, - job_result_id=job_result_id, + try: + base_rows_started_at = time.perf_counter() + await db.flush() + logfire.info( + "Demo materialization base rows completed", + demo_source_id=source.demo_source_id, + duration_seconds=time.perf_counter() - base_rows_started_at, + ) + published_state = await db.run_sync( + lambda sync_db: self._publication_service.publish_document_state( + sync_db, + job_id=job_id, + job_result_id=job_result_id, + chunks=[dict(chunk) for chunk in chunks], + update_namespace_snapshot=False, + ) + ) + await db.run_sync( + lambda sync_db: self._publication_service.publish_document_graph( + sync_db, + job_id=job_id, + job_result_id=job_result_id, + ) + ) + await db.flush() + logfire.info( + "Demo materialization sections chunks and map index completed", + demo_source_id=source.demo_source_id, + duration_seconds=time.perf_counter() - stage_started_at, + chunk_count=len(chunks), ) - ) - await db.flush() - if existing is None: - db.add( - DemoMaterialization( - user_id=user_id, - namespace=namespace, - demo_source_id=source.demo_source_id, - document_id=document_id, - created_at=timestamp, - updated_at=timestamp, + if published_state is None or published_state.document_id != document_id: + raise RuntimeError("Demo publication did not create its requested document") + if published_state.manifest_payload is None: + raise RuntimeError("Demo publication did not create a serving manifest") + manifest_payload = published_state.manifest_payload + stage_started_at = time.perf_counter() + await db.run_sync( + lambda sync_db: self._publication_service.update_namespace_snapshot( + sync_db, + scope=DocumentPublicationScope( + user_id=user_id, + namespace=namespace, + document_id=document_id, + job_result_id=job_result_id, + source_file_name=source.title, + ), + manifest_payload=manifest_payload, ) ) - else: - existing.document_id = document_id - existing.updated_at = timestamp + logfire.info( + "Demo materialization namespace snapshot completed", + demo_source_id=source.demo_source_id, + duration_seconds=time.perf_counter() - stage_started_at, + ) + commit_started_at = time.perf_counter() + await db.commit() + logfire.info( + "Demo materialization publication commit completed", + demo_source_id=source.demo_source_id, + duration_seconds=time.perf_counter() - commit_started_at, + ) + finally: + await semaphore.release() + + claim.document_id = document_id + claim.status = "ready" + claim.claimed_at = None + claim.updated_at = timestamp await db.flush() + await db.commit() return _materialized_source_payload( source=source, document_id=document_id, status="created", ) - async def _get_existing_materialization( + async def _claim_sources( self, db: AsyncSession, *, user_id: str, namespace: str, - demo_source_id: str, - ) -> DemoMaterialization | None: - result = await db.execute( - select(DemoMaterialization) - .where(DemoMaterialization.user_id == user_id) - .where(DemoMaterialization.namespace == namespace) - .where(DemoMaterialization.demo_source_id == demo_source_id) - .with_for_update() - .limit(1) - ) - return result.scalar_one_or_none() + sources: list[DemoSourceDefinition], + ) -> dict[str, DemoMaterialization]: + now = _utc_now() + claims: dict[str, DemoMaterialization] = {} + for source in sorted(sources, key=lambda item: item.demo_source_id): + lock_acquired = await db.scalar( + select( + func.pg_try_advisory_xact_lock( + _materialization_lock_id( + user_id=user_id, + namespace=namespace, + demo_source_id=source.demo_source_id, + ) + ) + ) + ) + if not lock_acquired: + raise ConflictException( + user_message="This demo source is currently being materialized.", + reason="ABORTED", + resource="Demo materialization", + resource_id=source.demo_source_id, + ) + result = await db.execute( + select(DemoMaterialization) + .where(DemoMaterialization.user_id == user_id) + .where(DemoMaterialization.namespace == namespace) + .where(DemoMaterialization.demo_source_id == source.demo_source_id) + ) + existing = result.scalar_one_or_none() + if existing is not None: + if not _is_stale_claim(existing, now=now): + raise _materialization_conflict(existing, source) + claim = existing + claim.status = "materializing" + claim.document_id = None + claim.claimed_at = now + claim.updated_at = now + else: + claim = DemoMaterialization( + user_id=user_id, + namespace=namespace, + demo_source_id=source.demo_source_id, + status="materializing", + document_id=None, + claimed_at=now, + created_at=now, + updated_at=now, + ) + db.add(claim) + claims[source.demo_source_id] = claim + await db.flush() + return claims - async def _is_active_document( + async def _release_claims( self, db: AsyncSession, *, - document_id: str, - ) -> bool: - result = await db.execute( - select(Document.document_id) - .where(Document.document_id == document_id) - .where(Document.status == "active") - .limit(1) + claims: Iterable[DemoMaterialization], + ) -> None: + claim_ids = [claim.id for claim in claims] + if not claim_ids: + return + await db.execute( + delete(DemoMaterialization).where( + DemoMaterialization.status == "materializing", + DemoMaterialization.id.in_(claim_ids), + ) ) - return result.scalar_one_or_none() is not None + await db.commit() def _deduplicate_source_ids(demo_source_ids: list[str]) -> list[str]: @@ -252,19 +366,12 @@ def _deduplicate_source_ids(demo_source_ids: list[str]) -> list[str]: return selected -async def _lock_materialization_scope( - db: AsyncSession, - *, - user_id: str, - namespace: str, - demo_source_id: str, -) -> None: - lock_id = _materialization_lock_id( - user_id=user_id, - namespace=namespace, - demo_source_id=demo_source_id, +def _is_stale_claim(materialization: DemoMaterialization, *, now: datetime) -> bool: + return ( + materialization.status == "materializing" + and materialization.claimed_at is not None + and materialization.claimed_at < now - timedelta(minutes=10) ) - await db.execute(select(func.pg_advisory_xact_lock(lock_id))) def _materialization_lock_id( @@ -278,6 +385,23 @@ def _materialization_lock_id( return int.from_bytes(digest, byteorder="big", signed=True) +def _materialization_conflict( + materialization: DemoMaterialization, + source: DemoSourceDefinition, +) -> ConflictException: + reason = "ALREADY_EXISTS" if materialization.status == "ready" else "ABORTED" + return ConflictException( + user_message=( + "This demo source has already been materialized." + if reason == "ALREADY_EXISTS" + else "This demo source is currently being materialized." + ), + reason=reason, + resource="Demo materialization", + resource_id=source.demo_source_id, + ) + + def _materialized_source_payload( *, source: DemoSourceDefinition, diff --git a/apps/api/tests/contract/test_demo_documents_contract.py b/apps/api/tests/contract/test_demo_documents_contract.py index bfe13e10..f2eab17e 100644 --- a/apps/api/tests/contract/test_demo_documents_contract.py +++ b/apps/api/tests/contract/test_demo_documents_contract.py @@ -365,7 +365,7 @@ async def test_should_materialize_demo_source_without_parse_or_credit_charge( assert empty_cached_response.status_code == 200 assert first_response.status_code == 200 - assert retry_response.status_code == 200 + assert retry_response.status_code == 409 assert retrieval_response.status_code == 200 assert document_chunks_response.status_code == 200 @@ -373,12 +373,10 @@ async def test_should_materialize_demo_source_without_parse_or_credit_charge( assert cast(list[dict[str, Any]], empty_cached_body["results"]) == [] first_source = cast(dict[str, Any], first_response.json()["sources"][0]) - retry_source = cast(dict[str, Any], retry_response.json()["sources"][0]) document_id = str(first_source["document_id"]) assert first_source["status"] == "created" - assert retry_source["status"] == "existing" - assert retry_source["document_id"] == document_id + assert retry_response.json()["error"]["details"]["reason"] == "ALREADY_EXISTS" materialization_rows = await ContractDatabase.fetch_all( """ @@ -517,7 +515,39 @@ async def test_should_materialize_each_normalized_demo_source_once_per_request( @pytest.mark.asyncio -async def test_should_serialize_concurrent_first_demo_materialization( +async def test_should_reject_sequential_duplicate_demo_materialization( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + monkeypatch: MonkeyPatch, +) -> None: + fake_result_storage = FakeResultStorage() + + async with developer_api_client_factory() as api_client: + import app.services.demo.source_materializer as source_materializer_module + + monkeypatch.setattr( + source_materializer_module, + "get_result_storage", + lambda: fake_result_storage, + ) + first_response = await api_client.post( + "/api/v1/demo/materializations", + json={"namespace": "contract-demo-duplicate", "demo_source_ids": [DEMO_SOURCE_ID]}, + ) + second_response = await api_client.post( + "/api/v1/demo/materializations", + json={"namespace": "contract-demo-duplicate", "demo_source_ids": [DEMO_SOURCE_ID]}, + ) + + assert first_response.status_code == 200 + assert second_response.status_code == 409 + assert second_response.json()["error"]["details"]["reason"] == "ALREADY_EXISTS" + assert len(fake_result_storage.raw_files_by_job_id) == 1 + + +@pytest.mark.asyncio +async def test_should_reject_concurrent_duplicate_demo_materialization( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] ], @@ -550,19 +580,12 @@ async def test_should_serialize_concurrent_first_demo_materialization( ), ) - assert first_response.status_code == 200 - assert second_response.status_code == 200 - - first_source = cast(dict[str, Any], first_response.json()["sources"][0]) - second_source = cast(dict[str, Any], second_response.json()["sources"][0]) - document_ids = { - str(first_source["document_id"]), - str(second_source["document_id"]), - } - statuses = {str(first_source["status"]), str(second_source["status"])} - - assert len(document_ids) == 1 - assert statuses == {"created", "existing"} + statuses = {first_response.status_code, second_response.status_code} + assert statuses == {200, 409} + conflict_response = ( + first_response if first_response.status_code == 409 else second_response + ) + assert conflict_response.json()["error"]["details"]["reason"] == "ABORTED" materialization_rows = await ContractDatabase.fetch_all( """ @@ -585,12 +608,8 @@ async def test_should_serialize_concurrent_first_demo_materialization( {"demo_source_id": DEMO_SOURCE_ID}, ) - assert materialization_rows == [ - { - "demo_source_id": DEMO_SOURCE_ID, - "document_id": next(iter(document_ids)), - } - ] + assert len(materialization_rows) == 1 + assert materialization_rows[0]["document_id"] is not None assert len(job_rows) == 1 diff --git a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py index d42148d8..57cc5077 100644 --- a/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py +++ b/apps/api/tests/contract/test_retrieval_classic_map_unit_contract.py @@ -781,6 +781,252 @@ async def test_classic_route_image_filter_scores_only_units_with_images( assert results[0]["chunk_type"] == "image" +async def test_classic_route_mixed_section_asset_filters_do_not_mark_index_unusable( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + identifier = uuid4().hex[:8] + namespace = f"classic-mixed-{identifier}" + async with developer_api_client_factory() as api_client: + mixed = await _publish_document( + namespace=namespace, + source_file_name="mixed-assets.pdf", + chunks=[ + { + "chunk_id": f"plain-{identifier}", + "type": "text", + "content": "mixedsection plaintext filler with no assets", + "path": "mixed-assets.pdf/Root/Plain/body", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"chart-body-{identifier}", + "type": "text", + "content": "mixedsection chartmarker beside a plot", + "path": "mixed-assets.pdf/Root/Chart/body", + "order": 2, + "metadata": {"connect_to": [{"target": f"chart-{identifier}"}]}, + }, + { + "chunk_id": f"chart-{identifier}", + "type": "image", + "content": "mixedsection chartmarker plot", + "path": "images/mixed-chart.png", + "order": 3, + "file_path": "images/mixed-chart.png", + "metadata": {}, + }, + { + "chunk_id": f"grid-body-{identifier}", + "type": "text", + "content": "mixedsection tablemarker beside a grid", + "path": "mixed-assets.pdf/Root/Grid/body", + "order": 4, + "metadata": {"connect_to": [{"target": f"grid-{identifier}"}]}, + }, + { + "chunk_id": f"grid-{identifier}", + "type": "table", + "content": "
mixedsection tablemarker grid
", + "path": "tables/mixed-grid.html", + "order": 5, + "file_path": "tables/mixed-grid.html", + "metadata": {}, + }, + ], + ) + for extra_kind, extra_type, extra_path, extra_content in ( + ("photo", "image", "images/extra-photo.png", "unrelated landscape photo"), + ("diagram", "image", "images/extra-diagram.png", "unrelated diagram"), + ("sheet", "table", "tables/extra-sheet.html", "
unrelated sheet
"), + ("grid", "table", "tables/extra-grid.html", "
unrelated grid
"), + ): + extra_chunk_id = f"{extra_kind}-{identifier}" + await _publish_document( + namespace=namespace, + source_file_name=f"extra-{extra_kind}.pdf", + chunks=[ + { + "chunk_id": f"{extra_kind}-body-{identifier}", + "type": "text", + "content": f"unrelated {extra_kind} caption", + "path": f"extra-{extra_kind}.pdf/Root/Section/body", + "order": 1, + "metadata": {"connect_to": [{"target": extra_chunk_id}]}, + }, + { + "chunk_id": extra_chunk_id, + "type": extra_type, + "content": extra_content, + "path": extra_path, + "order": 2, + "file_path": extra_path, + "metadata": {}, + }, + ], + ) + async with contract_db_session() as db: + mixed_units = list( + ( + await db.execute( + select(DocumentMapUnit).where( + DocumentMapUnit.document_id == mixed["document_id"] + ) + ) + ).scalars() + ) + assert len(mixed_units) >= 3 + assert any(unit.has_image and not unit.has_table for unit in mixed_units) + assert any(unit.has_table and not unit.has_image for unit in mixed_units) + assert any(not unit.has_image and not unit.has_table for unit in mixed_units) + + image_response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "mixedsection chartmarker", + "top_k": 1, + "use_agentic": False, + "chunk_types": ["image"], + }, + ) + table_response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "mixedsection tablemarker", + "top_k": 1, + "use_agentic": False, + "chunk_types": ["table"], + }, + ) + unfiltered_response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "mixedsection plaintext", + "top_k": 1, + "use_agentic": False, + }, + ) + + assert image_response.status_code == 200 + image_body = cast(dict[str, object], image_response.json()) + image_results = cast(list[dict[str, object]], image_body["results"]) + assert image_body["router_used"] == "classic_topk" + assert [row["chunk_id"] for row in image_results] == [f"chart-{identifier}"] + assert image_results[0]["chunk_type"] == "image" + + assert table_response.status_code == 200 + table_body = cast(dict[str, object], table_response.json()) + table_results = cast(list[dict[str, object]], table_body["results"]) + assert table_body["router_used"] == "classic_topk" + assert [row["chunk_id"] for row in table_results] == [f"grid-{identifier}"] + assert table_results[0]["chunk_type"] == "table" + + assert unfiltered_response.status_code == 200 + unfiltered_body = cast(dict[str, object], unfiltered_response.json()) + assert unfiltered_body["router_used"] == "classic_topk" + + +@pytest.mark.parametrize( + "incomplete_index_kind", ["legacy_format", "missing_index"] +) +async def test_classic_route_image_filter_raises_when_index_is_unusable( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], + incomplete_index_kind: str, +) -> None: + identifier = uuid4().hex[:8] + namespace = f"classic-image-unusable-{incomplete_index_kind}-{identifier}" + async with developer_api_client_factory() as api_client: + document = await _publish_document( + namespace=namespace, + source_file_name="unusable-image.pdf", + chunks=[ + { + "chunk_id": f"plain-{identifier}", + "type": "text", + "content": "unusable image plaintext filler", + "path": "unusable-image.pdf/Root/Plain/body", + "order": 1, + "metadata": {}, + }, + { + "chunk_id": f"body-{identifier}", + "type": "text", + "content": "unusable image marker next to a chart", + "path": "unusable-image.pdf/Root/Chart/body", + "order": 2, + "metadata": {"connect_to": [{"target": f"chart-{identifier}"}]}, + }, + { + "chunk_id": f"chart-{identifier}", + "type": "image", + "content": "unusable image marker chart", + "path": "images/unusable-chart.png", + "order": 3, + "file_path": "images/unusable-chart.png", + "metadata": {}, + }, + ], + ) + await _publish_document( + namespace=namespace, + source_file_name="extra-image.pdf", + chunks=[ + { + "chunk_id": f"extra-body-{identifier}", + "type": "text", + "content": "unrelated extra caption", + "path": "extra-image.pdf/Root/Section/body", + "order": 1, + "metadata": {"connect_to": [{"target": f"extra-{identifier}"}]}, + }, + { + "chunk_id": f"extra-{identifier}", + "type": "image", + "content": "unrelated extra photo", + "path": "images/extra.png", + "order": 2, + "file_path": "images/extra.png", + "metadata": {}, + }, + ], + ) + if incomplete_index_kind == "legacy_format": + await ContractDatabase.execute( + """ + UPDATE document_map_unit_indexes + SET format_version = 1 + WHERE document_id = :document_id + """, + {"document_id": document["document_id"]}, + ) + else: + await ContractDatabase.execute( + """ + DELETE FROM document_map_unit_indexes + WHERE document_id = :document_id + """, + {"document_id": document["document_id"]}, + ) + with pytest.raises(RuntimeError, match="map-unit index is incomplete"): + await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": namespace, + "query": "unusable image marker", + "top_k": 1, + "use_agentic": False, + "chunk_types": ["image"], + }, + ) + + async def test_connected_hydration_does_not_load_legacy_job_chunks( developer_api_client_factory: Callable[ [], AbstractAsyncContextManager[AsyncClient] diff --git a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py index 1319db3f..e9346c7e 100644 --- a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py +++ b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py @@ -2,6 +2,8 @@ import os import shutil +import threading +import time from pathlib import Path import pytest @@ -29,6 +31,7 @@ from shared.services.retrieval.execution.reference_resolver import ( # noqa: E402 resolve_workflow_references, ) +from shared.core.exceptions.domain_exceptions import StorageServiceException # noqa: E402 from shared.services.storage.result_storage import JobResultStorage # noqa: E402 from shared.services.storage.page_pdf_crop import crop_source_pdf_pages # noqa: E402 @@ -267,7 +270,9 @@ class FakeStorageAdapter: def __init__(self) -> None: self.uploaded_keys: list[str] = [] - def upload_file(self, local_path: str, key: str, bucket: str | None = None): + def upload_file( + self, local_path: str, key: str, bucket: str | None = None + ) -> dict[str, str]: del local_path, bucket self.uploaded_keys.append(key) return {"key": key} @@ -322,6 +327,173 @@ def generate_presigned_url(self, *args, **kwargs) -> str: assert "results/job-1/debug.csv" not in adapter.uploaded_keys +def test_result_storage_uploads_raw_files_with_bounded_concurrency(tmp_path) -> None: + class ConcurrentStorageAdapter: + def __init__(self) -> None: + self.uploaded_keys: list[str] = [] + self.active_uploads = 0 + self.maximum_active_uploads = 0 + self.lock = threading.Lock() + + def upload_file( + self, local_path: str, key: str, bucket: str | None = None + ) -> dict[str, str]: + del local_path, bucket + with self.lock: + self.active_uploads += 1 + self.maximum_active_uploads = max( + self.maximum_active_uploads, + self.active_uploads, + ) + time.sleep(0.01) + with self.lock: + self.uploaded_keys.append(key) + self.active_uploads -= 1 + return {"key": key} + + def generate_presigned_url(self, *args, **kwargs) -> str: + del args, kwargs + return "https://assets.example.test/file" + + result_dir = tmp_path / "result" + result_dir.mkdir() + for index in range(8): + (result_dir / f"asset-{index}.bin").write_bytes(b"asset") + zip_path = tmp_path / "result.zip" + zip_path.write_bytes(b"zip") + adapter = ConcurrentStorageAdapter() + storage = JobResultStorage( + results_bucket="test-results", + storage_adapter=adapter, # type: ignore[arg-type] + upload_concurrency=3, + ) + + bundle = storage.upload( + job_id="job-concurrent", + result_dir=str(result_dir), + zip_file_path=str(zip_path), + ) + + assert len(bundle.raw_files) == 8 + assert adapter.maximum_active_uploads <= 3 + assert adapter.maximum_active_uploads > 1 + + +def test_result_storage_bounds_uploads_across_concurrent_tasks(tmp_path) -> None: + class ProcessStorageAdapter: + def __init__(self) -> None: + self.active_uploads = 0 + self.maximum_active_uploads = 0 + self.lock = threading.Lock() + + def upload_file( + self, local_path: str, key: str, bucket: str | None = None + ) -> dict[str, str]: + del local_path, bucket + if key.endswith(".zip"): + return {"key": key} + with self.lock: + self.active_uploads += 1 + self.maximum_active_uploads = max( + self.maximum_active_uploads, + self.active_uploads, + ) + time.sleep(0.02) + with self.lock: + self.active_uploads -= 1 + return {"key": key} + + def generate_presigned_url(self, *args, **kwargs) -> str: + del args, kwargs + return "https://assets.example.test/file" + + def create_result_directory(name: str) -> tuple[Path, Path]: + result_dir = tmp_path / name + result_dir.mkdir() + for index in range(8): + (result_dir / f"asset-{index}.bin").write_bytes(b"asset") + zip_path = tmp_path / f"{name}.zip" + zip_path.write_bytes(b"zip") + return result_dir, zip_path + + adapter = ProcessStorageAdapter() + first_result_dir, first_zip_path = create_result_directory("first-result") + second_result_dir, second_zip_path = create_result_directory("second-result") + first_storage = JobResultStorage( + storage_adapter=adapter, # type: ignore[arg-type] + upload_concurrency=3, + ) + second_storage = JobResultStorage( + storage_adapter=adapter, # type: ignore[arg-type] + upload_concurrency=3, + ) + failures: list[Exception] = [] + + def upload_result( + storage: JobResultStorage, + job_id: str, + result_dir: Path, + zip_path: Path, + ) -> None: + try: + storage.upload( + job_id=job_id, + result_dir=str(result_dir), + zip_file_path=str(zip_path), + ) + except Exception as exc: + failures.append(exc) + + first_thread = threading.Thread( + target=upload_result, + args=(first_storage, "first-job", first_result_dir, first_zip_path), + ) + second_thread = threading.Thread( + target=upload_result, + args=(second_storage, "second-job", second_result_dir, second_zip_path), + ) + + first_thread.start() + second_thread.start() + first_thread.join() + second_thread.join() + + assert failures == [] + assert adapter.maximum_active_uploads == 3 + + +def test_result_storage_upload_propagates_raw_file_failure(tmp_path) -> None: + class FailingStorageAdapter: + def upload_file(self, local_path: str, key: str, bucket: str | None = None): + del local_path, bucket + if key.endswith("asset-2.bin"): + raise RuntimeError("upload failed") + return {"key": key} + + def generate_presigned_url(self, *args, **kwargs) -> str: + del args, kwargs + return "https://assets.example.test/file" + + result_dir = tmp_path / "result" + result_dir.mkdir() + for index in range(4): + (result_dir / f"asset-{index}.bin").write_bytes(b"asset") + zip_path = tmp_path / "result.zip" + zip_path.write_bytes(b"zip") + storage = JobResultStorage( + results_bucket="test-results", + storage_adapter=FailingStorageAdapter(), # type: ignore[arg-type] + upload_concurrency=2, + ) + + with pytest.raises(StorageServiceException, match="Storage upload failed|upload failed"): + storage.upload( + job_id="job-failure", + result_dir=str(result_dir), + zip_file_path=str(zip_path), + ) + + def test_crop_source_pdf_pages_uploads_and_reuses_page_pdf_cache(tmp_path) -> None: from pypdf import PdfReader, PdfWriter diff --git a/packages/shared-python/shared/core/config/database.py b/packages/shared-python/shared/core/config/database.py index a9baac5f..53d1b97d 100644 --- a/packages/shared-python/shared/core/config/database.py +++ b/packages/shared-python/shared/core/config/database.py @@ -47,6 +47,22 @@ class DatabaseConfig(BaseModel): default=50, description="Celery gevent worker concurrency" ) + MATERIALIZATION_DB_PUBLICATION_CONCURRENCY: int = Field( + default=2, + ge=1, + description="Maximum concurrent materialization database publications", + ) + MATERIALIZATION_DB_PUBLICATION_LEASE_SECONDS: int = Field( + default=900, + ge=60, + description="Redis lease duration for a materialization publication permit", + ) + MATERIALIZATION_DB_PUBLICATION_ACQUIRE_TIMEOUT_SECONDS: float = Field( + default=30.0, + ge=0.1, + description="Maximum time to wait for a materialization publication permit", + ) + def get_ssl_connect_args(self) -> dict: """Return SSL connect args for psycopg2.""" ssl_args = {"sslmode": self.DB_SSL_MODE} diff --git a/packages/shared-python/shared/core/config/storage.py b/packages/shared-python/shared/core/config/storage.py index e04d69d6..930ae06c 100644 --- a/packages/shared-python/shared/core/config/storage.py +++ b/packages/shared-python/shared/core/config/storage.py @@ -52,6 +52,12 @@ class StorageConfig(BaseModel): S3_ADDRESSING_STYLE: str = Field( default="auto", description="S3 addressing style: auto, path, or virtual" ) + S3_MAX_POOL_CONNECTIONS: int = Field( + default=20, + ge=1, + le=256, + description="Maximum connections retained by the shared S3 client pool.", + ) # OSS-only configuration. OSS_ENDPOINT: str = Field( @@ -96,6 +102,15 @@ class StorageConfig(BaseModel): le=10, description="Maximum concurrent MinerU API calls for shard parsing.", ) + RESULT_UPLOAD_CONCURRENCY: int = Field( + default=20, + ge=1, + le=64, + description=( + "Maximum concurrent raw result-file uploads per materialization. " + "ZIP bundles remain single-object uploads." + ), + ) SUPPORTED_EXTENSIONS: str = Field( default=".doc,.docx,.pdf,.txt,.xls,.xlsx,.pptx,.jpg,.jpeg,.png,.md,.html,.htm", description="Supported file extensions", @@ -146,6 +161,7 @@ def get_s3_client(self) -> BaseClient: # Configure retries. config_kwargs["retries"] = {"max_attempts": 5, "mode": "standard"} + config_kwargs["max_pool_connections"] = self.S3_MAX_POOL_CONNECTIONS config = Config(**config_kwargs) if config_kwargs else None diff --git a/packages/shared-python/shared/models/database/demo_materialization.py b/packages/shared-python/shared/models/database/demo_materialization.py index d4cc1a7a..32bbc950 100644 --- a/packages/shared-python/shared/models/database/demo_materialization.py +++ b/packages/shared-python/shared/models/database/demo_materialization.py @@ -29,10 +29,16 @@ class DemoMaterialization(Base): String(255), nullable=False, default="default" ) demo_source_id: Mapped[str] = mapped_column(String(128), nullable=False) - document_id: Mapped[str] = mapped_column( + status: Mapped[str] = mapped_column( + String(32), nullable=False, default="ready" + ) + document_id: Mapped[str | None] = mapped_column( String(36), ForeignKey("documents.document_id", ondelete="CASCADE"), - nullable=False, + nullable=True, + ) + claimed_at: Mapped[datetime | None] = mapped_column( + DateTime, nullable=True ) created_at: Mapped[datetime] = mapped_column( DateTime, default=utc_now_naive, nullable=False diff --git a/packages/shared-python/shared/services/jobs/lifecycle/publication.py b/packages/shared-python/shared/services/jobs/lifecycle/publication.py index a8a3e747..4ed9df81 100644 --- a/packages/shared-python/shared/services/jobs/lifecycle/publication.py +++ b/packages/shared-python/shared/services/jobs/lifecycle/publication.py @@ -1,5 +1,6 @@ from __future__ import annotations +from contextlib import nullcontext from dataclasses import dataclass from typing import Any @@ -11,6 +12,8 @@ from shared.models.schemas.job_metadata import JobMetadataHelper from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.redis.redis_sync_service import SyncRedisServiceFactory +from shared.services.redis.publication_semaphore import SyncRedisPublicationSemaphore +from shared.core.config import settings from shared.services.retrieval.publication_service import RetrievalPublicationService from shared.services.retrieval.publication_models import ( ExistingDocumentScope, @@ -53,25 +56,38 @@ def publish_result( section_summaries: dict[str, str] | None, document_top_summary: str | None = None, ) -> JobPublicationOutcome: - previous_document_scope = self._retrieval_publication.get_existing_document_scope( - db, - job_id=job_id, + semaphore = SyncRedisPublicationSemaphore( + SyncRedisServiceFactory.get_service(), + concurrency=settings.MATERIALIZATION_DB_PUBLICATION_CONCURRENCY, + lease_seconds=settings.MATERIALIZATION_DB_PUBLICATION_LEASE_SECONDS, + acquire_timeout_seconds=settings.MATERIALIZATION_DB_PUBLICATION_ACQUIRE_TIMEOUT_SECONDS, ) - published_document_state = self._retrieval_publication.publish_document_state( - db, - job_id=job_id, - job_result_id=job_result_id, - chunks=chunks, - section_summaries=section_summaries, + job_type = db.execute( + select(Job.job_type).where(Job.job_id == job_id) + ).scalar_one_or_none() + publication_context = ( + semaphore if job_type == "demo_materialization" else nullcontext() ) - if _should_publish_document_graph(published_document_state): - assert published_document_state is not None - self._retrieval_publication.publish_document_graph( + with publication_context: + previous_document_scope = self._retrieval_publication.get_existing_document_scope( + db, + job_id=job_id, + ) + published_document_state = self._retrieval_publication.publish_document_state( db, job_id=job_id, job_result_id=job_result_id, - top_summary=document_top_summary, + chunks=chunks, + section_summaries=section_summaries, ) + if _should_publish_document_graph(published_document_state): + assert published_document_state is not None + self._retrieval_publication.publish_document_graph( + db, + job_id=job_id, + job_result_id=job_result_id, + top_summary=document_top_summary, + ) cache_invalidation = self._build_cache_invalidation( db, diff --git a/packages/shared-python/shared/services/redis/__init__.py b/packages/shared-python/shared/services/redis/__init__.py index a0cb3f46..befe58f3 100644 --- a/packages/shared-python/shared/services/redis/__init__.py +++ b/packages/shared-python/shared/services/redis/__init__.py @@ -11,6 +11,10 @@ from .key_builder import RedisKeyBuilder, RedisKeyType, redis_key_builder from .redis_alerts import AlertRule, RedisAlertManager, RedisAlertNotifier from .redis_monitor import RedisMonitor + from .publication_semaphore import ( + RedisPublicationSemaphore, + SyncRedisPublicationSemaphore, + ) from .redis_service import RedisService from .redis_service_factory import RedisServiceFactory from .retry_policy import RedisHealthChecker, RedisRetry @@ -20,6 +24,8 @@ __all__ = [ "RedisService", "RedisServiceFactory", + "RedisPublicationSemaphore", + "SyncRedisPublicationSemaphore", "RedisMonitor", "RedisAlertManager", "RedisAlertNotifier", @@ -38,6 +44,8 @@ _EXPORT_MODULES: dict[str, str] = { "RedisService": "shared.services.redis.redis_service", "RedisServiceFactory": "shared.services.redis.redis_service_factory", + "RedisPublicationSemaphore": "shared.services.redis.publication_semaphore", + "SyncRedisPublicationSemaphore": "shared.services.redis.publication_semaphore", "RedisMonitor": "shared.services.redis.redis_monitor", "RedisAlertManager": "shared.services.redis.redis_alerts", "RedisAlertNotifier": "shared.services.redis.redis_alerts", diff --git a/packages/shared-python/shared/services/redis/publication_semaphore.py b/packages/shared-python/shared/services/redis/publication_semaphore.py new file mode 100644 index 00000000..742b6f73 --- /dev/null +++ b/packages/shared-python/shared/services/redis/publication_semaphore.py @@ -0,0 +1,142 @@ +"""Cross-instance concurrency control for database publication.""" + +from __future__ import annotations + +import asyncio +import time +import uuid +from typing import Any + +from loguru import logger + +from shared.services.redis.redis_service import RedisService +from shared.services.redis.redis_sync_service import SyncRedisService + +_RELEASE_SCRIPT = """ +if redis.call('get', KEYS[1]) == ARGV[1] then + return redis.call('del', KEYS[1]) +end +return 0 +""" + + +class RedisPublicationSemaphore: + """Lease-based Redis semaphore using one key per permit.""" + + def __init__( + self, + redis_service: RedisService, + *, + concurrency: int, + lease_seconds: int, + acquire_timeout_seconds: float, + poll_interval_seconds: float = 0.05, + ) -> None: + if concurrency < 1: + raise ValueError("concurrency must be at least 1") + self._redis = redis_service + self._concurrency = concurrency + self._lease_seconds = lease_seconds + self._acquire_timeout_seconds = acquire_timeout_seconds + self._poll_interval_seconds = poll_interval_seconds + self._owner = uuid.uuid4().hex + self._permit_key: str | None = None + + async def acquire(self) -> float: + """Acquire a permit and return the time spent waiting in seconds.""" + started_at = time.perf_counter() + deadline = started_at + self._acquire_timeout_seconds + while time.perf_counter() < deadline: + for permit_number in range(self._concurrency): + key = f"lock:materialization_publication:{permit_number}" + acquired = await self._redis.set_nx( + key, + self._owner, + ex=self._lease_seconds, + ) + if acquired: + self._permit_key = key + return time.perf_counter() - started_at + await asyncio.sleep(self._poll_interval_seconds) + raise TimeoutError("Timed out waiting for a materialization publication permit") + + async def release(self) -> bool: + """Release this semaphore's permit only when still its owner.""" + if self._permit_key is None: + return False + key = self._permit_key + self._permit_key = None + try: + result = await self._redis.eval( + _RELEASE_SCRIPT, + keys=[key], + args=[self._owner], + ) + return bool(result) + except Exception as error: + logger.warning(f"Failed to release publication permit: {error}") + return False + + async def __aenter__(self) -> "RedisPublicationSemaphore": + await self.acquire() + return self + + async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + await self.release() + + +class SyncRedisPublicationSemaphore: + """Synchronous lease-based semaphore for gevent worker publication.""" + + def __init__( + self, + redis_service: SyncRedisService, + *, + concurrency: int, + lease_seconds: int, + acquire_timeout_seconds: float, + poll_interval_seconds: float = 0.05, + ) -> None: + if concurrency < 1: + raise ValueError("concurrency must be at least 1") + self._redis = redis_service + self._concurrency = concurrency + self._lease_seconds = lease_seconds + self._acquire_timeout_seconds = acquire_timeout_seconds + self._poll_interval_seconds = poll_interval_seconds + self._owner = uuid.uuid4().hex + self._permit_key: str | None = None + + def acquire(self) -> float: + """Acquire a permit and return the time spent waiting in seconds.""" + started_at = time.perf_counter() + deadline = started_at + self._acquire_timeout_seconds + while time.perf_counter() < deadline: + for permit_number in range(self._concurrency): + key = f"lock:materialization_publication:{permit_number}" + if self._redis.set_nx(key, self._owner, ex=self._lease_seconds): + self._permit_key = key + return time.perf_counter() - started_at + time.sleep(self._poll_interval_seconds) + raise TimeoutError("Timed out waiting for a materialization publication permit") + + def release(self) -> bool: + """Release this semaphore's permit only when still its owner.""" + if self._permit_key is None: + return False + key = self._permit_key + self._permit_key = None + try: + return bool( + self._redis.eval(_RELEASE_SCRIPT, keys=[key], args=[self._owner]) + ) + except Exception as error: + logger.warning(f"Failed to release publication permit: {error}") + return False + + def __enter__(self) -> "SyncRedisPublicationSemaphore": + self.acquire() + return self + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + self.release() diff --git a/packages/shared-python/shared/services/redis/redis_service.py b/packages/shared-python/shared/services/redis/redis_service.py index c25a3e1d..cb45d75d 100644 --- a/packages/shared-python/shared/services/redis/redis_service.py +++ b/packages/shared-python/shared/services/redis/redis_service.py @@ -298,6 +298,32 @@ async def _operation(): original_exception=e, ) + async def eval( + self, script: str, keys: list[str], args: list[Any] | None = None + ) -> Any: + """Execute a Lua script with consistently namespaced keys.""" + try: + client = await self._get_client() + full_keys = [self._build_key(key) for key in keys] + + async def _operation() -> Any: + return await _await_redis_result( + client.eval( + script, + len(full_keys), + *(full_keys + (args or [])), + ) + ) + + return await self._execute_with_retry(_operation) + except Exception as e: + logger.error(f"Redis EVAL operation failed: {e}") + raise RedisOperationError( + internal_message=f"EVAL operation failed: {str(e)}", + operation="EVAL", + original_exception=e, + ) + # ==================== List Operations ==================== async def lpush(self, key: str, *values: Any) -> int: diff --git a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py index 610bf99f..96a38fd6 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py +++ b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py @@ -30,20 +30,15 @@ async def assemble_retrieval_results( allowed_chunk_types: set[str] | None = None, revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: - filtered_rows = filter_excluded_rows( + scoped_rows = filter_excluded_rows( rows, exclude_document_ids=exclude_document_ids, exclude_sections=exclude_sections, document_scope=document_scope, ) - if allowed_chunk_types is not None: - filtered_rows = [ - row for row in filtered_rows - if normalize_chunk_type(row.get('chunk_type')) in allowed_chunk_types - ] hydrated_rows = await hydrate_connected_target_rows( db=db, - rows=filtered_rows, + rows=scoped_rows, exclude_document_ids=exclude_document_ids, exclude_sections=exclude_sections, document_scope=document_scope, @@ -51,9 +46,14 @@ async def assemble_retrieval_results( ) rows_by_chunk_id = { str(row.get('chunk_id') or ''): row - for row in [*filtered_rows, *hydrated_rows] + for row in [*scoped_rows, *hydrated_rows] if row.get('chunk_id') } + filtered_rows = _filter_rows_by_allowed_chunk_types( + scoped_rows, + allowed_chunk_types=allowed_chunk_types, + rows_by_chunk_id=rows_by_chunk_id, + ) embedded_targets: set[str] = set() for row in filtered_rows: @@ -88,6 +88,27 @@ async def assemble_retrieval_results( return assembled +def _filter_rows_by_allowed_chunk_types( + rows: list[dict[str, Any]], + *, + allowed_chunk_types: set[str] | None, + rows_by_chunk_id: Mapping[str, dict[str, Any]], +) -> list[dict[str, Any]]: + if allowed_chunk_types is None: + return rows + + return [ + row + for row in rows + if normalize_chunk_type(row.get('chunk_type')) in allowed_chunk_types + or any( + normalize_chunk_type(rows_by_chunk_id.get(target_id, {}).get('chunk_type')) + in allowed_chunk_types + for target_id in iter_connected_target_ids(row) + ) + ] + + def _page_summary(row: dict[str, Any]) -> str: metadata = row.get('chunk_metadata') or row.get('metadata') or {} if not isinstance(metadata, dict): diff --git a/packages/shared-python/shared/services/retrieval/map_unit_index.py b/packages/shared-python/shared/services/retrieval/map_unit_index.py index c58e5f8e..19f95200 100644 --- a/packages/shared-python/shared/services/retrieval/map_unit_index.py +++ b/packages/shared-python/shared/services/retrieval/map_unit_index.py @@ -6,7 +6,7 @@ from hashlib import sha256 from uuid import uuid4 -from sqlalchemy import delete, select +from sqlalchemy import Table, delete, insert, select from sqlalchemy.orm import Session from shared.models.database.document import ( @@ -29,6 +29,8 @@ __all__ = ["MAP_UNIT_INDEX_FORMAT_VERSION", "replace_document_map_units"] +_BULK_INSERT_BATCH_SIZE = 5_000 + def replace_document_map_units( db: Session, @@ -86,6 +88,8 @@ def replace_document_map_units( ) persisted_count = 0 token_count = 0 + map_unit_rows: list[dict[str, object]] = [] + token_rows: list[dict[str, object]] = [] path_unit_df: Counter[str] = Counter() content_unit_df: Counter[str] = Counter() path_document_count: int = 0 @@ -113,39 +117,43 @@ def replace_document_map_units( # embeds them), so this is the same ownership the query-time scorer # sees, not a new computation. section_chunk_types = {u.chunk_type for u in provider.self_units(section_id)} - db.add( - DocumentMapUnit( - id=map_unit_id, - document_id=scope.document_id, - job_result_id=scope.job_result_id, - unit_id=unit_id, - section_id=section_id, - unit_kind=str(unit.get("kind") or "leaf"), - path_token_count=len(path_tokens), - content_token_count=len(content_tokens), - term_search_text_lower=str(unit.get("term_search_text") or "").lower(), - has_image="image" in section_chunk_types, - has_table="table" in section_chunk_types, - sort_order=sort_order, - ) + map_unit_rows.append( + { + "id": map_unit_id, + "document_id": scope.document_id, + "job_result_id": scope.job_result_id, + "unit_id": unit_id, + "section_id": section_id, + "unit_kind": str(unit.get("kind") or "leaf"), + "path_token_count": len(path_tokens), + "content_token_count": len(content_tokens), + "term_search_text_lower": str( + unit.get("term_search_text") or "" + ).lower(), + "has_image": "image" in section_chunk_types, + "has_table": "table" in section_chunk_types, + "sort_order": sort_order, + } ) for channel, frequencies in ( ("path", Counter(path_tokens)), ("content", Counter(content_tokens)), ): for token, frequency in frequencies.items(): - db.add( - DocumentMapUnitToken( - id=f"dmut_{uuid4().hex[:31]}", - map_unit_id=map_unit_id, - channel=channel, - token=token, - token_hash=sha256(token.encode("utf-8")).hexdigest(), - frequency=frequency, - ) + token_rows.append( + { + "id": f"dmut_{uuid4().hex[:31]}", + "map_unit_id": map_unit_id, + "channel": channel, + "token": token, + "token_hash": sha256(token.encode("utf-8")).hexdigest(), + "frequency": frequency, + } ) token_count += len(frequencies) persisted_count += 1 + _execute_bulk_insert(db, DocumentMapUnit, map_unit_rows) + _execute_bulk_insert(db, DocumentMapUnitToken, token_rows) db.add( DocumentMapUnitIndex( id=f"dmui_{uuid4().hex}", @@ -170,6 +178,19 @@ def replace_document_map_units( ) +def _execute_bulk_insert( + db: Session, + model: type[DocumentMapUnit] | type[DocumentMapUnitToken], + rows: list[dict[str, object]], +) -> None: + """Insert derived index rows in bounded Core batches.""" + if not rows: + return + table: Table = model.__table__ + for start in range(0, len(rows), _BULK_INSERT_BATCH_SIZE): + db.execute(insert(table), rows[start : start + _BULK_INSERT_BATCH_SIZE]) + + def _to_section_row(section: DocumentSection) -> SectionRow: return SectionRow( section_id=section.section_id, diff --git a/packages/shared-python/shared/services/retrieval/publication_content.py b/packages/shared-python/shared/services/retrieval/publication_content.py index 7ce75b80..90ab74e7 100644 --- a/packages/shared-python/shared/services/retrieval/publication_content.py +++ b/packages/shared-python/shared/services/retrieval/publication_content.py @@ -13,9 +13,6 @@ DocumentSection, ) from shared.services.retrieval.map_unit_index import replace_document_map_units -from shared.services.retrieval.namespace_map_snapshot import ( - patch_namespace_map_snapshot, -) from shared.services.retrieval.publication_models import DocumentPublicationScope from shared.services.retrieval.serving_manifest import persist_revision_serving_state from shared.services.retrieval.search.lexical_text import ( @@ -63,7 +60,7 @@ def replace_document_revision_content( scope: DocumentPublicationScope, chunks: list[dict[str, Any]], section_summaries: dict[str, str] | None = None, -) -> None: +) -> dict[str, Any]: """Replace retrieval sections and chunks for one published document revision.""" _delete_existing_revision_content(db, scope=scope) section_publisher = DocumentSectionPublisher( @@ -71,6 +68,7 @@ def replace_document_revision_content( scope=scope, section_summaries=section_summaries, ) + prepared_chunks: list[tuple[int, dict[str, Any], dict[str, Any], str | None, DocumentSection]] = [] for index, chunk in enumerate(chunks): safe_chunk = cast(dict[str, Any], remove_nul_characters(chunk)) chunk_metadata = _get_chunk_metadata(safe_chunk) @@ -83,6 +81,9 @@ def replace_document_revision_content( source_file_name=scope.source_file_name, ) section = section_publisher.ensure_section(section_path) + prepared_chunks.append((index, safe_chunk, chunk_metadata, source_path, section)) + db.flush() + for index, safe_chunk, chunk_metadata, source_path, section in prepared_chunks: db.add( _build_document_chunk( chunk=safe_chunk, @@ -97,7 +98,7 @@ def replace_document_revision_content( replace_document_map_units(db, scope=scope) db.flush() manifest_payload = persist_revision_serving_state(db, scope=scope) - patch_namespace_map_snapshot(db, scope=scope, manifest_payload=manifest_payload) + return manifest_payload class DocumentSectionPublisher: @@ -128,6 +129,7 @@ def ensure_section(self, section_path: str) -> DocumentSection: continue ancestor_section = DocumentSection( + section_id=f"sec_{uuid4().hex[:12]}", user_id=self._scope.user_id, namespace=self._scope.namespace, document_id=self._scope.document_id, @@ -141,7 +143,6 @@ def ensure_section(self, section_path: str) -> DocumentSection: summary=self._section_summaries.get(ancestor_path) or None, ) self._db.add(ancestor_section) - self._db.flush() self._sections_by_path[ancestor_path] = ancestor_section return self._sections_by_path[section_path] diff --git a/packages/shared-python/shared/services/retrieval/publication_models.py b/packages/shared-python/shared/services/retrieval/publication_models.py index 8c653212..0fdea6b1 100644 --- a/packages/shared-python/shared/services/retrieval/publication_models.py +++ b/packages/shared-python/shared/services/retrieval/publication_models.py @@ -1,6 +1,7 @@ from __future__ import annotations from dataclasses import dataclass +from typing import Any @dataclass(frozen=True) @@ -15,6 +16,7 @@ class PublishedDocumentState: namespace: str document_id: str | None skipped_all_duplicate: bool = False + manifest_payload: dict[str, Any] | None = None @dataclass(frozen=True) diff --git a/packages/shared-python/shared/services/retrieval/publication_service.py b/packages/shared-python/shared/services/retrieval/publication_service.py index 7ee5dec2..3cdd8ddc 100644 --- a/packages/shared-python/shared/services/retrieval/publication_service.py +++ b/packages/shared-python/shared/services/retrieval/publication_service.py @@ -23,6 +23,7 @@ from shared.models.schemas.retrieval_namespace import normalize_retrieval_namespace from shared.services.retrieval.graph.service import DocumentGraphService, GraphScope from shared.services.retrieval.namespace_map_snapshot import ( + patch_namespace_map_snapshot, remove_document_from_namespace_map_snapshot, ) from shared.services.retrieval.publication_content import ( @@ -81,6 +82,7 @@ def publish_document_state( job_result_id: str, chunks: list[dict[str, Any]], section_summaries: dict[str, str] | None = None, + update_namespace_snapshot: bool = True, ) -> PublishedDocumentState | None: job = db.execute(select(Job).where(Job.job_id == job_id)).scalar_one_or_none() if not job: @@ -93,6 +95,7 @@ def publish_document_state( job_result_id=job_result_id, chunks=chunks, section_summaries=section_summaries, + update_namespace_snapshot=update_namespace_snapshot, ) def _publish_document_state_for_job( @@ -103,6 +106,7 @@ def _publish_document_state_for_job( job_result_id: str, chunks: list[dict[str, Any]], section_summaries: dict[str, str] | None = None, + update_namespace_snapshot: bool = True, ) -> PublishedDocumentState | None: job_metadata = job.job_metadata or {} @@ -136,16 +140,6 @@ def _publish_document_state_for_job( Document.document_id == str(document_id) ) ).scalar_one_or_none() - namespaces_to_lock = {namespace} - if existing_namespace: - namespaces_to_lock.add(str(existing_namespace)) - for namespace_to_lock in sorted(namespaces_to_lock): - lock_namespace_generation( - db, - user_id=str(job.user_id), - namespace=namespace_to_lock, - ) - document = self._upsert_document_revision( db, job=job, @@ -172,7 +166,7 @@ def _publish_document_state_for_job( job_result_id=job_result_id, source_file_name=str(source_file_name) if source_file_name else None, ) - replace_document_revision_content( + manifest_payload = replace_document_revision_content( db, scope=scope, chunks=deduped_chunks, @@ -180,32 +174,61 @@ def _publish_document_state_for_job( ) db.flush() - if existing_namespace and str(existing_namespace) != scope.namespace: - remove_document_from_namespace_map_snapshot( + if update_namespace_snapshot: + self.update_namespace_snapshot( db, - user_id=scope.user_id, - namespace=str(existing_namespace), - document_id=document.document_id, + scope=scope, + manifest_payload=manifest_payload, + previous_namespace=( + str(existing_namespace) + if existing_namespace and str(existing_namespace) != scope.namespace + else None + ), ) - # A namespace move mutates both namespace snapshots. Advance the - # old namespace generation as well so request-scoped/process-local - # snapshot caches cannot reuse the pre-move generation. - advance_namespace_generation( - db, - user_id=scope.user_id, - namespace=str(existing_namespace), - ) - advance_namespace_generation( - db, - user_id=scope.user_id, - namespace=scope.namespace, - ) return PublishedDocumentState( user_id=str(job.user_id), namespace=namespace, document_id=document.document_id, + manifest_payload=manifest_payload, ) + def update_namespace_snapshot( + self, + db: Session, + *, + scope: DocumentPublicationScope, + manifest_payload: dict[str, Any], + previous_namespace: str | None = None, + ) -> None: + """Patch one namespace snapshot while holding only its short lock.""" + namespaces = [scope.namespace] + if previous_namespace: + namespaces.insert(0, previous_namespace) + for namespace in namespaces: + lock_namespace_generation( + db, + user_id=scope.user_id, + namespace=namespace, + ) + if namespace == scope.namespace: + patch_namespace_map_snapshot( + db, + scope=scope, + manifest_payload=manifest_payload, + ) + else: + remove_document_from_namespace_map_snapshot( + db, + user_id=scope.user_id, + namespace=namespace, + document_id=scope.document_id, + ) + advance_namespace_generation( + db, + user_id=scope.user_id, + namespace=namespace, + ) + def _upsert_document_revision( self, db: Session, diff --git a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py index 9a1561c9..790dc92d 100644 --- a/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py +++ b/packages/shared-python/shared/services/retrieval/search/map_unit_discovery.py @@ -451,9 +451,13 @@ async def map_unit_discovery( has_index_storage_mismatch = indexed_unit_count != int( actual_unit_count or 0 ) or indexed_token_count != int(actual_token_count or 0) + # Equality against indexes.unit_count is only valid when unit_rows is the + # full inventory of those revisions. Image/table type filters, signal + # paths, and exclude_sections load a subset, so require the looser + # undercount check used for unfiltered token projection. has_index_unit_count_mismatch = ( indexed_unit_count < len(unit_rows) - if is_unfiltered_scope + if is_unfiltered_scope or type_clause or signal_paths or exclude_sections else indexed_unit_count != len(unit_rows) ) is_index_format_incompatible = any( @@ -488,8 +492,9 @@ async def map_unit_discovery( _token_count, ) in index_parts ) + has_revision_coverage_mismatch = len(index_parts) != len(expected_revisions) has_unusable_index = ( - len(index_parts) != len(expected_revisions) + has_revision_coverage_mismatch or has_index_unit_count_mismatch or has_index_storage_mismatch or is_index_format_incompatible @@ -505,11 +510,24 @@ async def map_unit_discovery( ) except Exception as exc: logger.warning("retrieval index readiness publish failed: %s", exc) + if has_revision_coverage_mismatch: + unusable_reason = "revision_coverage" + elif is_index_format_incompatible: + unusable_reason = "format" + elif has_index_storage_mismatch: + unusable_reason = "storage" + else: + unusable_reason = "unit_count_mismatch" + chunk_types_label = ( + ",".join(sorted(chunk_types)) if chunk_types else "none" + ) raise RuntimeError( "retrieval map-unit index is incomplete or incompatible " - f"(user_id={user_id} namespace={namespace} " + f"(reason={unusable_reason} user_id={user_id} namespace={namespace} " f"expected_revisions={len(expected_revisions)} " - f"indexed_revisions={len(index_parts)})" + f"indexed_revisions={len(index_parts)} " + f"indexed_unit_count={indexed_unit_count} unit_rows={len(unit_rows)} " + f"unfiltered={is_unfiltered_scope} chunk_types={chunk_types_label})" ) if has_incomplete_index_statistics: logger.warning( diff --git a/packages/shared-python/shared/services/storage/result_storage.py b/packages/shared-python/shared/services/storage/result_storage.py index 51d3ab3c..9a5643fb 100644 --- a/packages/shared-python/shared/services/storage/result_storage.py +++ b/packages/shared-python/shared/services/storage/result_storage.py @@ -1,13 +1,16 @@ from __future__ import annotations import os +import threading from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from pathlib import Path from typing import Protocol from loguru import logger +from shared.core.config import settings from shared.services.storage.job_file_storage import JobFileStorage from shared.services.storage.storage_adapter import StorageAdapter @@ -15,6 +18,9 @@ _EXCLUDED_DIR_NAMES = {"tmp", "temp", "__pycache__"} _CLIENT_ARTIFACT_DIRS = {"images", "tables", "page_pdfs", "page_citation_assets"} _INTERNAL_RAW_FILES = {"source.pdf"} +_MAX_RESULT_UPLOAD_CONCURRENCY = 64 +_upload_executor_lock = threading.Lock() +_upload_executors: dict[int, ThreadPoolExecutor] = {} @dataclass(frozen=True) @@ -58,12 +64,17 @@ def __init__( *, results_bucket: str | None = None, storage_adapter: StorageAdapter | None = None, + upload_concurrency: int | None = None, ) -> None: self._job_file_storage = JobFileStorage( storage_adapter=storage_adapter, results_bucket=results_bucket, ) self.results_bucket = self._job_file_storage.results_bucket + configured_concurrency = getattr(settings, "RESULT_UPLOAD_CONCURRENCY", 20) + self.upload_concurrency = self._validate_upload_concurrency( + configured_concurrency if upload_concurrency is None else upload_concurrency + ) def build_zip_key(self, *, job_id: str) -> str: return self._job_file_storage.build_result_zip_key(job_id=job_id) @@ -109,19 +120,33 @@ def upload( ) self._cleanup_file(zip_path) - raw_files: dict[str, str] = {} artifact_ref_filter = self._normalize_artifact_refs(artifact_refs) - for file_path in self._iter_raw_files(result_path): - relative_path = file_path.relative_to(result_path).as_posix() - if artifact_ref_filter is not None and relative_path not in artifact_ref_filter: - continue - raw_key = self.build_raw_key(job_id=job_id, relative_path=relative_path) - self._job_file_storage.upload_local_file( - str(file_path), - raw_key, - bucket=self.results_bucket, + upload_items = [ + ( + file_path, + file_path.relative_to(result_path).as_posix(), + ) + for file_path in self._iter_raw_files(result_path) + if artifact_ref_filter is None + or file_path.relative_to(result_path).as_posix() in artifact_ref_filter + ] + executor = self._get_upload_executor() + futures = [ + executor.submit( + self._upload_raw_file, + job_id=job_id, + file_path=file_path, + relative_path=relative_path, ) - raw_files[relative_path] = raw_key + for file_path, relative_path in upload_items + ] + try: + uploaded_items = [future.result() for future in futures] + except Exception: + for future in futures: + future.cancel() + raise + raw_files = dict(uploaded_items) return UploadedResultBundle( zip_key=zip_key, @@ -129,6 +154,40 @@ def upload( raw_files=raw_files, ) + def _upload_raw_file( + self, + *, + job_id: str, + file_path: Path, + relative_path: str, + ) -> tuple[str, str]: + raw_key = self.build_raw_key(job_id=job_id, relative_path=relative_path) + self._job_file_storage.upload_local_file( + str(file_path), + raw_key, + bucket=self.results_bucket, + ) + return relative_path, raw_key + + def _validate_upload_concurrency(self, upload_concurrency: int) -> int: + if not 1 <= upload_concurrency <= _MAX_RESULT_UPLOAD_CONCURRENCY: + raise ValueError( + "upload_concurrency must be between 1 and " + f"{_MAX_RESULT_UPLOAD_CONCURRENCY}" + ) + return upload_concurrency + + def _get_upload_executor(self) -> ThreadPoolExecutor: + with _upload_executor_lock: + executor = _upload_executors.get(self.upload_concurrency) + if executor is None: + executor = ThreadPoolExecutor( + max_workers=self.upload_concurrency, + thread_name_prefix="result-upload", + ) + _upload_executors[self.upload_concurrency] = executor + return executor + def generate_url(self, *, storage_key: str, expires_in: int = 3600) -> str | None: return self._job_file_storage.generate_download_url( storage_key, diff --git a/packages/shared-python/shared/tests/test_asset_inline.py b/packages/shared-python/shared/tests/test_asset_inline.py index d8345c24..75b09c04 100644 --- a/packages/shared-python/shared/tests/test_asset_inline.py +++ b/packages/shared-python/shared/tests/test_asset_inline.py @@ -101,6 +101,66 @@ async def test_assemble_inserts_table_at_placeholder() -> None: assert "SHOULD NOT LEAK" not in content +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("asset_type", "file_path", "placeholder", "display_marker"), + [ + ("image", "images/a.png", "[images/a.png]", "[Image: images/a.png]"), + ("table", "tables/a.html", "[tables/a.html]", "[Table: tables/a.html]"), + ], +) +async def test_asset_type_filter_keeps_body_that_connects_to_requested_asset( + monkeypatch, + asset_type: str, + file_path: str, + placeholder: str, + display_marker: str, +) -> None: + body_row = { + "chunk_id": "text-1", + "chunk_type": "text", + "content": f"查看 {placeholder}", + "chunk_metadata": { + "connect_to": [ + { + "target": "asset-1", + "relation": "embeds", + "ref": placeholder, + } + ] + }, + } + asset_row = { + "chunk_id": "asset-1", + "chunk_type": asset_type, + "content": "资产说明" if asset_type == "image" else "
", + "file_path": file_path, + "chunk_metadata": {"summary": "资产说明"}, + } + + async def hydrate_connected_rows(**_kwargs: object) -> list[dict[str, object]]: + return [asset_row] + + monkeypatch.setattr( + "shared.services.retrieval.hydration.result_assembly.hydrate_connected_target_rows", + hydrate_connected_rows, + ) + + assembled = await assemble_retrieval_results( + rows=[ + body_row, + {"chunk_id": "text-2", "chunk_type": "text", "content": "无图"}, + ], + exclude_document_ids=[], + exclude_sections=[], + allowed_chunk_types={asset_type}, + ) + + assert [row["chunk_id"] for row in assembled] == ["text-1"] + assert display_marker in assembled[0]["content"] + assert "资产说明" in assembled[0]["content"] + + def test_node_unit_span_inlines_section_assets() -> None: provider = KnowhereProvider( doc_id="doc-1", diff --git a/packages/shared-python/shared/tests/test_storage_config_contract.py b/packages/shared-python/shared/tests/test_storage_config_contract.py index d7dce29e..8931688f 100644 --- a/packages/shared-python/shared/tests/test_storage_config_contract.py +++ b/packages/shared-python/shared/tests/test_storage_config_contract.py @@ -43,6 +43,13 @@ def test_aws_s3_uses_default_credential_chain_when_keys_are_empty( assert client_arguments["region_name"] == "us-east-1" assert "aws_access_key_id" not in client_arguments assert "aws_secret_access_key" not in client_arguments + assert client_arguments["config"].max_pool_connections == 20 + + +def test_result_upload_concurrency_defaults_to_twenty() -> None: + config: StorageConfig = create_storage_config() + + assert config.RESULT_UPLOAD_CONCURRENCY == 20 def test_aws_s3_passes_complete_explicit_credentials(