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_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_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/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/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("