diff --git a/apps/api/app/api/v2/api_v2.py b/apps/api/app/api/v2/api_v2.py index f558f8b8..37174339 100644 --- a/apps/api/app/api/v2/api_v2.py +++ b/apps/api/app/api/v2/api_v2.py @@ -1,5 +1,6 @@ """API v2 route registry.""" +from app.api.v1.routes import documents as v1_documents from app.api.v2.routes import documents, jobs, retrieval from fastapi import APIRouter @@ -8,5 +9,6 @@ api_router.include_router(jobs.router, prefix="/jobs", tags=["Jobs"]) api_router.include_router(retrieval.router, prefix="/retrieval", tags=["Retrieval"]) api_router.include_router(documents.router, prefix="/documents", tags=["Documents"]) +api_router.include_router(v1_documents.router, prefix="/documents", tags=["Documents"]) __all__ = ["api_router"] diff --git a/apps/api/app/api/v2/routes/documents.py b/apps/api/app/api/v2/routes/documents.py index 636002f4..7bde02b5 100644 --- a/apps/api/app/api/v2/routes/documents.py +++ b/apps/api/app/api/v2/routes/documents.py @@ -1,5 +1,40 @@ """Documents API v2 routes.""" -from app.api.v1.routes.documents import router +from __future__ import annotations + +from typing import Any + +from app.api.dependencies.current_user import with_current_user +from app.services.documents.lifecycle_service import DocumentService +from app.services.rate_limit.data_structures import CurrentUser +from fastapi import APIRouter, Depends +from sqlalchemy.ext.asyncio import AsyncSession + +from shared.core.database import get_db +from shared.core.exceptions.domain_exceptions import NotFoundException + +router = APIRouter(tags=["Documents"]) + +_document_service = DocumentService() + + +@router.get("/{document_id}/files/page-citation-source") +async def get_document_page_citation_source( + document_id: str, + current_user: CurrentUser = Depends(with_current_user), + db: AsyncSession = Depends(get_db), +) -> dict[str, Any]: + response = await _document_service.get_document_page_citation_source( + db, + user_id=current_user.user_id, + document_id=document_id, + ) + if response is None: + raise NotFoundException( + resource="Document page citation source", + resource_id=document_id, + internal_message="Document page citation source not found", + ) + return response __all__ = ["router"] diff --git a/apps/api/app/repositories/document_repository.py b/apps/api/app/repositories/document_repository.py index c25406db..865f1230 100644 --- a/apps/api/app/repositories/document_repository.py +++ b/apps/api/app/repositories/document_repository.py @@ -11,9 +11,11 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.models.database.document import Document, DocumentChunk, DocumentSection +from shared.models.database.job import Job from shared.models.database.job_result import JobResult DocumentChunkRow = tuple[DocumentChunk, DocumentSection | None, JobResult] +DocumentJobRevisionRow = tuple[Document, JobResult, Job] class DocumentRepository: @@ -66,6 +68,29 @@ async def get_document( ) return result.scalar_one_or_none() + async def get_current_document_job_revision( + self, + db: AsyncSession, + *, + document_id: str, + user_id: str, + ) -> DocumentJobRevisionRow | None: + stmt = ( + select(Document, JobResult, Job) + .join(JobResult, JobResult.id == Document.current_job_result_id) + .join(Job, Job.job_id == JobResult.job_id) + .where(Document.document_id == document_id) + .where(Document.user_id == user_id) + .where(JobResult.document_id == Document.document_id) + .where(Job.user_id == user_id) + .where(Document.status != "archived") + .limit(1) + ) + + result = await db.execute(stmt) + row = result.first() + return cast(DocumentJobRevisionRow | None, row) + async def archive_document( self, db: AsyncSession, diff --git a/apps/api/app/services/documents/lifecycle_service.py b/apps/api/app/services/documents/lifecycle_service.py index 7e262542..1bd5ae3a 100644 --- a/apps/api/app/services/documents/lifecycle_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -3,7 +3,7 @@ from __future__ import annotations import math -from datetime import datetime +from datetime import datetime, timedelta, timezone from typing import Any from app.repositories.document_repository import DocumentRepository @@ -19,6 +19,10 @@ _DOCUMENT_CHUNK_ASSET_URL_EXPIRES_SECONDS = 7 * 24 * 60 * 60 _MEDIA_CHUNK_TYPES = frozenset({"image", "table"}) +_PAGE_CITATION_SOURCE_EXPIRES_SECONDS = 60 * 60 +_PAGE_CITATION_SOURCE_FILE_NAME = "source.pdf" +_PAGE_CITATION_SOURCE_VARIANT = "normalized_pdf" +_PAGE_MEMORY_PARSE_TRACK = "page_memory" def _datetime_payload(value: datetime | None) -> str | None: @@ -53,6 +57,89 @@ def _document_chunk_asset_url( return None +def _document_page_assets( + *, + metadata: dict[str, Any] | None, + job_id: str | None, + include_asset_urls: bool, + result_storage: ResultStorage | None, +) -> list[dict[str, Any]]: + if not isinstance(metadata, dict): + return [] + raw_assets = metadata.get("page_assets") + if not isinstance(raw_assets, list): + return [] + + page_assets: list[dict[str, Any]] = [] + for raw_asset in raw_assets: + if not isinstance(raw_asset, dict): + continue + asset = _normalize_page_asset(raw_asset) + if asset is None: + continue + if include_asset_urls and job_id and result_storage is not None: + asset_url = _page_asset_url( + job_id=job_id, + artifact_ref=asset["artifact_ref"], + result_storage=result_storage, + ) + if asset_url: + asset["asset_url"] = asset_url + page_assets.append(asset) + return page_assets + + +def _normalize_page_asset(raw_asset: dict[str, Any]) -> dict[str, Any] | None: + page_num = _positive_int(raw_asset.get("page_num")) + artifact_ref = str(raw_asset.get("artifact_ref") or "").strip() + content_type = str(raw_asset.get("content_type") or "").strip() + source = str(raw_asset.get("source") or "").strip() + if page_num is None or not artifact_ref or not content_type or not source: + return None + + asset: dict[str, Any] = { + "page_num": page_num, + "artifact_ref": artifact_ref, + "content_type": content_type, + "source": source, + } + if (asset_url := str(raw_asset.get("asset_url") or "").strip()): + asset["asset_url"] = asset_url + if (width := _positive_int(raw_asset.get("width"))) is not None: + asset["width"] = width + if (height := _positive_int(raw_asset.get("height"))) is not None: + asset["height"] = height + return asset + + +def _page_asset_url( + *, + job_id: str, + artifact_ref: str, + result_storage: ResultStorage, +) -> str | None: + normalized_ref = result_storage.normalize_artifact_ref(artifact_ref) + if not normalized_ref or not normalized_ref.startswith("page_citation_assets/"): + return None + try: + return result_storage.generate_artifact_url( + job_id=job_id, + artifact_ref=normalized_ref, + expires_in=_DOCUMENT_CHUNK_ASSET_URL_EXPIRES_SECONDS, + ) + except Exception as exc: + logger.warning(f"Failed to generate page citation asset URL (ignored): {exc}") + return None + + +def _positive_int(value: Any) -> int | None: + try: + number = int(value) + except (TypeError, ValueError): + return None + return number if number > 0 else None + + def document_payload(document) -> dict[str, Any]: return { "document_id": document.document_id, @@ -75,9 +162,11 @@ def __init__( *, repository: DocumentRepository | None = None, graph_service: DocumentGraphService | None = None, + result_storage: ResultStorage | None = None, ) -> None: self._repository = repository or DocumentRepository() self._graph_service = graph_service or DocumentGraphService() + self._result_storage = result_storage async def list_documents( self, @@ -246,6 +335,55 @@ async def get_document( return None return document_payload(document) + async def get_document_page_citation_source( + self, + db: AsyncSession, + *, + user_id: str, + document_id: str, + ) -> dict[str, Any] | None: + row = await self._repository.get_current_document_job_revision( + db, + user_id=user_id, + document_id=document_id, + ) + if row is None: + return None + + document, job_result, job = row + if document.parse_track != _PAGE_MEMORY_PARSE_TRACK: + return None + + result_storage = self._result_storage or get_result_storage() + if not result_storage.verify_raw_exists( + job_id=job_result.job_id, + relative_path=_PAGE_CITATION_SOURCE_FILE_NAME, + ): + return None + + source_url = result_storage.generate_raw_file_url( + job_id=job_result.job_id, + relative_path=_PAGE_CITATION_SOURCE_FILE_NAME, + expires_in=_PAGE_CITATION_SOURCE_EXPIRES_SECONDS, + ) + if not source_url: + return None + + expires_at = datetime.now(timezone.utc) + timedelta( + seconds=_PAGE_CITATION_SOURCE_EXPIRES_SECONDS, + ) + return { + "document_id": document.document_id, + "namespace": document.namespace, + "job_id": job.job_id, + "job_result_id": job_result.id, + "variant": _PAGE_CITATION_SOURCE_VARIANT, + "file_name": _PAGE_CITATION_SOURCE_FILE_NAME, + "content_type": "application/pdf", + "url": source_url, + "expires_at": expires_at.isoformat(), + } + def _chunk_payload( self, *, @@ -257,7 +395,17 @@ def _chunk_payload( ) -> dict[str, Any]: chunk_type = _normalize_chunk_type(chunk.chunk_type) file_path = chunk.file_path - return { + raw_metadata = chunk.chunk_metadata or {} + page_assets = _document_page_assets( + metadata=raw_metadata, + job_id=job_id, + include_asset_urls=include_asset_urls, + result_storage=result_storage, + ) + metadata = dict(raw_metadata) + if include_asset_urls and page_assets: + metadata["page_assets"] = page_assets + payload = { "id": chunk.id, "chunk_id": chunk.chunk_id, "chunk_type": chunk_type, @@ -267,7 +415,7 @@ def _chunk_payload( "source_chunk_path": chunk.source_chunk_path, "file_path": file_path, "sort_order": chunk.sort_order, - "metadata": chunk.chunk_metadata, + "metadata": metadata, "asset_url": _document_chunk_asset_url( chunk_type=chunk_type, job_id=job_id, @@ -277,6 +425,7 @@ def _chunk_payload( ), "created_at": _datetime_payload(chunk.created_at), } + return payload async def archive_document( self, diff --git a/apps/api/tests/contract/test_agentic_discovery_selection_contract.py b/apps/api/tests/contract/test_agentic_discovery_selection_contract.py index a27dc654..09858a1d 100644 --- a/apps/api/tests/contract/test_agentic_discovery_selection_contract.py +++ b/apps/api/tests/contract/test_agentic_discovery_selection_contract.py @@ -1,4 +1,13 @@ from shared.services.retrieval.agentic.navigation.actions import build_legal_actions +from shared.services.retrieval.agentic.navigation.state import RejectionRecord + + +def _rejected(paths: dict[str, str]) -> dict[str, RejectionRecord]: + """Build a rejection ledger from {path: reason}.""" + return { + path: RejectionRecord(path=path, reason=reason, step=1, detail="") + for path, reason in paths.items() + } def test_discovery_hint_is_projected_as_collect_action() -> None: @@ -14,8 +23,7 @@ def test_discovery_hint_is_projected_as_collect_action() -> None: "chunk_type": "text", } ], - rejected_paths=set(), - rejected_collect_paths=set(), + rejected={}, total_images=0, total_tables=0, budget_snapshot=None, @@ -47,8 +55,7 @@ def test_discovery_hint_under_collected_path_is_not_repeated() -> None: "discovery_score": 0.82, } ], - rejected_paths=set(), - rejected_collect_paths=set(), + rejected={}, total_images=0, total_tables=0, budget_snapshot=None, @@ -57,7 +64,8 @@ def test_discovery_hint_under_collected_path_is_not_repeated() -> None: assert action_set.collect == [] -def test_discovery_hint_under_rejected_collect_path_is_not_repeated() -> None: +def test_discovery_hint_under_tool_adjudicated_path_is_not_repeated() -> None: + """tool_adjudicated rejections are not revived by discovery this round.""" action_set = build_legal_actions( items=[], current_scope=None, @@ -69,11 +77,239 @@ def test_discovery_hint_under_rejected_collect_path_is_not_repeated() -> None: "discovery_score": 0.7, } ], - rejected_paths=set(), - rejected_collect_paths={"1、2016：机构行为助推行情演绎"}, + rejected=_rejected({ + "1、2016：机构行为助推行情演绎": "tool_adjudicated", + }), total_images=0, total_tables=0, budget_snapshot=None, ) assert action_set.collect == [] + + +# ─── NavigationState ledger (Phase 0) ─────────────────────────────────────── + + +def test_mark_rejected_collect_records_tool_adjudicated_reason() -> None: + from shared.services.retrieval.agentic.navigation.state import NavigationState + + state = NavigationState( + document_id="d1", + document_name="doc.pdf", + job_result_id="j1", + ) + state.mark_rejected_collect("Chapter 1", step=3, detail="no matching asset") + + assert "Chapter 1" in state.rejected + record = state.rejected["Chapter 1"] + assert record.reason == "tool_adjudicated" + assert record.step == 3 + assert record.detail == "no matching asset" + + +def test_mark_rejected_if_unproductive_records_navigational_abandon() -> None: + from shared.services.retrieval.agentic.navigation.state import NavigationState + + state = NavigationState( + document_id="d1", + document_name="doc.pdf", + job_result_id="j1", + ) + state.mark_rejected_if_unproductive("Chapter 2", step=5, detail="back_from_unproductive") + + assert state.rejected["Chapter 2"].reason == "navigational_abandon" + + +def test_tool_adjudicated_overrides_weak_abandon_record() -> None: + from shared.services.retrieval.agentic.navigation.state import NavigationState + + state = NavigationState( + document_id="d1", + document_name="doc.pdf", + job_result_id="j1", + ) + state.mark_rejected_if_unproductive("Chapter 3", step=2) + state.mark_rejected_collect("Chapter 3", step=4, detail="asset mismatch") + + # Stronger reason wins. + assert state.rejected["Chapter 3"].reason == "tool_adjudicated" + assert state.rejected["Chapter 3"].step == 4 + + +def test_weak_abandon_does_not_overwrite_strong_record() -> None: + from shared.services.retrieval.agentic.navigation.state import NavigationState + + state = NavigationState( + document_id="d1", + document_name="doc.pdf", + job_result_id="j1", + ) + state.mark_rejected_collect("Chapter 4", step=1) + state.mark_rejected_if_unproductive("Chapter 4", step=5) + + assert state.rejected["Chapter 4"].reason == "tool_adjudicated" + assert state.rejected["Chapter 4"].step == 1 + + +def test_coverage_helpers_derive_from_collected_paths() -> None: + from shared.services.retrieval.agentic.navigation.state import NavigationState + + state = NavigationState( + document_id="d1", + document_name="doc.pdf", + job_result_id="j1", + ) + state.add_collected( + {"path": "A", "hydrate_mode": "chunks", "confidence": 0.9}, + step=1, + scope_context=None, + ) + state.add_collected( + {"path": "B", "hydrate_mode": "outline", "confidence": 0.6}, + step=2, + scope_context=None, + ) + # A path upgraded from outline to full counts as covered, not outline. + state.add_collected( + {"path": "C", "hydrate_mode": "outline", "confidence": 0.5}, + step=3, + scope_context=None, + ) + state.add_collected( + {"path": "C", "hydrate_mode": "chunks", "confidence": 0.8}, + step=4, + scope_context=None, + ) + + assert state.covered_paths() == {"A", "C"} + assert state.outline_paths() == {"B"} + + +def test_snapshot_delta_records_rejection_reasons() -> None: + from shared.services.retrieval.agentic.navigation.state import NavigationState + + state = NavigationState( + document_id="d1", + document_name="doc.pdf", + job_result_id="j1", + ) + state.step_count = 2 + state.mark_rejected_if_unproductive("X", step=2) + + state.step_count = 3 + state.mark_rejected_collect("Y", step=3, detail="asset mismatch") + + delta = state.snapshot_delta( + before_scope=None, + expanded_before=set(), + rejected_before={}, + collected_before_count=0, + ) + rejected_added = {item["path"]: item["reason"] for item in delta["rejected_added"]} + assert rejected_added == {"X": "navigational_abandon", "Y": "tool_adjudicated"} + + +def test_rejected_paths_with_reason_partitions_by_label() -> None: + from shared.services.retrieval.agentic.navigation.state import NavigationState + + state = NavigationState( + document_id="d1", + document_name="doc.pdf", + job_result_id="j1", + ) + state.mark_rejected_collect("A", step=1) + state.mark_rejected_if_unproductive("B", step=1) + state.mark_rejected_collect("C", step=1) + + assert state.rejected_paths_with_reason("tool_adjudicated") == {"A", "C"} + assert state.rejected_paths_with_reason("navigational_abandon") == {"B"} + + +# ─── Reason-aware action filtering (T7-style regression) ──────────────────── + + +def test_tool_adjudicated_rejection_blocks_collect_even_with_discovery() -> None: + """A tool-adjudicated path stays out of COLLECT even when discovery hints it.""" + action_set = build_legal_actions( + items=[], + current_scope=None, + collected_paths=[], + expanded_scopes=set(), + discovery_hints=[{"section_path": "X", "discovery_score": 0.95}], + rejected=_rejected({"X": "tool_adjudicated"}), + total_images=0, + total_tables=0, + budget_snapshot=None, + ) + assert action_set.collect == [] + + +def test_navigational_abandon_is_revived_by_discovery_for_collect() -> None: + """A soft-abandoned path CAN still be COLLECTed when discovery signals it.""" + action_set = build_legal_actions( + items=[], + current_scope=None, + collected_paths=[], + expanded_scopes=set(), + discovery_hints=[{"section_path": "Y", "discovery_score": 0.9}], + rejected=_rejected({"Y": "navigational_abandon"}), + total_images=0, + total_tables=0, + budget_snapshot=None, + ) + assert any(action.path == "Y" for action in action_set.collect) + + +def test_navigational_abandon_suppresses_expand_without_discovery_signal() -> None: + """EXPAND is suppressed for soft-abandoned scopes lacking any discovery signal.""" + items = [{"path": "Z", "level": 1, "is_leaf": False, "chunk_count": 5}] + action_set = build_legal_actions( + items=items, + current_scope=None, + collected_paths=[], + expanded_scopes=set(), + discovery_hints=[], + rejected=_rejected({"Z": "navigational_abandon"}), + total_images=0, + total_tables=0, + budget_snapshot=None, + ) + assert any(action.path == "Z" for action in action_set.collect) + assert not any(action.path == "Z" for action in action_set.expand) + + +def test_navigational_abandon_revives_expand_with_discovery_signal() -> None: + """EXPAND is offered for soft-abandoned scopes when a discovery signal exists.""" + items = [{"path": "Z", "level": 1, "is_leaf": False, "chunk_count": 5}] + action_set = build_legal_actions( + items=items, + current_scope=None, + collected_paths=[], + expanded_scopes=set(), + discovery_hints=[{"section_path": "Z / child", "discovery_score": 0.7}], + rejected=_rejected({"Z": "navigational_abandon"}), + total_images=0, + total_tables=0, + budget_snapshot=None, + ) + assert any(action.path == "Z" for action in action_set.expand) + + +def test_covered_path_excluded_from_actions() -> None: + """Regression: a path already collected as full evidence is not re-offered.""" + items = [{"path": "A", "level": 1, "is_leaf": False, "chunk_count": 3}] + action_set = build_legal_actions( + items=items, + current_scope=None, + collected_paths=[{"path": "A", "hydrate_mode": "chunks"}], + expanded_scopes=set(), + discovery_hints=[{"section_path": "A", "discovery_score": 0.9}], + rejected={}, + total_images=0, + total_tables=0, + budget_snapshot=None, + ) + assert not any(action.path == "A" for action in action_set.collect) + assert not any(action.path == "A" for action in action_set.expand) + diff --git a/apps/api/tests/contract/test_documents_contract.py b/apps/api/tests/contract/test_documents_contract.py index 92a80bc2..6b8e33b0 100644 --- a/apps/api/tests/contract/test_documents_contract.py +++ b/apps/api/tests/contract/test_documents_contract.py @@ -2,6 +2,7 @@ from collections.abc import Callable from contextlib import AbstractAsyncContextManager from datetime import datetime, timedelta, timezone +from pathlib import Path from typing import cast from uuid import uuid4 @@ -258,6 +259,8 @@ async def _insert_document_revision_with_chunks( namespace: str = "contract-documents", user_id: str = "local-dev-user", source_file_name: str = "contract-chunks.pdf", + parse_track: str = "chunk", + status: str = "active", ) -> dict[str, str]: engine = await _create_contract_engine() timestamp = datetime.now(timezone.utc).replace(tzinfo=None) @@ -284,23 +287,25 @@ async def _insert_document_revision_with_chunks( :document_id, :user_id, :namespace, - 'active', + :status, NULL, :source_file_name, :parse_track, :created_at, :updated_at, - NULL + :archived_at ) """), { "document_id": document_id, "user_id": user_id, "namespace": namespace, + "status": status, "source_file_name": source_file_name, - "parse_track": "chunk", + "parse_track": parse_track, "created_at": timestamp, "updated_at": timestamp, + "archived_at": timestamp if status == "archived" else None, }, ) await connection.execute( @@ -498,6 +503,36 @@ async def _insert_document_revision_with_chunks( } +def _upload_page_citation_source(*, job_id: str) -> None: + from shared.services.storage.result_storage import JobResultStorage + + source_pdf_path = Path("/tmp") / f"knowhere-contract-source-{job_id}.pdf" + source_pdf_path.write_bytes(b"%PDF-1.4\n%contract page citation source\n") + try: + JobResultStorage().upload_raw_file( + job_id=job_id, + relative_path="source.pdf", + local_file_path=str(source_pdf_path), + ) + finally: + source_pdf_path.unlink(missing_ok=True) + + +def _upload_page_citation_asset(*, job_id: str, artifact_ref: str) -> None: + from shared.services.storage.result_storage import JobResultStorage + + asset_path = Path("/tmp") / f"knowhere-contract-page-asset-{uuid4().hex}.png" + asset_path.write_bytes(b"\x89PNG\r\n\x1a\ncontract page citation asset\n") + try: + JobResultStorage().upload_raw_file( + job_id=job_id, + relative_path=artifact_ref, + local_file_path=str(asset_path), + ) + finally: + asset_path.unlink(missing_ok=True) + + @pytest.mark.asyncio async def test_should_list_only_the_authenticated_users_documents_for_the_effective_namespace( developer_api_client_factory: Callable[ @@ -681,6 +716,178 @@ async def test_should_return_not_found_when_requesting_a_missing_document( } +@pytest.mark.asyncio +async def test_should_return_page_citation_source_for_page_memory_document( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + document_id = f"doc_{uuid4().hex[:12]}" + + async with developer_api_client_factory() as api_client: + revision = await _insert_document_revision_with_chunks( + document_id=document_id, + parse_track="page_memory", + chunks=[ + { + "id": f"dchk_{uuid4().hex[:12]}", + "chunk_id": "page-memory-text", + "chunk_type": "text", + "content": "Page memory content", + "source_chunk_path": "Page 1", + "metadata": {"page_nums": [1]}, + } + ], + ) + _upload_page_citation_source(job_id=revision["job_id"]) + response = await api_client.get( + f"/api/v2/documents/{document_id}/files/page-citation-source" + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + assert response_json["document_id"] == document_id + assert response_json["namespace"] == "contract-documents" + assert response_json["job_id"] == revision["job_id"] + assert response_json["job_result_id"] == revision["job_result_id"] + assert response_json["variant"] == "normalized_pdf" + assert response_json["file_name"] == "source.pdf" + assert response_json["content_type"] == "application/pdf" + assert response_json["url"] == ( + "filesystem://knowhere-test-results/" + f"results/{revision['job_id']}/source.pdf" + "?method=GET&expires_in=3600" + ) + assert response_json["expires_at"] + + +@pytest.mark.asyncio +async def test_should_return_not_found_for_chunk_document_page_citation_source( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + document_id = f"doc_{uuid4().hex[:12]}" + + async with developer_api_client_factory() as api_client: + revision = await _insert_document_revision_with_chunks( + document_id=document_id, + chunks=[ + { + "id": f"dchk_{uuid4().hex[:12]}", + "chunk_id": "chunk-text", + "chunk_type": "text", + "content": "Chunk content", + "source_chunk_path": "Chunk 1", + "metadata": {"page_nums": []}, + } + ], + ) + _upload_page_citation_source(job_id=revision["job_id"]) + response = await api_client.get( + f"/api/v2/documents/{document_id}/files/page-citation-source" + ) + + assert response.status_code == 404 + + +@pytest.mark.asyncio +async def test_should_return_not_found_when_page_citation_source_is_missing( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + document_id = f"doc_{uuid4().hex[:12]}" + + async with developer_api_client_factory() as api_client: + await _insert_document_revision_with_chunks( + document_id=document_id, + parse_track="page_memory", + chunks=[ + { + "id": f"dchk_{uuid4().hex[:12]}", + "chunk_id": "page-memory-text", + "chunk_type": "text", + "content": "Page memory content", + "source_chunk_path": "Page 1", + "metadata": {"page_nums": [1]}, + } + ], + ) + response = await api_client.get( + f"/api/v2/documents/{document_id}/files/page-citation-source" + ) + + assert response.status_code == 404 + + +@pytest.mark.asyncio +async def test_should_return_not_found_for_archived_page_citation_source( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + document_id = f"doc_{uuid4().hex[:12]}" + + async with developer_api_client_factory() as api_client: + revision = await _insert_document_revision_with_chunks( + document_id=document_id, + parse_track="page_memory", + status="archived", + chunks=[ + { + "id": f"dchk_{uuid4().hex[:12]}", + "chunk_id": "page-memory-text", + "chunk_type": "text", + "content": "Page memory content", + "source_chunk_path": "Page 1", + "metadata": {"page_nums": [1]}, + } + ], + ) + _upload_page_citation_source(job_id=revision["job_id"]) + response = await api_client.get( + f"/api/v2/documents/{document_id}/files/page-citation-source" + ) + + assert response.status_code == 404 + + +@pytest.mark.asyncio +async def test_should_return_not_found_for_other_users_page_citation_source( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + document_id = f"doc_{uuid4().hex[:12]}" + other_user_id = f"contract-user-{uuid4().hex[:12]}" + + async with developer_api_client_factory() as api_client: + await ContractDatabase.insert_user(user_id=other_user_id) + revision = await _insert_document_revision_with_chunks( + document_id=document_id, + user_id=other_user_id, + parse_track="page_memory", + chunks=[ + { + "id": f"dchk_{uuid4().hex[:12]}", + "chunk_id": "page-memory-text", + "chunk_type": "text", + "content": "Page memory content", + "source_chunk_path": "Page 1", + "metadata": {"page_nums": [1]}, + } + ], + ) + _upload_page_citation_source(job_id=revision["job_id"]) + response = await api_client.get( + f"/api/v2/documents/{document_id}/files/page-citation-source" + ) + + assert response.status_code == 404 + + @pytest.mark.asyncio async def test_should_list_current_document_chunks_by_document_id( developer_api_client_factory: Callable[ @@ -818,6 +1025,96 @@ async def test_should_include_media_asset_urls_in_document_chunk_list_when_reque assert default_chunks[1]["asset_url"] is None +@pytest.mark.asyncio +async def test_should_include_page_citation_asset_urls_in_document_chunk_metadata_when_requested( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + document_id = f"doc_{uuid4().hex[:12]}" + page_chunk_id = f"dchk_{uuid4().hex[:12]}" + page_asset_ref = "page_citation_assets/page-4.png" + + async with developer_api_client_factory() as api_client: + revision = await _insert_document_revision_with_chunks( + document_id=document_id, + parse_track="page_memory", + chunks=[ + { + "id": page_chunk_id, + "chunk_id": "page-node-4", + "chunk_type": "page", + "content": "Page node content", + "source_chunk_path": "Chapter 1/Page 4", + "metadata": { + "summary": "Page node summary", + "page_nums": [4], + "page_assets": [ + { + "page_num": 4, + "artifact_ref": page_asset_ref, + "content_type": "image/png", + "width": 1200, + "height": 1800, + "source": "knowhere-rendered-page-citation-source", + } + ], + }, + }, + ], + ) + _upload_page_citation_asset( + job_id=revision["job_id"], + artifact_ref=page_asset_ref, + ) + response = await api_client.get( + f"/api/v1/documents/{document_id}/chunks", + params={ + "page": 1, + "page_size": 1, + "include_asset_urls": "true", + }, + ) + default_response = await api_client.get( + f"/api/v1/documents/{document_id}/chunks", + params={ + "page": 1, + "page_size": 1, + }, + ) + + assert response.status_code == 200 + assert default_response.status_code == 200 + + chunk = cast(list[dict[str, object]], response.json()["chunks"])[0] + default_chunk = cast(list[dict[str, object]], default_response.json()["chunks"])[0] + metadata = cast(dict[str, object], chunk["metadata"]) + metadata_page_assets = cast(list[dict[str, object]], metadata["page_assets"]) + + expected_asset_url = ( + "filesystem://knowhere-test-results/" + f"results/{revision['job_id']}/{page_asset_ref}" + "?method=GET&expires_in=604800" + ) + assert metadata_page_assets[0]["asset_url"] == expected_asset_url + assert "page_assets" not in chunk + assert default_chunk["metadata"] == { + "summary": "Page node summary", + "page_nums": [4], + "page_assets": [ + { + "page_num": 4, + "artifact_ref": page_asset_ref, + "content_type": "image/png", + "width": 1200, + "height": 1800, + "source": "knowhere-rendered-page-citation-source", + } + ], + } + assert "page_assets" not in default_chunk + + @pytest.mark.asyncio async def test_should_return_not_found_when_listing_chunks_for_missing_document( developer_api_client_factory: Callable[ diff --git a/apps/api/tests/contract/test_page_memory_parse_track_contract.py b/apps/api/tests/contract/test_page_memory_parse_track_contract.py index 5dc29eea..17e7752a 100644 --- a/apps/api/tests/contract/test_page_memory_parse_track_contract.py +++ b/apps/api/tests/contract/test_page_memory_parse_track_contract.py @@ -209,6 +209,14 @@ def test_v2_jobs_documents_and_retrieval_routes_are_registered_in_openapi() -> N assert any(path.endswith("/v2/jobs/{job_id}") for path in paths) assert any(path.endswith("/v2/jobs/{job_id}/confirm-upload") for path in paths) assert any(path.endswith("/v2/documents") for path in paths) + assert any( + path.endswith("/v2/documents/{document_id}/files/page-citation-source") + for path in paths + ) + assert not any( + path.endswith("/v1/documents/{document_id}/files/page-citation-source") + for path in paths + ) assert any(path.endswith("/v2/retrieval/query") for path in paths) @@ -219,6 +227,7 @@ def test_v2_jobs_documents_and_retrieval_routes_are_registered_in_openapi() -> N "/v2/jobs/job_123", "/v2/documents", "/v2/documents/doc_123", + "/v2/documents/doc_123/files/page-citation-source", "/v2/retrieval/query", ], ) diff --git a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py index 3dabcdf0..28697c45 100644 --- a/apps/worker/app/services/document_agent/structure/hierarchy_locator.py +++ b/apps/worker/app/services/document_agent/structure/hierarchy_locator.py @@ -302,6 +302,10 @@ def _resolve_siblings( printed_page=node.printed_page, page_offset_hint=page_offset_hint, ) + if match is None and node.children and match_overrides: + match = _infer_start_from_descendant_overrides( + node, parent_titles, match_overrides, pages, + ) if match is None: start_page = lower_bound else: @@ -405,11 +409,58 @@ def _find_next_located_sibling( printed_page=sibling.printed_page, page_offset_hint=page_offset_hint, ) + if match is None and sibling.children and match_overrides: + match = _infer_start_from_descendant_overrides( + sibling, parent_titles, match_overrides, pages, + ) if match is not None: return match return None +def _infer_start_from_descendant_overrides( + node: TitleNode, + parent_titles: tuple[str, ...], + match_overrides: dict[tuple[str, ...], TitleMatch], + scope_pages: list[int], +) -> TitleMatch | None: + """Infer a parent node's start page from its earliest located descendant leaf. + + When a non-leaf node cannot be directly located (no printed_page, no grep + match), its descendant leaves may already be in match_overrides from + offset-guided bulk anchoring. Use the minimum page among those descendants + as a synthetic match so the resolver can cap the previous sibling's end_page. + """ + if not node.children or not match_overrides: + return None + leaves = iter_leaf_title_nodes([node], parent_titles=parent_titles) + min_page: int | None = None + min_match: TitleMatch | None = None + for leaf_path, _leaf_node in leaves: + m = match_overrides.get(leaf_path) + if m is None: + continue + if m.page not in scope_pages: + continue + if min_page is None or m.page < min_page: + min_page = m.page + min_match = m + if min_match is None: + return None + return TitleMatch( + page=min_match.page, + confidence=min(min_match.confidence, 0.80), + source=min_match.source, + matched_line="", + score=min(min_match.score, 0.80), + candidates=[min_match.page], + evidence={ + "inferred_from": "descendant_leaf_override", + "original_confidence": min_match.confidence, + }, + ) + + def _match_override( path_titles: tuple[str, ...], match_overrides: dict[tuple[str, ...], TitleMatch], diff --git a/apps/worker/app/services/document_agent/structure/page_locate_agent.py b/apps/worker/app/services/document_agent/structure/page_locate_agent.py index d31014f7..ad033f54 100644 --- a/apps/worker/app/services/document_agent/structure/page_locate_agent.py +++ b/apps/worker/app/services/document_agent/structure/page_locate_agent.py @@ -132,6 +132,22 @@ def prepare(self, nodes: list[TitleNode]) -> PageLocatePrepareResult: summary=summary, ) + def _entry_scope(self, node: TitleNode) -> list[int]: + """Narrow search scope for a node using its printed page + offset hint. + + When a node carries a printed_page and we have a page_offset_hint, + restrict grep to a small window around the expected physical page. + The window spans [printed_page, printed_page + offset] (inclusive, + order-independent) to cover the uncertainty between printed numbering + and physical page positions. + """ + if node.printed_page is None or self.page_offset_hint is None: + return self.body_pages + expected = node.printed_page + self.page_offset_hint + lo = min(node.printed_page, expected) + hi = max(node.printed_page, expected) + return [p for p in self.body_pages if lo <= p <= hi] + def _classify_residuals( self, nodes: list[TitleNode], @@ -139,9 +155,10 @@ def _classify_residuals( direct: dict[tuple[str, ...], TitleMatch] = {} residuals: list[ResidualRequest] = [] for path_titles, node in iter_leaf_title_nodes(nodes): + scope = self._entry_scope(node) match = locate_title_strict_exact( node.title, - scope_pages=self.body_pages, + scope_pages=scope, page_texts=self.page_texts, ) if match is not None: @@ -180,9 +197,10 @@ def _resolve_residuals( verify_page_cap=max(self.config.vlm_candidate_page_cap, 1), ) for residual in residuals: + scope = self._entry_scope(residual.node) agent = PageLocateSubAgent( ctx=self.ctx, - scope_pages=self.body_pages, + scope_pages=scope, page_count=self.page_count, config=sub_config, page_offset_hint=self.page_offset_hint, diff --git a/apps/worker/app/services/document_ingestion/artifact_refs.py b/apps/worker/app/services/document_ingestion/artifact_refs.py index a72ab633..df90073f 100644 --- a/apps/worker/app/services/document_ingestion/artifact_refs.py +++ b/apps/worker/app/services/document_ingestion/artifact_refs.py @@ -19,6 +19,24 @@ def collect_referenced_artifact_refs(chunks: list[dict[str, Any]]) -> set[str]: metadata.get("file_path") or chunk.get("file_path") or chunk.get("path"), allowed_roots={allowed_root}, ) + elif chunk_type == "page": + for page_asset in _iter_page_asset_refs(metadata.get("page_assets")): + _add_artifact_ref( + refs, + page_asset, + allowed_roots={"page_citation_assets"}, + ) + return refs + + +def _iter_page_asset_refs(raw_page_assets: object) -> list[object]: + if not isinstance(raw_page_assets, list): + return [] + refs: list[object] = [] + for item in raw_page_assets: + if not isinstance(item, dict): + continue + refs.append(item.get("artifact_ref")) return refs @@ -47,6 +65,6 @@ def _normalize_client_artifact_ref(raw_ref: object) -> str | None: ] if len(parts) < 2: return None - if parts[0] not in {"images", "tables"}: + if parts[0] not in {"images", "tables", "page_citation_assets"}: return None return "/".join(parts) diff --git a/apps/worker/app/services/page_memory/_utils.py b/apps/worker/app/services/page_memory/_utils.py index 816dee91..27d9d6c5 100644 --- a/apps/worker/app/services/page_memory/_utils.py +++ b/apps/worker/app/services/page_memory/_utils.py @@ -3,6 +3,7 @@ from __future__ import annotations import re +from dataclasses import dataclass, replace from typing import Any _NATURAL_RE = re.compile(r"(\d+)") @@ -57,3 +58,162 @@ def page_scope_info(pages: Any) -> dict[str, Any]: "page_count": len(normalized), "page_ranges": collapse_page_ranges(normalized), } + + +@dataclass(frozen=True) +class CoarseScope: + scope_id: str + skeletons: list[Any] + strategy: str + start_page: int + end_page: int + + +def build_hierarchy_scopes( + *, + skeletons: list[Any], + filename: str, + page_count: int, +) -> list[CoarseScope]: + """Split skeleton list into coarse scopes for independent processing. + + Groups skeletons by top-level section_path prefix. Skeletons not claimed + by any top-node are grouped by their root ancestor (first path segment + after filename) as orphan scopes. + """ + if not skeletons: + return [] + + root_fallback = all( + getattr(item, "title", "") == "Root" + or (getattr(item, "evidence", {}) or {}).get("source") == "fallback_root" + for item in skeletons + ) + if root_fallback: + ordered = sort_skeletons(skeletons) + return [ + CoarseScope( + scope_id=scope_id_for_pages(1, page_count), + skeletons=ordered, + strategy="fallback_root", + start_page=1, + end_page=page_count, + ) + ] + + min_level = min(int(getattr(item, "level", 0) or 0) for item in skeletons) + top_nodes = [ + item + for item in skeletons + if int(getattr(item, "level", 0) or 0) == min_level + or not str(getattr(item, "parent_path", "") or "").startswith(f"{filename}/") + ] + seen_top_paths: set[str] = set() + unique_top_nodes: list[Any] = [] + for item in sort_skeletons(top_nodes): + if item.section_path in seen_top_paths: + continue + seen_top_paths.add(item.section_path) + unique_top_nodes.append(item) + + scopes: list[CoarseScope] = [] + for index, top in enumerate(unique_top_nodes): + prefix = f"{top.section_path}/" + members = [ + item + for item in skeletons + if item.section_path == top.section_path + or str(item.section_path).startswith(prefix) + ] + if not members: + continue + start_page = max(1, int(getattr(top, "start_page", 1) or 1)) + next_top_start = ( + int(getattr(unique_top_nodes[index + 1], "start_page", page_count + 1) or page_count + 1) + if index + 1 < len(unique_top_nodes) + else page_count + 1 + ) + end_page = min( + page_count, + max( + start_page, + min( + int(getattr(top, "end_page", page_count) or page_count), + next_top_start - 1, + ), + ), + ) + bounded_members = [ + replace( + item, + start_page=max( + start_page, + int(getattr(item, "start_page", start_page) or start_page), + ), + end_page=min( + end_page, + int(getattr(item, "end_page", end_page) or end_page), + ), + ) + for item in members + if int(getattr(item, "start_page", 1) or 1) <= end_page + and int(getattr(item, "end_page", end_page) or end_page) >= start_page + ] + bounded_members = sort_skeletons(bounded_members) + scopes.append( + CoarseScope( + scope_id=scope_id_for_pages(start_page, end_page), + skeletons=bounded_members, + strategy=f"coarse_scope_{index + 1}", + start_page=start_page, + end_page=end_page, + ) + ) + + if scopes: + claimed_paths: set[str] = set() + for scope in scopes: + for item in scope.skeletons: + claimed_paths.add(item.section_path) + orphans = [item for item in skeletons if item.section_path not in claimed_paths] + if orphans: + root_groups: dict[str, list[Any]] = {} + for item in orphans: + parts = str(getattr(item, "section_path", "") or "").split("/") + root_key = "/".join(parts[:2]) if len(parts) >= 2 else item.section_path + root_groups.setdefault(root_key, []).append(item) + for root_key, group in sorted(root_groups.items()): + group = sort_skeletons(group) + sp = max(1, min(int(getattr(g, "start_page", 1) or 1) for g in group)) + ep = min(page_count, max(int(getattr(g, "end_page", page_count) or page_count) for g in group)) + scopes.append( + CoarseScope( + scope_id=scope_id_for_pages(sp, ep), + skeletons=group, + strategy="orphan_scope", + start_page=sp, + end_page=ep, + ) + ) + return scopes + + ordered = sort_skeletons(skeletons) + all_pages: set[int] = set() + for item in ordered: + sp = max(1, int(getattr(item, "start_page", 1) or 1)) + ep = min(page_count, int(getattr(item, "end_page", sp) or sp)) + if ep >= sp: + all_pages.update(range(sp, ep + 1)) + pages = sorted(all_pages) or list(range(1, page_count + 1)) + return [ + CoarseScope( + scope_id=scope_id_for_pages( + pages[0] if pages else 1, + pages[-1] if pages else page_count, + ), + skeletons=ordered, + strategy="full_coarse_hierarchy", + start_page=pages[0] if pages else 1, + end_page=pages[-1] if pages else page_count, + ) + ] diff --git a/apps/worker/app/services/page_memory/memory_service.py b/apps/worker/app/services/page_memory/memory_service.py index 9e697aed..d1397bc3 100644 --- a/apps/worker/app/services/page_memory/memory_service.py +++ b/apps/worker/app/services/page_memory/memory_service.py @@ -2,7 +2,7 @@ import os import shutil -from dataclasses import dataclass, field, replace +from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -21,9 +21,10 @@ from app.services.document_parser.support.stage_profiler import stage_timer from app.services.page_memory.normalizer import normalize_to_pdf from app.services.page_memory._utils import ( + CoarseScope, + build_hierarchy_scopes, collapse_page_ranges, page_scope_info, - scope_id_for_pages, sort_skeletons, ) from app.services.page_memory._serialization import ( @@ -52,13 +53,7 @@ class PageMemoryInput: ) -@dataclass(frozen=True) -class _HierarchyScope: - scope_id: str - skeletons: list[Any] - strategy: str - start_page: int - end_page: int +_HierarchyScope = CoarseScope @dataclass(frozen=True) @@ -263,7 +258,6 @@ def _build_page_dataframe( filename=filename, page_texts=page_texts, ctx=ctx, - page_memory_config=page_memory_config, ) else: skeletons = [] @@ -505,111 +499,11 @@ def _build_hierarchy_scopes( filename: str, page_count: int, ) -> list[_HierarchyScope]: - if not skeletons: - return [] - - root_fallback = all( - getattr(item, "title", "") == "Root" - or (getattr(item, "evidence", {}) or {}).get("source") == "fallback_root" - for item in skeletons + return build_hierarchy_scopes( + skeletons=skeletons, + filename=filename, + page_count=page_count, ) - if root_fallback: - ordered = sort_skeletons(skeletons) - return [ - _HierarchyScope( - scope_id=scope_id_for_pages(1, page_count), - skeletons=ordered, - strategy="fallback_root", - start_page=1, - end_page=page_count, - ) - ] - - min_level = min(int(getattr(item, "level", 0) or 0) for item in skeletons) - top_nodes = [ - item - for item in skeletons - if int(getattr(item, "level", 0) or 0) == min_level - or not str(getattr(item, "parent_path", "") or "").startswith(f"{filename}/") - ] - seen_top_paths: set[str] = set() - unique_top_nodes: list[Any] = [] - for item in sort_skeletons(top_nodes): - if item.section_path in seen_top_paths: - continue - seen_top_paths.add(item.section_path) - unique_top_nodes.append(item) - - scopes: list[_HierarchyScope] = [] - for index, top in enumerate(unique_top_nodes): - prefix = f"{top.section_path}/" - members = [ - item - for item in skeletons - if item.section_path == top.section_path - or str(item.section_path).startswith(prefix) - ] - if not members: - continue - start_page = max(1, int(getattr(top, "start_page", 1) or 1)) - next_top_start = ( - int(getattr(unique_top_nodes[index + 1], "start_page", page_count + 1) or page_count + 1) - if index + 1 < len(unique_top_nodes) - else page_count + 1 - ) - end_page = min( - page_count, - max( - start_page, - min( - int(getattr(top, "end_page", page_count) or page_count), - next_top_start - 1, - ), - ), - ) - bounded_members = [ - replace( - item, - start_page=max( - start_page, - int(getattr(item, "start_page", start_page) or start_page), - ), - end_page=min( - end_page, - int(getattr(item, "end_page", end_page) or end_page), - ), - ) - for item in members - if int(getattr(item, "start_page", 1) or 1) <= end_page - and int(getattr(item, "end_page", end_page) or end_page) >= start_page - ] - bounded_members = sort_skeletons(bounded_members) - scopes.append( - _HierarchyScope( - scope_id=scope_id_for_pages(start_page, end_page), - skeletons=bounded_members, - strategy=f"coarse_scope_{index + 1}", - start_page=start_page, - end_page=end_page, - ) - ) - if scopes: - return scopes - - ordered = sort_skeletons(skeletons) - pages = _derive_hierarchy_page_scope(skeletons=ordered, page_count=page_count) - return [ - _HierarchyScope( - scope_id=scope_id_for_pages( - pages[0] if pages else 1, - pages[-1] if pages else page_count, - ), - skeletons=ordered, - strategy="full_coarse_hierarchy", - start_page=pages[0] if pages else 1, - end_page=pages[-1] if pages else page_count, - ) - ] def _allocate_asset_pages( diff --git a/apps/worker/app/services/page_memory/node_assembler.py b/apps/worker/app/services/page_memory/node_assembler.py index 5c50fab5..ff4191a2 100644 --- a/apps/worker/app/services/page_memory/node_assembler.py +++ b/apps/worker/app/services/page_memory/node_assembler.py @@ -22,7 +22,9 @@ import json import os +import shutil from dataclasses import dataclass, field +from pathlib import Path from typing import Any from loguru import logger @@ -40,6 +42,8 @@ SAME_AS_PREFIX = "SAME-AS" _NODE_SUMMARY_MAX_PAGES_DEFAULT = 5 +_PAGE_CITATION_ASSET_SOURCE = "knowhere-rendered-page-citation-source" +_PAGE_CITATION_ASSET_CONTENT_TYPE = "image/png" @dataclass(frozen=True) @@ -462,7 +466,10 @@ def build_node_rows( "page_nums": ",".join(str(page) for page in view.pages), "entities": serialize_entities(entities), "asset_title": "", - "extra_metadata": {}, + "extra_metadata": _build_page_extra_metadata( + pages=view.pages, + image_path_by_page=image_path_by_page, + ), } rows.append(row) rows_by_path[leaf.section_path] = row @@ -487,6 +494,88 @@ def build_node_rows( return asset_rows + rows +def _build_page_extra_metadata( + *, + pages: list[int], + image_path_by_page: dict[int, str], +) -> dict[str, Any]: + page_assets = _build_page_citation_assets( + pages=pages, + image_path_by_page=image_path_by_page, + ) + if not page_assets: + return {} + return {"page_assets": page_assets} + + +def _build_page_citation_assets( + *, + pages: list[int], + image_path_by_page: dict[int, str], +) -> list[dict[str, Any]]: + assets: list[dict[str, Any]] = [] + seen_pages: set[int] = set() + for page in pages: + if page in seen_pages: + continue + seen_pages.add(page) + image_path = image_path_by_page.get(page) + if not image_path or not os.path.exists(image_path): + continue + artifact_ref = _promote_page_citation_asset(page=page, image_path=image_path) + if not artifact_ref: + continue + width, height = _read_image_dimensions(image_path) + asset = { + "page_num": page, + "artifact_ref": artifact_ref, + "content_type": _PAGE_CITATION_ASSET_CONTENT_TYPE, + "source": _PAGE_CITATION_ASSET_SOURCE, + } + if width is not None: + asset["width"] = width + if height is not None: + asset["height"] = height + assets.append(asset) + return assets + + +def _promote_page_citation_asset(*, page: int, image_path: str) -> str: + source_path = Path(image_path) + output_dir = source_path.parent.parent + target_dir = output_dir / "page_citation_assets" + target_path = target_dir / f"page-{page}.png" + artifact_ref = f"page_citation_assets/page-{page}.png" + try: + target_dir.mkdir(parents=True, exist_ok=True) + if source_path.resolve() != target_path.resolve(): + shutil.copyfile(source_path, target_path) + return artifact_ref + except Exception as exc: + logger.warning( + "[node_assembler] failed to promote page citation asset page={} path={}: {}", + page, + image_path, + exc, + ) + return "" + + +def _read_image_dimensions(image_path: str) -> tuple[int | None, int | None]: + try: + from PIL import Image + + with Image.open(image_path) as image: + return int(image.width), int(image.height) + except Exception as exc: + logger.debug( + "[node_assembler] failed to read page citation image dimensions {}: {}", + image_path, + exc, + ) + return None, None + + def _attach_asset_connections( *, page_assets_by_page: dict[int, list[PageAsset]], diff --git a/apps/worker/app/services/page_memory/skeleton_extractor.py b/apps/worker/app/services/page_memory/skeleton_extractor.py index 7a7d4ad6..d0778137 100644 --- a/apps/worker/app/services/page_memory/skeleton_extractor.py +++ b/apps/worker/app/services/page_memory/skeleton_extractor.py @@ -16,13 +16,14 @@ ToolContext, ) from app.services.document_agent.structure.page_locate_agent import ( - PageLocateConfig, - PageLocateResidualAgent, + verify_section_page_choice, ) from app.services.document_agent.structure.hierarchy_locator import ( ResolvedHierarchyRange, + TitleMatch, TitleNode, extract_toc_nodes, + iter_leaf_title_nodes, resolve_hierarchy_page_ranges, ) from app.services.document_parser.structure.body_boundary import ( @@ -30,11 +31,62 @@ normalize_heading_text, ) from loguru import logger -from shared.models.schemas.page_memory_config import PageMemoryConfig _FRONT_TOC_REGION_GAP_PAGES = 5 +def _prune_out_of_scope_nodes( + nodes: list[TitleNode], + *, + offset: int, + page_count: int, +) -> tuple[list[TitleNode], int]: + """Remove leaf nodes whose printed_page + offset exceeds page_count. + + Bottom-up: prune out-of-scope leaves, then remove intermediate nodes + that become childless after pruning. Returns (pruned_tree, removed_count). + """ + from dataclasses import replace as _replace + + removed = 0 + + def _prune(node: TitleNode) -> TitleNode | None: + nonlocal removed + if not node.children: + if node.printed_page is not None: + expected = node.printed_page + offset + if expected > page_count or expected < 1: + removed += 1 + return None + return node + pruned_children = [] + for child in node.children: + result = _prune(child) + if result is not None: + pruned_children.append(result) + if not pruned_children: + removed += 1 + return None + return _replace(node, children=pruned_children) + + pruned = [] + for node in nodes: + result = _prune(node) + if result is not None: + pruned.append(result) + + if removed: + logger.info( + "[page_memory.skeleton] pruned {} out-of-scope TOC nodes " + "(printed_page + offset={} exceeds page_count={})", + removed, + offset, + page_count, + ) + + return pruned, removed + + @dataclass(frozen=True) class SectionSkeleton: section_path: str @@ -56,7 +108,6 @@ def extract_section_skeletons( page_texts: dict[int, str], ctx: ToolContext | None = None, hierarchy_nodes: list[TitleNode] | None = None, - page_memory_config: PageMemoryConfig | None = None, ) -> list[SectionSkeleton]: """Convert PageAnatomyMap hierarchy evidence into section skeletons. @@ -69,10 +120,12 @@ def extract_section_skeletons( return [_root_skeleton(root_path=root_path, filename=filename, page_count=0)] toc_selection: dict[str, Any] = {} + pending_tocs: list[dict[str, Any]] = [] + toc_hierarchies: list[dict[str, Any]] | None = None if hierarchy_nodes: nodes = hierarchy_nodes else: - toc_hierarchies, toc_selection = _select_global_toc_hierarchies( + toc_hierarchies, pending_tocs, toc_selection = _select_global_toc_hierarchies( anatomy=anatomy, filename=filename, ) @@ -91,30 +144,86 @@ def extract_section_skeletons( # Collapse degenerate single-child intermediate chains before locate. # Rule: only merge a parent with its only child when that child is NOT a # leaf (i.e. the child still has children of its own). This preserves the - # original leaf title so PageLocateResidualAgent can find it in the PDF. + # original leaf title so offset-guided anchoring can find it in the PDF. nodes = _collapse_intermediate_single_child_chains(nodes) body_pages = _body_pages(anatomy=anatomy, page_count=page_count) - offset_hint = _estimate_page_offset(nodes=nodes, anatomy=anatomy) - locate_result = PageLocateResidualAgent( + + # When pending TOCs exist, limit primary scope so the last sibling's + # end_page doesn't extend into the pending TOC region. + primary_page_count = page_count + primary_body_pages = body_pages + if pending_tocs: + pending_starts: list[int] = [] + for t in pending_tocs: + start = _toc_range_start(t) + if start is not None: + pending_starts.append(start) + if pending_starts: + primary_page_count = min(pending_starts) - 1 + primary_body_pages = [p for p in body_pages if p <= primary_page_count] + + offset_hint, calibration_overrides = _calibrate_offset_via_vlm( + nodes=nodes, + toc_hierarchies=toc_hierarchies if not hierarchy_nodes else None, ctx=ctx, page_texts=page_texts, - body_pages=body_pages, page_count=page_count, - page_offset_hint=offset_hint, - config=( - PageLocateConfig.from_page_memory_config(page_memory_config) - if page_memory_config is not None - else None - ), - ).prepare(nodes) + ) + + # Prune TOC nodes whose printed_page + offset exceeds the physical PDF. + pruned_count = 0 + if offset_hint is not None: + nodes, pruned_count = _prune_out_of_scope_nodes( + nodes, offset=offset_hint, page_count=page_count, + ) + if not nodes: + return [ + _root_skeleton( + root_path=root_path, + filename=filename, + page_count=page_count, + reason="all_toc_nodes_out_of_scope", + ) + ] + + # Phase A3: try offset-guided bulk anchoring before expensive residual agent. + offset_matches: dict[tuple[str, ...], TitleMatch] | None = None + if offset_hint is not None and ctx is not None: + offset_matches = _offset_guided_anchoring( + nodes=nodes, + offset=offset_hint, + ctx=ctx, + page_count=page_count, + calibration_overrides=calibration_overrides, + ) + + if offset_matches is not None: + match_overrides = offset_matches + locate_summary: dict[str, Any] = { + "agent": "offset_guided_bulk", + "offset": offset_hint, + "bulk_count": len(offset_matches), + "pruned_out_of_scope": pruned_count, + } + else: + match_overrides = calibration_overrides + locate_summary = { + "agent": "offset_only", + "offset": offset_hint, + "reason": "offset_guided_anchoring_skipped_or_empty", + "pruned_out_of_scope": pruned_count, + } + resolve_nodes = nodes + + ranges = resolve_hierarchy_page_ranges( - locate_result.nodes, - page_count=page_count, + resolve_nodes, + page_count=primary_page_count, page_texts=page_texts, - body_pages=body_pages, + body_pages=primary_body_pages, page_offset_hint=offset_hint, - match_overrides=locate_result.match_overrides, + match_overrides=match_overrides, ) if not ranges: return [ @@ -131,11 +240,25 @@ def extract_section_skeletons( item, filename=filename, page_count=page_count, - locate_summary=locate_result.summary, + locate_summary=locate_summary, toc_selection=toc_selection, ) for item in ranges ] + + # Phase B: graft pending TOCs (appendix / parallel sections) + if pending_tocs: + secondary_skeletons = _resolve_pending_tocs( + pending_tocs=pending_tocs, + primary_ranges=ranges, + ctx=ctx, + page_texts=page_texts, + page_count=page_count, + filename=filename, + body_pages=body_pages, + ) + skeletons.extend(secondary_skeletons) + _log_unlocated_title_warnings(filename=filename, skeletons=skeletons) return skeletons @@ -177,11 +300,11 @@ def _range_to_skeleton( # Motivation: TOC hierarchies often contain "structural" intermediate nodes # (category codes, volume identifiers) that add depth but carry no locatable # text. Compressing them before locate keeps emit_depth small and lets the -# VLM/grep focus on meaningful leaf titles. +# offset-guided anchoring focus on meaningful leaf titles. # # Critical invariant: a node whose only child is a LEAF (no grandchildren) is # NOT merged, so the leaf's original title survives unchanged into -# PageLocateResidualAgent. Only pure-intermediate chains are compressed. +# offset-guided anchoring. Only pure-intermediate chains are compressed. def _collapse_intermediate_single_child_chains( @@ -256,17 +379,19 @@ def _select_global_toc_hierarchies( *, anatomy: Any | None, filename: str, -) -> tuple[list[dict[str, Any]] | None, dict[str, Any]]: - """Keep the front/global TOC cluster and skip later embedded TOCs. +) -> tuple[list[dict[str, Any]] | None, list[dict[str, Any]], dict[str, Any]]: + """Split TOC hierarchies into primary (front cluster) and pending (for probe). + + Profile-time TOC extraction can find multiple TOCs in a long document. + The front cluster is selected by physical page proximity. Remaining TOCs + are returned as *pending* for downstream independent calibration rather + than being unconditionally discarded. - Profile-time TOC extraction can find local TOCs inside a long document - (for example, an embedded standard with its own English outline). Page - memory C4 currently emits a document-level skeleton, so later page-based - TOC regions must not be concatenated as root siblings. + Returns (primary_hierarchies, pending_hierarchies, summary). """ hierarchies = list(_toc_hierarchies(anatomy) or []) if len(hierarchies) <= 1: - return (hierarchies or None), {} + return (hierarchies or None), [], {} page_based = [ hierarchy @@ -274,11 +399,11 @@ def _select_global_toc_hierarchies( if hierarchy.get("toc_range_unit") == "page" and _toc_range_start(hierarchy) is not None ] if not page_based or len(page_based) != len(hierarchies): - return hierarchies, {} + return hierarchies, [], {} sorted_items = sorted(enumerate(hierarchies), key=lambda item: _toc_range_start(item[1]) or 0) selected_indices: set[int] = set() - skipped: list[dict[str, Any]] = [] + pending_indices: list[int] = [] cluster_end: int | None = None for original_index, hierarchy in sorted_items: @@ -295,34 +420,29 @@ def _select_global_toc_hierarchies( selected_indices.add(original_index) cluster_end = max(cluster_end, end) continue - skipped.append( - { - "index": original_index, - "toc_range": [start, end], - "scan_range": hierarchy.get("scan_range"), - "reason": "embedded_toc_region_outside_front_cluster", - } - ) + pending_indices.append(original_index) selected = [ hierarchy for index, hierarchy in enumerate(hierarchies) if index in selected_indices ] - if skipped: - logger.warning( - "[page_memory.skeleton] skipped embedded toc regions filename={} skipped={}", + pending = [hierarchies[i] for i in pending_indices] + + if pending: + logger.info( + "[page_memory.skeleton] toc split: primary={} pending={} filename={}", + len(selected), + len(pending), filename, - skipped, ) summary = { - "strategy": "front_page_toc_cluster", + "strategy": "front_cluster_with_pending", "input_count": len(hierarchies), - "selected_count": len(selected), - "skipped_count": len(skipped), - "skipped": skipped, + "primary_count": len(selected), + "pending_count": len(pending), } - return (selected or None), summary + return (selected or None), pending, summary def _toc_range_start(hierarchy: dict[str, Any]) -> int | None: @@ -364,35 +484,594 @@ def _body_pages(*, anatomy: Any | None, page_count: int) -> list[int]: return [page for page in range(1, page_count + 1) if page not in excluded] -def _estimate_page_offset(*, nodes: list[TitleNode], anatomy: Any | None) -> int | None: - printed_by_title: dict[str, int] = {} - for node in _walk_nodes(nodes): +# ── VLM offset calibration (Phase A1) ─────────────────────────────────────── + +_CALIBRATION_WINDOW_PAGES = 10 +_CALIBRATION_LEAF_PROBE_COUNT = 3 + + +def _calibrate_offset_via_vlm( + *, + nodes: list[TitleNode], + toc_hierarchies: list[dict[str, Any]] | None, + ctx: ToolContext | None, + page_texts: dict[int, str], + page_count: int, +) -> tuple[int | None, dict[tuple[str, ...], TitleMatch]]: + """Scan pages after the TOC to find the first leaf entry via VLM. + + Computes offset = confirmed_physical_page - printed_page. + Returns (offset, match_overrides) where match_overrides contains the + confirmed entry so downstream locate doesn't re-process it. + """ + if ctx is None: + return None, {} + + toc_physical_end = _toc_cluster_end_page(toc_hierarchies) + if toc_physical_end is None: + return None, {} + + scan_start = toc_physical_end + 1 + scan_end = min(scan_start + _CALIBRATION_WINDOW_PAGES - 1, page_count) + if scan_start > page_count: + return None, {} + + leaves = list(iter_leaf_title_nodes(nodes)) + probe_leaves = [ + (path_titles, node) + for path_titles, node in leaves + if node.printed_page is not None + ][:_CALIBRATION_LEAF_PROBE_COUNT] + + if not probe_leaves: + return None, {} + + scan_pages = list(range(scan_start, scan_end + 1)) + candidates = [ + TitleMatch( + page=page, + confidence=0.4, + source="agent_heuristic", + matched_line="", + score=0.4, + candidates=scan_pages, + evidence={"calibration_probe": True}, + ) + for page in scan_pages + ] + + for path_titles, node in probe_leaves: + result = verify_section_page_choice( + ctx=ctx, + title=node.title, + candidate_matches=candidates, + candidate_page_cap=len(scan_pages), + ) + selected = result.get("selected_page") + if selected is not None and result.get("confidence", 0) >= 0.6: + offset = selected - node.printed_page + match = TitleMatch( + page=selected, + confidence=result.get("confidence", 0.75), + source="agent_vlm", + matched_line="", + score=result.get("confidence", 0.75), + candidates=[selected], + evidence={ + "calibration": True, + "printed_page": node.printed_page, + "reason": result.get("reason", ""), + }, + ) + logger.info( + "[page_memory.skeleton] calibration confirmed: title={!r} " + "printed_page={} physical_page={} offset={}", + node.title, + node.printed_page, + selected, + offset, + ) + return offset, {path_titles: match} + + logger.info("[page_memory.skeleton] calibration: no leaf confirmed in scan window") + return None, {} + + +def _toc_cluster_end_page(toc_hierarchies: list[dict[str, Any]] | None) -> int | None: + """Get the last physical page of the primary TOC cluster.""" + if not toc_hierarchies: + return None + end_pages: list[int] = [] + for hierarchy in toc_hierarchies: + end = _toc_range_end(hierarchy) + if end is not None: + end_pages.append(end) + return max(end_pages) if end_pages else None + + +# ── Offset-guided bulk anchoring with recursive recalibrate (Phase A3) ─────── + +_TAIL_VERIFY_CONFIDENCE_THRESHOLD = 0.6 +_MAX_RECALIBRATE_DEPTH = 5 +_MAX_RECALIBRATE_DELTA = 5 + + +def _verify_offset_tail( + *, + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, + ctx: ToolContext, + page_count: int, +) -> bool: + """VLM-verify that the offset holds for the last leaf entry (Theorem 1). + + If head offset == tail offset, monotonicity guarantees all intermediate + entries share the same offset. + + Prefers a tail leaf whose expected page is strictly less than page_count + (boundary pages are unreliable for VLM verification). + """ + tail_leaves = [ + (path, node) for path, node in reversed(leaves) if node.printed_page is not None + ] + if not tail_leaves: + return True + + # Prefer non-boundary: printed_page + offset < page_count + selected = None + for path, node in tail_leaves: + pp = node.printed_page + if pp is None: + continue + expected = pp + offset + if 1 <= expected < page_count: + selected = (path, node) + break + if selected is None: + # All leaves are at the boundary; fall back to the last one + selected = tail_leaves[0] + + path, node = selected + printed_page = node.printed_page + if printed_page is None: + return True + expected_page = printed_page + offset + if expected_page < 1 or expected_page > page_count: + return False + + candidate = TitleMatch( + page=expected_page, + confidence=0.4, + source="agent_heuristic", + matched_line="", + score=0.4, + candidates=[expected_page], + evidence={"tail_verify_probe": True}, + ) + result = verify_section_page_choice( + ctx=ctx, + title=node.title, + candidate_matches=[candidate], + candidate_page_cap=1, + ) + confirmed = ( + result.get("selected_page") == expected_page + and result.get("confidence", 0) >= _TAIL_VERIFY_CONFIDENCE_THRESHOLD + ) + logger.info( + "[page_memory.skeleton] tail verify: title={!r} expected_page={} confirmed={} confidence={}", + node.title, + expected_page, + confirmed, + result.get("confidence", 0), + ) + return confirmed + + +def _vlm_confirm_single_page( + *, + ctx: ToolContext, + title: str, + expected_page: int, + page_count: int, +) -> bool: + """Single-page VLM confirmation for binary search steps.""" + if expected_page < 1 or expected_page > page_count: + return False + candidate = TitleMatch( + page=expected_page, + confidence=0.4, + source="agent_heuristic", + matched_line="", + score=0.4, + candidates=[expected_page], + evidence={"bisect_probe": True}, + ) + result = verify_section_page_choice( + ctx=ctx, + title=title, + candidate_matches=[candidate], + candidate_page_cap=1, + ) + return ( + result.get("selected_page") == expected_page + and result.get("confidence", 0) >= _TAIL_VERIFY_CONFIDENCE_THRESHOLD + ) + + +def _bisect_offset_breakpoint( + *, + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, + ctx: ToolContext, + page_count: int, +) -> int: + """Binary search for the last leaf index where offset is valid. O(log n) VLM calls.""" + lo, hi = 0, len(leaves) - 1 + while lo < hi: + mid = (lo + hi + 1) // 2 + _, node = leaves[mid] if node.printed_page is None: + hi = mid - 1 continue - printed_by_title[_title_key(node.title)] = node.printed_page + expected = node.printed_page + offset + if _vlm_confirm_single_page( + ctx=ctx, title=node.title, expected_page=expected, page_count=page_count + ): + lo = mid + else: + hi = mid - 1 + logger.info( + "[page_memory.skeleton] bisect breakpoint: last_valid_index={} / total={}", + lo, + len(leaves), + ) + return lo - offsets: list[int] = [] - h1_result = getattr(anatomy, "h1_result", None) - for candidate in getattr(h1_result, "h1_candidates", []) or []: - printed_page = printed_by_title.get(_title_key(candidate.title)) - if printed_page is not None: - offsets.append(int(candidate.page) - printed_page) - if not offsets: + +def _bulk_offset_matches( + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, +) -> dict[tuple[str, ...], TitleMatch]: + """Generate TitleMatch overrides for all leaves using offset. No VLM calls.""" + matches: dict[tuple[str, ...], TitleMatch] = {} + for path_titles, node in leaves: + if node.printed_page is None: + continue + page = node.printed_page + offset + matches[path_titles] = TitleMatch( + page=page, + confidence=0.88, + source="agent_vlm", + matched_line="", + score=0.88, + candidates=[page], + evidence={ + "bulk_offset": True, + "offset": offset, + "printed_page": node.printed_page, + }, + ) + return matches + + +def _recalibrate_after_breakpoint( + *, + entry_node: TitleNode, + old_offset: int, + ctx: ToolContext, + page_count: int, +) -> int | None: + """Probe offsets old_offset+1, +2, ... to find new offset after breakpoint. + + Monotonicity guarantees new offset > old offset, so search space is tiny. + """ + entry_printed_page = entry_node.printed_page + if entry_printed_page is None: return None - offsets.sort() - return offsets[len(offsets) // 2] + for delta in range(1, _MAX_RECALIBRATE_DELTA + 1): + new_offset = old_offset + delta + if _vlm_confirm_single_page( + ctx=ctx, + title=entry_node.title, + expected_page=entry_printed_page + new_offset, + page_count=page_count, + ): + logger.info( + "[page_memory.skeleton] recalibrate: title={!r} new_offset={} (delta=+{})", + entry_node.title, + new_offset, + delta, + ) + return new_offset + return None -def _walk_nodes(nodes: list[TitleNode]) -> list[TitleNode]: - walked: list[TitleNode] = [] - for node in nodes: - walked.append(node) - walked.extend(_walk_nodes(node.children)) - return walked +def _offset_guided_anchoring( + *, + nodes: list[TitleNode], + offset: int, + ctx: ToolContext, + page_count: int, + calibration_overrides: dict[tuple[str, ...], TitleMatch], +) -> dict[tuple[str, ...], TitleMatch] | None: + """Offset-guided bulk anchoring with recursive recalibrate on breakpoints. + + Strategy: + 1. Tail verify last leaf with current offset + 2. If pass → bulk apply all leaves (Theorem 1) + 3. If fail → binary search for breakpoint + 4. Bulk apply leaves before breakpoint + 5. Recalibrate: probe remaining[0] with offset+1, +2, ... (monotonicity) + 6. Recurse on remaining segment with new offset + 7. If recalibrate fails → return partial (caller falls back for remainder) + + Returns match_overrides for all anchored leaves, or None for full fallback. + """ + leaves = [ + (path, node) + for path, node in iter_leaf_title_nodes(nodes) + if node.printed_page is not None + ] + if len(leaves) < 2: + return None + + all_matches: dict[tuple[str, ...], TitleMatch] = {} + all_matches.update(calibration_overrides) + + _anchor_segment_recursive( + leaves=leaves, + offset=offset, + ctx=ctx, + page_count=page_count, + matches=all_matches, + depth=0, + ) + + if not all_matches: + return None + + logger.info( + "[page_memory.skeleton] offset bulk anchoring: {} / {} leaves anchored", + len(all_matches), + len(leaves), + ) + return all_matches + + +def _anchor_segment_recursive( + *, + leaves: list[tuple[tuple[str, ...], TitleNode]], + offset: int, + ctx: ToolContext, + page_count: int, + matches: dict[tuple[str, ...], TitleMatch], + depth: int, +) -> None: + """Recursively anchor a segment of leaves, handling multiple breakpoints.""" + if not leaves or depth >= _MAX_RECALIBRATE_DEPTH: + return + + if _verify_offset_tail(leaves=leaves, offset=offset, ctx=ctx, page_count=page_count): + bulk = _bulk_offset_matches(leaves, offset) + matches.update(bulk) + return + + bp = _bisect_offset_breakpoint(leaves=leaves, offset=offset, ctx=ctx, page_count=page_count) + confirmed_leaves = leaves[: bp + 1] + if confirmed_leaves: + bulk = _bulk_offset_matches(confirmed_leaves, offset) + matches.update(bulk) + + remaining = leaves[bp + 1:] + if not remaining: + return + + _, first_remaining_node = remaining[0] + new_offset = _recalibrate_after_breakpoint( + entry_node=first_remaining_node, + old_offset=offset, + ctx=ctx, + page_count=page_count, + ) + if new_offset is None: + return + + _anchor_segment_recursive( + leaves=remaining, + offset=new_offset, + ctx=ctx, + page_count=page_count, + matches=matches, + depth=depth + 1, + ) + + +# ── Multi-TOC grafting (Track B) ───────────────────────────────────────────── + + +def _resolve_pending_tocs( + *, + pending_tocs: list[dict[str, Any]], + primary_ranges: list[ResolvedHierarchyRange], + ctx: ToolContext | None, + page_texts: dict[int, str], + page_count: int, + filename: str, + body_pages: list[int], +) -> list[SectionSkeleton]: + """Independently calibrate and anchor each pending TOC, then graft results. + + Each pending TOC gets its own offset via VLM calibration + tail verify, + then entries are bulk-anchored (or fallback to residual agent). + Classification is PARALLEL (append at root level) or CONTAINED (skip). + """ + if not pending_tocs or ctx is None: + return [] + + all_secondary_skeletons: list[SectionSkeleton] = [] + + for i, pending_toc in enumerate(pending_tocs): + toc_range = pending_toc.get("toc_range") + nodes = extract_toc_nodes([pending_toc]) + if not nodes: + continue + nodes = _collapse_intermediate_single_child_chains(nodes) + + # Each TOC's content scope: [toc_range_end + 1, next_toc_start - 1] + toc_end = _toc_range_end(pending_toc) + toc_scope_start = (toc_end + 1) if toc_end is not None else None + next_starts: list[int] = [] + for j in range(i + 1, len(pending_tocs)): + start = _toc_range_start(pending_tocs[j]) + if start is not None: + next_starts.append(start) + toc_scope_end = (min(next_starts) - 1) if next_starts else page_count + toc_body_pages = [ + p for p in body_pages + if p <= toc_scope_end and (toc_scope_start is None or p >= toc_scope_start) + ] + + offset, cal_overrides = _calibrate_offset_via_vlm( + nodes=nodes, + toc_hierarchies=[pending_toc], + ctx=ctx, + page_texts=page_texts, + page_count=toc_scope_end, + ) + + if offset is None: + logger.info( + "[page_memory.skeleton] pending TOC toc_range={}: calibration failed, skipping", + toc_range, + ) + continue + + relationship = _classify_toc_relationship( + offset=offset, + nodes=nodes, + primary_ranges=primary_ranges, + page_count=page_count, + ) + if relationship == "unresolvable": + logger.info( + "[page_memory.skeleton] pending TOC toc_range={}: unresolvable, skipping", + toc_range, + ) + continue + + offset_matches = _offset_guided_anchoring( + nodes=nodes, + offset=offset, + ctx=ctx, + page_count=toc_scope_end, + calibration_overrides=cal_overrides, + ) + + if offset_matches is not None: + match_overrides = offset_matches + locate_summary: dict[str, Any] = { + "agent": "offset_guided_bulk", + "offset": offset, + "bulk_count": len(offset_matches), + "toc_relationship": relationship, + } + else: + match_overrides = cal_overrides + locate_summary = { + "agent": "offset_only", + "offset": offset, + "toc_relationship": relationship, + "reason": "offset_guided_anchoring_skipped_or_empty", + } + + ranges = resolve_hierarchy_page_ranges( + nodes, + page_count=toc_scope_end, + page_texts=page_texts, + body_pages=toc_body_pages, + page_offset_hint=offset, + match_overrides=match_overrides, + ) + + toc_selection_info: dict[str, Any] = { + "toc_range": toc_range, + "offset": offset, + "relationship": relationship, + } + for item in ranges: + skeleton = _range_to_skeleton( + item, + filename=filename, + page_count=toc_scope_end, + locate_summary=locate_summary, + toc_selection=toc_selection_info, + ) + all_secondary_skeletons.append(skeleton) + + logger.info( + "[page_memory.skeleton] pending TOC toc_range={}: " + "relationship={} offset={} skeletons={}", + toc_range, + relationship, + offset, + len(ranges), + ) + + return all_secondary_skeletons + + +def _classify_toc_relationship( + *, + offset: int, + nodes: list[TitleNode], + primary_ranges: list[ResolvedHierarchyRange], + page_count: int, +) -> str: + """Classify a pending TOC as parallel or contained vs primary ranges. + + parallel: the pending TOC covers pages beyond the primary tree's *anchored* + content (i.e. the last explicitly-located section start page). + contained: the pending TOC's content falls strictly within a primary + section's explicitly-anchored range. + """ + leaves = [ + node for _, node in iter_leaf_title_nodes(nodes) if node.printed_page is not None + ] + if not leaves: + return "unresolvable" + + first_printed = leaves[0].printed_page + last_printed = leaves[-1].printed_page + if first_printed is None or last_printed is None: + return "unresolvable" + first_physical = first_printed + offset + last_physical = last_printed + offset + + if first_physical < 1 or first_physical > page_count: + return "unresolvable" + + if not primary_ranges: + return "parallel" + + # Use the last *start_page* among primary ranges as the boundary of + # explicitly-anchored content. The end_page of the last section is often + # extended to page_count by default and doesn't reflect real content coverage. + last_anchored_start = max( + (r.start_page for r in primary_ranges if r.start_page is not None), default=0 + ) + + if first_physical > last_anchored_start: + return "parallel" + min_level = min(r.level for r in primary_ranges) + top_level_ranges = [r for r in primary_ranges if r.level == min_level] + for r in top_level_ranges: + if r.start_page and r.end_page: + if r.start_page <= first_physical and last_physical <= r.end_page: + return "contained" -def _title_key(title: str) -> str: - return normalize_heading_text(clean_toc_title(title) or title).casefold() + return "parallel" def _clamp_page(page: int, page_count: int) -> int: diff --git a/apps/worker/tests/contract/test_document_agent_budget_contract.py b/apps/worker/tests/contract/test_document_agent_budget_contract.py index 036f7ae8..d92909f5 100644 --- a/apps/worker/tests/contract/test_document_agent_budget_contract.py +++ b/apps/worker/tests/contract/test_document_agent_budget_contract.py @@ -96,6 +96,16 @@ def test_zip_chunk_schema_preserves_page_memory_node_metadata() -> None: "metadata": { "summary": "page summary", "page_nums": [231], + "page_assets": [ + { + "page_num": 231, + "artifact_ref": "page_citation_assets/page-231.png", + "content_type": "image/png", + "width": 1200, + "height": 1800, + "source": "knowhere-rendered-page-citation-source", + } + ], }, } ] @@ -108,6 +118,17 @@ def test_zip_chunk_schema_preserves_page_memory_node_metadata() -> None: metadata = formatted[0]["metadata"] assert metadata["page_nums"] == [231] + assert metadata["page_assets"] == [ + { + "page_num": 231, + "artifact_ref": "page_citation_assets/page-231.png", + "content_type": "image/png", + "width": 1200, + "height": 1800, + "source": "knowhere-rendered-page-citation-source", + } + ] + assert "page_assets" not in formatted[0] assert "page_image_uris" not in metadata assert "granularity" not in metadata assert "page_indices" not in metadata diff --git a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py index 76521963..c66fc217 100644 --- a/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py +++ b/apps/worker/tests/contract/test_page_memory_node_assembler_contract.py @@ -23,6 +23,7 @@ from shared.services.chunks.dataframe_chunk_converter import dataframe_to_chunks import pandas as pd +from PIL import Image def _same_page_sibling_skeletons() -> list[SectionSkeleton]: @@ -124,6 +125,42 @@ def test_build_node_rows_reuses_tags_without_vlm() -> None: assert leaf_b["extra_metadata"] == {} +def test_build_node_rows_attaches_page_citation_assets_for_rendered_pages(tmp_path) -> None: + page_image = tmp_path / "pages" / "page-231.png" + page_image.parent.mkdir() + Image.new("RGB", (2, 3), color=(255, 255, 255)).save(page_image) + + rows = build_node_rows( + skeletons=_same_page_sibling_skeletons(), + raw_text_by_page={231: "text-231", 232: "text-232"}, + image_path_by_page={231: str(page_image)}, + kind_by_page={}, + tag_by_page={ + 231: PageTagResult(page_index=231, summary="s231", keywords=["k1"]), + 232: PageTagResult(page_index=232, summary="s232", keywords=["k2"]), + }, + filename="demo.pdf", + verdict="page", + budget=None, + vlm_model=None, + ) + + first_page_chunk = next(row for row in rows if row["type"] == "page") + page_assets = first_page_chunk["extra_metadata"]["page_assets"] + + assert page_assets == [ + { + "page_num": 231, + "artifact_ref": "page_citation_assets/page-231.png", + "content_type": "image/png", + "width": 2, + "height": 3, + "source": "knowhere-rendered-page-citation-source", + } + ] + assert (tmp_path / "page_citation_assets" / "page-231.png").is_file() + + def test_build_node_rows_keeps_internal_section_body_pages() -> None: parent = SectionSkeleton( section_path="demo.pdf/4 风险辨识与分级管控", 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 26a1065d..922aff4c 100644 --- a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py +++ b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py @@ -179,10 +179,81 @@ async def test_page_asset_url_is_generated_from_page_nums(monkeypatch) -> None: assert url_map["page-node-1"] == enriched[0]["asset_url"] -def test_result_storage_allows_page_pdf_artifact_refs_not_page_pngs() -> None: +@pytest.mark.asyncio +async def test_page_citation_asset_precedes_lazy_page_pdf_fallback(monkeypatch) -> None: + page_pdf_calls: list[tuple[str, list[int]]] = [] + + def fake_crop_source_pdf_pages(*, job_id, pages): + page_pdf_calls.append((job_id, pages)) + return f"https://assets.example.com/{job_id}/page_pdfs/{'-'.join(map(str, pages))}.pdf" + + class FakeResultStorage: + def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: + if artifact_ref == "page_citation_assets/page-225.png": + return artifact_ref + return None + + def generate_artifact_url( + self, + *, + job_id: str, + artifact_ref: str, + expires_in: int = 3600, + ) -> str: + del expires_in + return f"https://assets.example.com/{job_id}/{artifact_ref}" + + monkeypatch.setattr( + "shared.services.retrieval.hydration.assets.crop_source_pdf_pages", + fake_crop_source_pdf_pages, + ) + monkeypatch.setattr( + "shared.services.retrieval.hydration.assets.get_result_storage", + lambda: FakeResultStorage(), + ) + + rows = [ + { + "chunk_id": "page-node-1", + "chunk_type": "page", + "job_id": "job-1", + "chunk_metadata": { + "page_nums": [225, 226], + "page_assets": [ + { + "page_num": 225, + "artifact_ref": "page_citation_assets/page-225.png", + "content_type": "image/png", + "width": 1200, + "height": 1800, + "source": "knowhere-rendered-page-citation-source", + } + ], + }, + } + ] + + enriched = await enrich_rows_with_retrieval_asset_url( + rows, + log_context="contract", + ) + url_map = await build_retrieval_asset_url_map(rows, log_context="contract") + + expected_url = "https://assets.example.com/job-1/page_citation_assets/page-225.png" + assert enriched[0]["asset_url"] == expected_url + assert enriched[0]["metadata"]["page_assets"][0]["asset_url"] == expected_url + assert url_map["page-node-1"] == expected_url + assert page_pdf_calls == [] + + +def test_result_storage_allows_page_citation_and_page_pdf_artifact_refs_not_debug_page_pngs() -> None: storage = JobResultStorage(results_bucket="test-results") assert storage.normalize_artifact_ref("pages/page-225.png") is None + assert ( + storage.normalize_artifact_ref("page_citation_assets/page-225.png") + == "page_citation_assets/page-225.png" + ) assert ( storage.normalize_artifact_ref("page_pdfs/page-225.pdf") == "page_pdfs/page-225.pdf" @@ -205,10 +276,13 @@ def generate_presigned_url(self, *args, **kwargs) -> str: result_dir = tmp_path / "result" (result_dir / "pages").mkdir(parents=True) + (result_dir / "page_citation_assets").mkdir() (result_dir / "tables").mkdir() (result_dir / "source.pdf").write_bytes(b"source") (result_dir / "pages" / "page-225.png").write_bytes(b"anchored") (result_dir / "pages" / "page-999.png").write_bytes(b"unanchored") + (result_dir / "page_citation_assets" / "page-225.png").write_bytes(b"citation") + (result_dir / "page_citation_assets" / "page-999.png").write_bytes(b"unreferenced") (result_dir / "tables" / "table-1.html").write_text("
") (result_dir / "debug.csv").write_text("debug") zip_path = tmp_path / "result.zip" @@ -224,14 +298,25 @@ def generate_presigned_url(self, *args, **kwargs) -> str: job_id="job-1", result_dir=str(result_dir), zip_file_path=str(zip_path), - artifact_refs={"source.pdf", "pages/page-225.png", "tables/table-1.html"}, + artifact_refs={ + "source.pdf", + "pages/page-225.png", + "page_citation_assets/page-225.png", + "tables/table-1.html", + }, ) - assert set(bundle.raw_files) == {"source.pdf", "tables/table-1.html"} + assert set(bundle.raw_files) == { + "source.pdf", + "page_citation_assets/page-225.png", + "tables/table-1.html", + } assert "results/job-1/source.pdf" in adapter.uploaded_keys assert "results/job-1/tables/table-1.html" in adapter.uploaded_keys + assert "results/job-1/page_citation_assets/page-225.png" in adapter.uploaded_keys assert "results/job-1/pages/page-225.png" not in adapter.uploaded_keys assert "results/job-1/pages/page-999.png" not in adapter.uploaded_keys + assert "results/job-1/page_citation_assets/page-999.png" not in adapter.uploaded_keys assert "results/job-1/debug.csv" not in adapter.uploaded_keys diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/actions.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/actions.py index 3b151378..8858056f 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/actions.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/actions.py @@ -6,6 +6,9 @@ from shared.services.retrieval.agentic.core.budget import budget_status_from_snapshot from shared.services.retrieval.agentic.navigation.path_ledger import PathLedger +from shared.services.retrieval.agentic.navigation.state import ( + RejectionRecord, +) from shared.services.retrieval.search.lexical_text import normalize_section_path from shared.utils.text_utils import truncate_content_preview @@ -68,8 +71,7 @@ def build_legal_actions( collected_paths: list[dict[str, Any]], expanded_scopes: set[str], discovery_hints: list[dict[str, Any]] | None = None, - rejected_paths: set[str] | None = None, - rejected_collect_paths: set[str] | None = None, + rejected: dict[str, RejectionRecord] | None = None, total_images: int, total_tables: int, disabled_asset_types: set[str] | None = None, @@ -79,11 +81,14 @@ def build_legal_actions( covered_paths = _covered_paths(collected_paths) outline_paths = _outline_paths(collected_paths) budget_mode = budget_status_from_snapshot(budget_snapshot) - rejected = {PathLedger.normalize(path) for path in rejected_paths or set()} - rejected_collects = { - normalized - for path in rejected_collect_paths or set() - if (normalized := PathLedger.normalize(path)) + rejection_ledger = rejected or {} + tool_adjudicated_paths = { + path for path, record in rejection_ledger.items() + if record.reason == "tool_adjudicated" + } + navigational_abandoned_paths = { + path for path, record in rejection_ledger.items() + if record.reason == "navigational_abandon" } discovery_scores = _discovery_scores_by_path(discovery_hints or []) scored_items = _score_items(items, discovery_scores) @@ -100,9 +105,12 @@ def build_legal_actions( path = str(item.get("path") or "").strip() if not path or path == "Root": continue - if PathLedger.is_covered(path, covered_paths): + normalized_path = PathLedger.normalize(path) + if PathLedger.is_covered(normalized_path, covered_paths): continue - if PathLedger.is_covered(path, rejected_collects): + # tool_adjudicated rejections are content-level negative and are not + # revived this round (TODO: future strong-signal revival). + if PathLedger.is_covered(normalized_path, tool_adjudicated_paths): continue action_set.add(LegalAction( @@ -112,7 +120,7 @@ def build_legal_actions( target_scope=path, note=( "upgrade outline to full evidence" - if path in outline_paths + if normalized_path in outline_paths else _item_note(item) ), score=float(item.get("relevance_score") or 0.0), @@ -128,18 +136,22 @@ def build_legal_actions( critical_expand = True if item.get("is_leaf"): continue - if path == current_scope: + if normalized_path == current_scope: continue - if current_scope and PathLedger.is_ancestor(path, current_scope): + if current_scope and PathLedger.is_ancestor(normalized_path, current_scope): continue - if path in expanded_scopes: + if normalized_path in expanded_scopes: continue if ( budget_mode == "TIGHT" - and path not in expand_allowlist + and normalized_path not in expand_allowlist ): continue - if path in rejected and not _path_has_discovery_signal(path, discovery_scores): + # EXPAND suppression: navigational_abandon (weak) suppresses unless a + # discovery / lexical signal revives the path. + if normalized_path in navigational_abandoned_paths and not _has_discovery_signal( + normalized_path, discovery_scores + ): continue action_set.add(LegalAction( @@ -166,9 +178,10 @@ def build_legal_actions( seen_discovery_paths.add(path) if PathLedger.is_covered(path, covered_paths): continue + # tool_adjudicated rejections are not revived by discovery this round. # TODO: allow tool-specific LLM adjudicators to revive rejected # collects when validity cannot be determined structurally. - if PathLedger.is_covered(path, rejected_collects): + if PathLedger.is_covered(path, tool_adjudicated_paths): continue if any(action.path == path for action in action_set.collect): continue @@ -236,9 +249,8 @@ def format_agent_state_block( current_scope: str | None, query_intent: str, expanded_scopes: set[str], - rejected_paths: set[str], + rejected: dict[str, RejectionRecord], collected_paths: list[dict[str, Any]], - rejected_collect_paths: set[str] | None = None, prior_tool_result: dict[str, Any] | None, search_context: str, budget_snapshot: dict[str, Any] | None, @@ -267,15 +279,26 @@ def format_agent_state_block( lines.append(f' - "{path}"') else: lines.append("Expanded scopes: none") - rejected_collects = set(rejected_collect_paths or set()) - low_value_rejected = set(rejected_paths) - rejected_collects - if low_value_rejected: - lines.append("Low-value scopes avoided unless revived by discovery:") - for path in sorted(low_value_rejected): + + navigational_abandoned = sorted( + path for path, record in rejected.items() + if record.reason == "navigational_abandon" + ) + tool_adjudicated = sorted( + path for path, record in rejected.items() + if record.reason == "tool_adjudicated" + ) + if navigational_abandoned: + lines.append( + "Scopes avoided (soft; revived by discovery):" + ) + for path in navigational_abandoned: lines.append(f' - "{path}"') - if rejected_collect_paths: - lines.append("Collects rejected by tool reconciliation:") - for path in sorted(rejected_collect_paths): + if tool_adjudicated: + lines.append( + "Collects rejected by tool reconciliation (content-level; not revived):" + ) + for path in tool_adjudicated: lines.append(f' - "{path}"') full_paths, outline_paths = _dedupe_collection_modes(collected_paths) @@ -648,10 +671,11 @@ def _expand_allowlist( return set(candidates[:limit]) -def _path_has_discovery_signal( +def _has_discovery_signal( path: str, discovery_scores: dict[str, float], ) -> bool: + """A path has a discovery signal if it, an ancestor, or a descendant appears in discovery.""" return any( candidate == path or PathLedger.is_ancestor(path, candidate) diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py index bc6c38fa..22d9c273 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/document.py @@ -169,7 +169,6 @@ async def _navigate_collector( Returns (doc_pending_assets, collected_paths). """ - doc_exclude: set[str] = set() doc_discovery_hints = self._discovery_by_doc.get(doc.document_id, []) doc_pending_assets: list[dict[str, Any]] = [] nav_state = NavigationState( @@ -201,8 +200,7 @@ async def _navigate_collector( nav_state.step_count += 1 before_scope = nav_state.current_scope expanded_before = set(nav_state.expanded_scopes) - rejected_before = set(nav_state.rejected_paths) - rejected_collect_before = set(nav_state.rejected_collect_paths) + rejected_before = dict(nav_state.rejected) collected_before_count = len(nav_state.collected_paths) doc_llm_fn = self._llm_budget.for_document( @@ -226,13 +224,11 @@ async def _navigate_collector( namespace=self._namespace, doc_name=doc_name, scope_path=nav_state.current_scope, - exclude_paths=doc_exclude, budget_snapshot=self._state.ledger.snapshot() if self._state.ledger else None, nav_trace=nav_state.nav_trace if nav_state.nav_trace else None, collected_paths=nav_state.collected_paths, expanded_scopes=nav_state.expanded_scopes, - rejected_paths=nav_state.rejected_paths, - rejected_collect_paths=nav_state.rejected_collect_paths, + rejected=nav_state.rejected, disabled_asset_types=self._disabled_asset_types | nav_state.blocked_asset_types_for_scope( nav_state.current_scope ), @@ -292,8 +288,11 @@ async def _navigate_collector( if rejected_collects: nav_result.collect = collect_reconcile["accepted_collects"] for path in rejected_collects: - nav_state.mark_rejected_collect(path) - doc_exclude.add(path) + nav_state.mark_rejected_collect( + path, + step=nav_state.step_count, + detail=collect_reconcile.get("reason", ""), + ) logger.info( " agentic: tool reconciliation rejected collects: " f"{rejected_collects}" @@ -309,10 +308,10 @@ async def _navigate_collector( scope_context=nav_state.current_scope, ) collected_in_step.append(path) - # Outline collections should NOT exclude children — the intent - # is "see structure, then drill deeper for full content". - if coll_item.get("hydrate_mode") != "outline": - doc_exclude.add(path) + # Outline collections keep children visible — the intent is + # "see structure, then drill deeper for full content". + # Coverage / action filtering now derives from collected_paths + # via the navigation state ledger (no physical exclusion). # ── Process navigation action ──────────────────────────────── should_break = False @@ -336,7 +335,11 @@ async def _navigate_collector( else: back_target = nav_result.back_to # None = root if PathLedger.valid_back_target(nav_state.current_scope, back_target): - nav_state.mark_rejected_if_unproductive(nav_state.current_scope) + nav_state.mark_rejected_if_unproductive( + nav_state.current_scope, + step=nav_state.step_count, + detail="back_from_unproductive_scope", + ) nav_state.current_scope = back_target else: logger.warning( @@ -363,7 +366,6 @@ async def _navigate_collector( before_scope=before_scope, expanded_before=expanded_before, rejected_before=rejected_before, - rejected_collect_before=rejected_collect_before, collected_before_count=collected_before_count, ) trace_entry: dict[str, Any] = { diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/state.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/state.py index 841059e5..83b6ac57 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/state.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/state.py @@ -2,45 +2,111 @@ from __future__ import annotations from dataclasses import dataclass, field -from typing import Any +from typing import Any, Literal from shared.services.retrieval.agentic.navigation.path_ledger import PathLedger +RejectReason = Literal["tool_adjudicated", "navigational_abandon"] + +# Strength ordering: a stronger reason overrides a weaker one so we keep the +# most informative record for a given path. +_REASON_STRENGTH: dict[RejectReason, int] = { + "tool_adjudicated": 2, + "navigational_abandon": 1, +} + + +@dataclass +class RejectionRecord: + """A single rejection entry in the navigation state ledger. + + Two reasons are distinguished: + + - ``tool_adjudicated``: a SEARCH_* tool reconciliation proved the path has + no matching assets. Strong (content-level) negative signal. Not revived + this round; future strong-signal revival is a TODO. + - ``navigational_abandon``: BACK left an unproductive scope. Weak (soft) + signal; any discovery / lexical hit revives the path. + """ + + path: str + reason: RejectReason + step: int + detail: str = "" + @dataclass class NavigationState: - """Mutable state for one document navigation loop.""" + """Mutable state for one document navigation loop. + + Path state is tracked in two orthogonal dimensions: + + - **Coverage** (positive "already taken as evidence"): derived from + ``collected_paths`` via :meth:`covered_paths` / :meth:`outline_paths`. + - **Rejection** (negative "evaluated, not taken"): a single labelled + ledger in :attr:`rejected` keyed by normalized path. Replaces the + former ``rejected_paths`` / ``rejected_collect_paths`` dual sets. + """ document_id: str document_name: str job_result_id: str current_scope: str | None = None expanded_scopes: set[str] = field(default_factory=set) - rejected_paths: set[str] = field(default_factory=set) - rejected_collect_paths: set[str] = field(default_factory=set) + rejected: dict[str, RejectionRecord] = field(default_factory=dict) collected_paths: list[dict[str, Any]] = field(default_factory=list) nav_trace: list[dict[str, Any]] = field(default_factory=list) tool_history: list[dict[str, Any]] = field(default_factory=list) blocked_asset_searches: set[str] = field(default_factory=set) step_count: int = 0 + # ── Coverage helpers (single source: collected_paths) ──────────────── + + def covered_paths(self) -> set[str]: + """Full-evidence paths (hydrate_mode != 'outline').""" + return { + PathLedger.normalize(str(item.get("path") or "")) + for item in self.collected_paths + if item.get("path") and item.get("hydrate_mode") != "outline" + } + + def outline_paths(self) -> set[str]: + """Outline-only paths; excludes any path also collected as full.""" + full = self.covered_paths() + return { + PathLedger.normalize(str(item.get("path") or "")) + for item in self.collected_paths + if item.get("path") + and item.get("hydrate_mode") == "outline" + and PathLedger.normalize(str(item.get("path") or "")) not in full + } + + # ── Snapshot / delta for replayable traces ────────────────────────── + def snapshot_delta( self, *, before_scope: str | None, expanded_before: set[str], - rejected_before: set[str], - rejected_collect_before: set[str], + rejected_before: dict[str, RejectionRecord], collected_before_count: int, ) -> dict[str, Any]: + rejected_added: list[dict[str, Any]] = [] + for path, record in self.rejected.items(): + if path in rejected_before: + continue + rejected_added.append({ + "path": record.path, + "reason": record.reason, + "step": record.step, + "detail": record.detail, + }) + rejected_added.sort(key=lambda item: (item["step"], item["path"])) return { "current_scope_before": before_scope or "root", "current_scope_after": self.current_scope or "root", "expanded_added": sorted(self.expanded_scopes - expanded_before), - "rejected_added": sorted(self.rejected_paths - rejected_before), - "rejected_collect_added": sorted( - self.rejected_collect_paths - rejected_collect_before - ), + "rejected_added": rejected_added, "collected_added": [ item.get("path", "") for item in self.collected_paths[collected_before_count:] @@ -48,6 +114,8 @@ def snapshot_delta( ], } + # ── Mutation helpers ──────────────────────────────────────────────── + def add_collected( self, item: dict[str, Any], @@ -66,13 +134,35 @@ def mark_expanded(self, path: str | None) -> None: if normalized: self.expanded_scopes.add(normalized) - def mark_rejected_collect(self, path: str | None) -> None: + def mark_rejected_collect( + self, + path: str | None, + *, + step: int, + detail: str = "", + ) -> None: + """Record a tool-adjudicated rejection (strong, content-level).""" normalized = PathLedger.normalize(path) - if normalized: - self.rejected_paths.add(normalized) - self.rejected_collect_paths.add(normalized) + if not normalized: + return + self._upsert_rejection( + normalized, + reason="tool_adjudicated", + step=step, + detail=detail, + ) - def mark_rejected_if_unproductive(self, path: str | None) -> None: + def mark_rejected_if_unproductive( + self, + path: str | None, + *, + step: int, + detail: str = "", + ) -> None: + """Record a soft navigational abandon when leaving an unproductive scope. + + Only written when no stronger record exists for the path. + """ normalized = PathLedger.normalize(path) if not normalized: return @@ -81,8 +171,40 @@ def mark_rejected_if_unproductive(self, path: str | None) -> None: and PathLedger.is_same_or_descendant(item.get("path"), normalized) for item in self.collected_paths ) - if not has_full_collect: - self.rejected_paths.add(normalized) + if has_full_collect: + return + self._upsert_rejection( + normalized, + reason="navigational_abandon", + step=step, + detail=detail, + ) + + def _upsert_rejection( + self, + normalized_path: str, + *, + reason: RejectReason, + step: int, + detail: str, + ) -> None: + existing = self.rejected.get(normalized_path) + if existing is not None and _REASON_STRENGTH[existing.reason] >= _REASON_STRENGTH[reason]: + # Keep the stronger prior record. + return + self.rejected[normalized_path] = RejectionRecord( + path=normalized_path, + reason=reason, + step=step, + detail=detail, + ) + + def rejected_paths_with_reason(self, reason: RejectReason) -> set[str]: + """All paths rejected with a specific reason label.""" + return { + path for path, record in self.rejected.items() + if record.reason == reason + } def blocked_asset_types_for_scope(self, scope: str | None) -> set[str]: prefix = f"{PathLedger.normalize(scope) or 'root'}:" diff --git a/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py index 2873a2e2..066fa130 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/navigation/tools.py @@ -27,6 +27,7 @@ format_agent_state_block, ) from shared.services.retrieval.agentic.navigation.section_tree import load_child_sections +from shared.services.retrieval.agentic.navigation.state import RejectionRecord from shared.services.retrieval.agentic.core.types import DocTreeNode, NavigateStepResult from shared.services.retrieval.llm_adapter import LLMFn @@ -48,8 +49,7 @@ async def navigate_step( nav_trace: list[dict[str, Any]] | None = None, collected_paths: list[dict[str, Any]] | None = None, expanded_scopes: set[str] | None = None, - rejected_paths: set[str] | None = None, - rejected_collect_paths: set[str] | None = None, + rejected: dict[str, RejectionRecord] | None = None, disabled_asset_types: set[str] | None = None, discovery_hints: list[dict[str, Any]] | None = None, section_rows: list | None = None, @@ -89,14 +89,14 @@ async def navigate_step( expanded_path_set = set(expanded_scopes or _expanded_paths_from_trace(nav_trace or [])) if scope_path: expanded_path_set.add(scope_path) + rejection_ledger = rejected or {} provisional_action_set = build_legal_actions( items=items, current_scope=scope_path, collected_paths=collected_paths or [], expanded_scopes=expanded_path_set, discovery_hints=discovery_hints, - rejected_paths=rejected_paths or set(), - rejected_collect_paths=rejected_collect_paths or set(), + rejected=rejection_ledger, total_images=total_images, total_tables=total_tables, disabled_asset_types=disabled_asset_types or set(), @@ -137,8 +137,7 @@ async def navigate_step( collected_paths=collected_paths or [], expanded_scopes=expanded_path_set, discovery_hints=discovery_hints, - rejected_paths=rejected_paths or set(), - rejected_collect_paths=rejected_collect_paths or set(), + rejected=rejection_ledger, total_images=total_images, total_tables=total_tables, disabled_asset_types=disabled_asset_types or set(), @@ -166,16 +165,17 @@ async def navigate_step( "search": [item.id for item in action_set.search], "finish": [action_set.finish.id] if action_set.finish else [], }, - "rejected_paths": sorted(rejected_paths or set()), - "rejected_collect_paths": sorted(rejected_collect_paths or set()), + "rejected": { + path: {"reason": record.reason, "step": record.step} + for path, record in rejection_ledger.items() + }, } agent_state_block = format_agent_state_block( current_scope=scope_path, query_intent=query_intent, expanded_scopes=expanded_path_set, - rejected_paths=rejected_paths or set(), + rejected=rejection_ledger, collected_paths=collected_paths or [], - rejected_collect_paths=rejected_collect_paths or set(), prior_tool_result=prior_tool_result, search_context=search_context, budget_snapshot=adjusted_snapshot, diff --git a/packages/shared-python/shared/services/retrieval/agentic/tools.py b/packages/shared-python/shared/services/retrieval/agentic/tools.py index d9d6b331..65e16187 100644 --- a/packages/shared-python/shared/services/retrieval/agentic/tools.py +++ b/packages/shared-python/shared/services/retrieval/agentic/tools.py @@ -90,8 +90,7 @@ async def navigate_step( nav_trace: list[dict[str, Any]] | None = None, collected_paths: list[dict[str, Any]] | None = None, expanded_scopes: set[str] | None = None, - rejected_paths: set[str] | None = None, - rejected_collect_paths: set[str] | None = None, + rejected: dict[str, Any] | None = None, disabled_asset_types: set[str] | None = None, discovery_hints: list[dict[str, Any]] | None = None, section_rows: list | None = None, @@ -114,8 +113,7 @@ async def navigate_step( nav_trace=nav_trace, collected_paths=collected_paths, expanded_scopes=expanded_scopes, - rejected_paths=rejected_paths, - rejected_collect_paths=rejected_collect_paths, + rejected=rejected, disabled_asset_types=disabled_asset_types, discovery_hints=discovery_hints, section_rows=section_rows, diff --git a/packages/shared-python/shared/services/retrieval/hydration/assets.py b/packages/shared-python/shared/services/retrieval/hydration/assets.py index 8d305512..8b679e5e 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/assets.py +++ b/packages/shared-python/shared/services/retrieval/hydration/assets.py @@ -41,6 +41,40 @@ def _resolve_asset_request(row: dict[str, Any]) -> tuple[str, str] | None: return job_id, artifact_ref +def _metadata_for_row(row: dict[str, Any]) -> dict[str, Any]: + metadata = row.get("chunk_metadata") or row.get("metadata") or {} + return metadata if isinstance(metadata, dict) else {} + + +def _resolve_page_citation_asset_request(row: dict[str, Any]) -> tuple[str, str] | None: + job_id = str(row.get("job_id") or "").strip() + if not job_id or not _is_page_row(row): + return None + + for page_asset in _iter_page_assets(_metadata_for_row(row).get("page_assets")): + artifact_ref = _normalize_artifact_ref(page_asset.get("artifact_ref")) + if artifact_ref is None: + continue + return job_id, artifact_ref + return None + + +def _resolve_direct_page_citation_asset_url(row: dict[str, Any]) -> str | None: + if not _is_page_row(row): + return None + for page_asset in _iter_page_assets(_metadata_for_row(row).get("page_assets")): + asset_url = str(page_asset.get("asset_url") or "").strip() + if asset_url: + return asset_url + return None + + +def _iter_page_assets(raw_assets: object) -> list[dict[str, Any]]: + if not isinstance(raw_assets, list): + return [] + return [item for item in raw_assets if isinstance(item, dict)] + + def _coerce_page_nums(value: object) -> list[int]: if isinstance(value, list): raw_values = value @@ -63,10 +97,10 @@ def _resolve_page_pdf_request(row: dict[str, Any]) -> PagePdfRequestKey | None: if not job_id or not _is_page_row(row): return None - metadata = row.get("chunk_metadata") or row.get("metadata") or {} - if not isinstance(metadata, dict): + if _resolve_direct_page_citation_asset_url(row) or _resolve_page_citation_asset_request(row): return None + metadata = _metadata_for_row(row) pages = _coerce_page_nums(metadata.get("page_nums") or row.get("page_nums")) if not pages: return None @@ -96,6 +130,79 @@ async def _generate_retrieval_asset_url( return None +async def _generate_page_citation_asset_url( + *, + row: dict[str, Any], + log_context: str, +) -> str | None: + direct_url = _resolve_direct_page_citation_asset_url(row) + if direct_url: + return direct_url + + request = _resolve_page_citation_asset_request(row) + if request is None: + return None + + job_id, artifact_ref = request + try: + return get_result_storage().generate_artifact_url( + job_id=job_id, + artifact_ref=artifact_ref, + ) + except Exception as exc: + logger.warning( + f"Failed to generate {log_context} page citation asset URL (ignored): {exc}" + ) + return None + + +async def _enrich_page_assets( + *, + row: dict[str, Any], + log_context: str, +) -> list[dict[str, Any]]: + if not _is_page_row(row): + return [] + page_assets = _iter_page_assets(_metadata_for_row(row).get("page_assets")) + if not page_assets: + return [] + + enriched_assets: list[dict[str, Any]] = [] + job_id = str(row.get("job_id") or "").strip() + for page_asset in page_assets: + enriched = dict(page_asset) + if not str(enriched.get("asset_url") or "").strip() and job_id: + artifact_ref = _normalize_artifact_ref(enriched.get("artifact_ref")) + if artifact_ref is not None: + try: + asset_url = get_result_storage().generate_artifact_url( + job_id=job_id, + artifact_ref=artifact_ref, + ) + except Exception as exc: + logger.warning( + f"Failed to generate {log_context} page citation metadata URL (ignored): {exc}" + ) + asset_url = None + if asset_url: + enriched["asset_url"] = asset_url + enriched_assets.append(enriched) + return enriched_assets + + +def _attach_enriched_page_assets( + *, + row: dict[str, Any], + page_assets: list[dict[str, Any]], +) -> None: + if not page_assets: + return + metadata = dict(_metadata_for_row(row)) + metadata["page_assets"] = page_assets + row["chunk_metadata"] = metadata + row["metadata"] = metadata + + async def _generate_page_pdf_asset_url_for_request( request: PagePdfRequestKey, *, @@ -151,16 +258,27 @@ async def enrich_rows_with_retrieval_asset_url( enriched_rows: list[dict[str, Any]] = [] for row in rows: enriched = dict(row) + page_assets = await _enrich_page_assets( + row=row, + log_context=log_context, + ) + _attach_enriched_page_assets(row=enriched, page_assets=page_assets) asset_url = await _generate_retrieval_asset_url( row=row, log_context=log_context, ) if asset_url: enriched["asset_url"] = asset_url + page_citation_asset_url = await _generate_page_citation_asset_url( + row=row, + log_context=log_context, + ) + if page_citation_asset_url: + enriched["asset_url"] = page_citation_asset_url page_pdf_url = None if page_request := _resolve_page_pdf_request(row): page_pdf_url = page_pdf_urls.get(page_request) - if page_pdf_url: + if page_pdf_url and not page_citation_asset_url: enriched["asset_url"] = page_pdf_url enriched_rows.append(enriched) return enriched_rows @@ -178,6 +296,14 @@ async def build_retrieval_asset_url_map( if not chunk_id: continue + page_citation_asset_url = await _generate_page_citation_asset_url( + row=row, + log_context=log_context, + ) + if page_citation_asset_url: + url_map[chunk_id] = page_citation_asset_url + continue + page_pdf_url = None if page_request := _resolve_page_pdf_request(row): page_pdf_url = page_pdf_urls.get(page_request) diff --git a/packages/shared-python/shared/services/storage/result_storage.py b/packages/shared-python/shared/services/storage/result_storage.py index 529dde30..51d3ab3c 100644 --- a/packages/shared-python/shared/services/storage/result_storage.py +++ b/packages/shared-python/shared/services/storage/result_storage.py @@ -13,7 +13,7 @@ _EXCLUDED_FILE_NAMES = {".DS_Store", "Thumbs.db"} _EXCLUDED_DIR_NAMES = {"tmp", "temp", "__pycache__"} -_CLIENT_ARTIFACT_DIRS = {"images", "tables", "page_pdfs"} +_CLIENT_ARTIFACT_DIRS = {"images", "tables", "page_pdfs", "page_citation_assets"} _INTERNAL_RAW_FILES = {"source.pdf"} @@ -40,6 +40,14 @@ def generate_artifact_url( ) -> str | None: raise NotImplementedError + def generate_raw_file_url( + self, *, job_id: str, relative_path: str, expires_in: int = 3600 + ) -> str | None: + raise NotImplementedError + + def verify_raw_exists(self, *, job_id: str, relative_path: str) -> bool: + raise NotImplementedError + def normalize_artifact_ref(self, artifact_ref: str | None) -> str | None: raise NotImplementedError @@ -133,6 +141,14 @@ def verify_raw_exists(self, *, job_id: str, relative_path: str) -> bool: result = self._job_file_storage.verify_exists(key, bucket=self.results_bucket) return bool(result.get("exists")) + def generate_raw_file_url( + self, *, job_id: str, relative_path: str, expires_in: int = 3600 + ) -> str | None: + return self.generate_url( + storage_key=self.build_raw_key(job_id=job_id, relative_path=relative_path), + expires_in=expires_in, + ) + def upload_raw_file( self, *, diff --git a/packages/shared-python/shared/services/storage/zip_chunk_schema.py b/packages/shared-python/shared/services/storage/zip_chunk_schema.py index 6f8fe765..2cb36b43 100644 --- a/packages/shared-python/shared/services/storage/zip_chunk_schema.py +++ b/packages/shared-python/shared/services/storage/zip_chunk_schema.py @@ -108,16 +108,19 @@ def format_chunks( elif chunk_type == "page": metadata["keywords"] = existing_metadata.get("keywords") or [] metadata["connect_to"] = existing_metadata.get("connect_to") or [] + page_assets = _normalize_page_assets(existing_metadata.get("page_assets")) + if page_assets: + metadata["page_assets"] = page_assets - formatted.append( - { - "chunk_id": chunk_id, - "type": chunk_type, - "content": content, - "path": path, - "metadata": metadata, - } - ) + formatted_chunk = { + "chunk_id": chunk_id, + "type": chunk_type, + "content": content, + "path": path, + "metadata": metadata, + } + + formatted.append(formatted_chunk) return formatted @@ -190,3 +193,40 @@ def _resolve_table_file_path( table_name = path.split("/")[-1] if "/" in path else f"table_{chunk_id}.html" return f"tables/{table_name}" + + +def _normalize_page_assets(value: Any) -> list[dict[str, Any]]: + if not isinstance(value, list): + return [] + assets: list[dict[str, Any]] = [] + for item in value: + if not isinstance(item, dict): + continue + page_num = _positive_int_or_none(item.get("page_num")) + artifact_ref = str(item.get("artifact_ref") or "").strip() + content_type = str(item.get("content_type") or "").strip() + source = str(item.get("source") or "").strip() + if page_num is None or not artifact_ref or not content_type or not source: + continue + asset = { + "page_num": page_num, + "artifact_ref": artifact_ref, + "content_type": content_type, + "source": source, + } + if (asset_url := str(item.get("asset_url") or "").strip()): + asset["asset_url"] = asset_url + if (width := _positive_int_or_none(item.get("width"))) is not None: + asset["width"] = width + if (height := _positive_int_or_none(item.get("height"))) is not None: + asset["height"] = height + assets.append(asset) + return assets + + +def _positive_int_or_none(value: Any) -> int | None: + try: + number = int(value) + except (TypeError, ValueError): + return None + return number if number > 0 else None