From a457ca802a8cb28a970c103d3697e3dd88e43509 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Thu, 2 Jul 2026 22:15:38 +0800 Subject: [PATCH 1/4] feat: add page citation source API --- apps/api/app/api/v2/api_v2.py | 2 + apps/api/app/api/v2/routes/documents.py | 37 +++- .../app/repositories/document_repository.py | 25 +++ .../services/documents/lifecycle_service.py | 57 ++++- .../tests/contract/test_documents_contract.py | 198 +++++++++++++++++- .../test_page_memory_parse_track_contract.py | 9 + .../shared/services/storage/result_storage.py | 16 ++ 7 files changed, 339 insertions(+), 5 deletions(-) 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..31a9b044 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: @@ -75,9 +79,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 +252,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, *, diff --git a/apps/api/tests/contract/test_documents_contract.py b/apps/api/tests/contract/test_documents_contract.py index 92a80bc2..1395720d 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,21 @@ 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) + + @pytest.mark.asyncio async def test_should_list_only_the_authenticated_users_documents_for_the_effective_namespace( developer_api_client_factory: Callable[ @@ -681,6 +701,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[ 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/packages/shared-python/shared/services/storage/result_storage.py b/packages/shared-python/shared/services/storage/result_storage.py index 529dde30..ab607b14 100644 --- a/packages/shared-python/shared/services/storage/result_storage.py +++ b/packages/shared-python/shared/services/storage/result_storage.py @@ -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, *, From 45ce5459db0b58fee66339ffc98f4ba1cb91761b Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 3 Jul 2026 17:40:37 +0800 Subject: [PATCH 2/4] Add server page citation assets --- .../services/documents/lifecycle_service.py | 100 ++++++++++++- .../tests/contract/test_documents_contract.py | 106 ++++++++++++++ .../document_ingestion/artifact_refs.py | 20 ++- .../services/page_memory/node_assembler.py | 91 +++++++++++- ...est_page_memory_node_assembler_contract.py | 37 +++++ .../test_page_memory_retrieval_contract.py | 91 +++++++++++- .../retrieval/execution/reference_resolver.py | 3 + .../services/retrieval/hydration/assets.py | 132 +++++++++++++++++- .../shared/services/storage/result_storage.py | 2 +- .../services/storage/zip_chunk_schema.py | 64 +++++++-- 10 files changed, 625 insertions(+), 21 deletions(-) diff --git a/apps/api/app/services/documents/lifecycle_service.py b/apps/api/app/services/documents/lifecycle_service.py index 31a9b044..1c13679c 100644 --- a/apps/api/app/services/documents/lifecycle_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -57,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, @@ -312,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, @@ -322,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, @@ -332,6 +425,9 @@ def _chunk_payload( ), "created_at": _datetime_payload(chunk.created_at), } + if include_asset_urls and page_assets: + payload["page_assets"] = page_assets + 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 1395720d..5463999a 100644 --- a/apps/api/tests/contract/test_documents_contract.py +++ b/apps/api/tests/contract/test_documents_contract.py @@ -518,6 +518,21 @@ def _upload_page_citation_source(*, job_id: str) -> None: 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[ @@ -1010,6 +1025,97 @@ 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] + page_assets = cast(list[dict[str, object]], chunk["page_assets"]) + 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 page_assets[0]["asset_url"] == expected_asset_url + assert metadata_page_assets[0]["asset_url"] == expected_asset_url + 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/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_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/execution/reference_resolver.py b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py index d5e2c358..f292247b 100644 --- a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py +++ b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py @@ -116,6 +116,9 @@ def _merge_reference_asset_url( if row is not None: if row.get("asset_url"): merged["asset_url"] = row["asset_url"] + metadata = row.get("metadata") or row.get("chunk_metadata") + if isinstance(metadata, dict): + merged["metadata"] = metadata merged_refs.append(merged) return merged_refs 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 ab607b14..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"} 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..60e82043 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,23 @@ def format_chunks( elif chunk_type == "page": metadata["keywords"] = existing_metadata.get("keywords") or [] metadata["connect_to"] = existing_metadata.get("connect_to") or [] - - formatted.append( - { - "chunk_id": chunk_id, - "type": chunk_type, - "content": content, - "path": path, - "metadata": metadata, - } - ) + page_assets = _normalize_page_assets(existing_metadata.get("page_assets")) + if page_assets: + metadata["page_assets"] = page_assets + + formatted_chunk = { + "chunk_id": chunk_id, + "type": chunk_type, + "content": content, + "path": path, + "metadata": metadata, + } + if chunk_type == "page": + page_assets = _normalize_page_assets(metadata.get("page_assets")) + if page_assets: + formatted_chunk["page_assets"] = page_assets + + formatted.append(formatted_chunk) return formatted @@ -190,3 +197,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 From 4f7b9c35e04ff6b0bbecb0ed499d49c283f55807 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Fri, 3 Jul 2026 21:14:10 +0800 Subject: [PATCH 3/4] Preserve page citation assets in metadata --- .../services/documents/lifecycle_service.py | 2 -- .../tests/contract/test_documents_contract.py | 3 +-- .../test_document_agent_budget_contract.py | 21 +++++++++++++++++++ .../services/storage/zip_chunk_schema.py | 4 ---- 4 files changed, 22 insertions(+), 8 deletions(-) diff --git a/apps/api/app/services/documents/lifecycle_service.py b/apps/api/app/services/documents/lifecycle_service.py index 1c13679c..1bd5ae3a 100644 --- a/apps/api/app/services/documents/lifecycle_service.py +++ b/apps/api/app/services/documents/lifecycle_service.py @@ -425,8 +425,6 @@ def _chunk_payload( ), "created_at": _datetime_payload(chunk.created_at), } - if include_asset_urls and page_assets: - payload["page_assets"] = page_assets return payload async def archive_document( diff --git a/apps/api/tests/contract/test_documents_contract.py b/apps/api/tests/contract/test_documents_contract.py index 5463999a..6b8e33b0 100644 --- a/apps/api/tests/contract/test_documents_contract.py +++ b/apps/api/tests/contract/test_documents_contract.py @@ -1088,7 +1088,6 @@ async def test_should_include_page_citation_asset_urls_in_document_chunk_metadat chunk = cast(list[dict[str, object]], response.json()["chunks"])[0] default_chunk = cast(list[dict[str, object]], default_response.json()["chunks"])[0] - page_assets = cast(list[dict[str, object]], chunk["page_assets"]) metadata = cast(dict[str, object], chunk["metadata"]) metadata_page_assets = cast(list[dict[str, object]], metadata["page_assets"]) @@ -1097,8 +1096,8 @@ async def test_should_include_page_citation_asset_urls_in_document_chunk_metadat f"results/{revision['job_id']}/{page_asset_ref}" "?method=GET&expires_in=604800" ) - assert page_assets[0]["asset_url"] == expected_asset_url 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], 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/packages/shared-python/shared/services/storage/zip_chunk_schema.py b/packages/shared-python/shared/services/storage/zip_chunk_schema.py index 60e82043..2cb36b43 100644 --- a/packages/shared-python/shared/services/storage/zip_chunk_schema.py +++ b/packages/shared-python/shared/services/storage/zip_chunk_schema.py @@ -119,10 +119,6 @@ def format_chunks( "path": path, "metadata": metadata, } - if chunk_type == "page": - page_assets = _normalize_page_assets(metadata.get("page_assets")) - if page_assets: - formatted_chunk["page_assets"] = page_assets formatted.append(formatted_chunk) From b8c2034e719ca2dcd3ad5a2a45d5fd3f4bf5a020 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Sun, 5 Jul 2026 12:52:48 +0800 Subject: [PATCH 4/4] fix: keep workflow references minimal --- .../shared/services/retrieval/execution/reference_resolver.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py index f292247b..d5e2c358 100644 --- a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py +++ b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py @@ -116,9 +116,6 @@ def _merge_reference_asset_url( if row is not None: if row.get("asset_url"): merged["asset_url"] = row["asset_url"] - metadata = row.get("metadata") or row.get("chunk_metadata") - if isinstance(metadata, dict): - merged["metadata"] = metadata merged_refs.append(merged) return merged_refs