From ec82460051b83c68ef51cf83bd76ebd801e88591 Mon Sep 17 00:00:00 2001 From: suguanYang Date: Tue, 7 Jul 2026 22:35:19 +0800 Subject: [PATCH] fix: include page citation assets in result zip --- .../contract/test_parse_task_contract.py | 105 ++++++++++++++++++ .../services/storage/zip_package_writer.py | 6 + .../services/storage/zip_result_resources.py | 92 +++++++++++++++ .../services/storage/zip_result_service.py | 1 + 4 files changed, 204 insertions(+) diff --git a/apps/worker/tests/contract/test_parse_task_contract.py b/apps/worker/tests/contract/test_parse_task_contract.py index f25ab383b..59cde7310 100644 --- a/apps/worker/tests/contract/test_parse_task_contract.py +++ b/apps/worker/tests/contract/test_parse_task_contract.py @@ -1,11 +1,14 @@ from __future__ import annotations +import json from concurrent.futures import ThreadPoolExecutor from pathlib import Path from uuid import uuid4 +import pandas as pd import pytest +from app.services.document_parser.orchestration.parse_output import ParseOutput from support.worker_parse_contract import WorkerParseContract _REPO_ROOT: Path = Path(__file__).resolve().parents[4] @@ -118,6 +121,108 @@ def test_parse_task_should_process_uploaded_file_through_real_contract_boundarie assert contract.find_task_workspaces(tmp_path, job["job_id"]) == [] +def test_parse_task_result_zip_includes_page_citation_assets( + worker_contract_environment: None, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + contract = WorkerParseContract.create() + contract.use_workspace_root(monkeypatch, tmp_path) + contract.use_billing(monkeypatch, is_enabled=False) + + source_pdf = tmp_path / "page-citation-source.pdf" + _write_blank_pdf(source_pdf, page_count=1) + + job = contract.create_file_job( + source_file_name=source_pdf.name, + job_id_prefix="job_parse_page_citation", + ) + contract.upload_source_file( + local_file_path=source_pdf, + s3_key=job["s3_key"], + ) + + def fake_execute_document_parse( + *, + job_id: str, + job_context: object, + prepared_source: object, + output_dir: str, + ) -> ParseOutput: + del job_id, job_context, prepared_source + output_path = Path(output_dir) + page_citation_assets_dir = output_path / "page_citation_assets" + page_citation_assets_dir.mkdir(parents=True, exist_ok=True) + (page_citation_assets_dir / "page-1.png").write_bytes(b"page image") + (page_citation_assets_dir / "page-999.png").write_bytes( + b"unreferenced page image" + ) + parsed_df = pd.DataFrame( + [ + { + "know_id": "page-node-1", + "type": "page", + "content": "Page text", + "path": f"{source_pdf.name}/Root/Introduction", + "length": 9, + "keywords": "", + "summary": "Page summary", + "tokens": "", + "connectto": "", + "page_nums": "1", + "extra_metadata": json.dumps( + { + "page_assets": [ + { + "page_num": 1, + "artifact_ref": "page_citation_assets/page-1.png", + "content_type": "image/png", + "source": "knowhere-rendered-page-citation-source", + "width": 1200, + "height": 1800, + } + ], + } + ), + } + ] + ) + return ParseOutput(output_dir=output_dir, parsed_df=parsed_df) + + monkeypatch.setattr( + "app.services.document_ingestion.processing_run.execute_document_parse", + fake_execute_document_parse, + ) + + celery_result = contract.enqueue_parse_task( + job_id=job["job_id"], + user_id=job["user_id"], + ) + + assert celery_result.successful() + assert celery_result.result["status"] == "success" + + observed = contract.observe_successful_job(job["job_id"]) + result_row = observed["result"] + result_zip = contract.read_result_zip( + result_s3_key=result_row["result_s3_key"], + tmp_path=tmp_path, + ) + + assert "page_citation_assets/page-1.png" in result_zip["members"] + assert "page_citation_assets/page-999.png" not in result_zip["members"] + assert result_zip["chunks"]["chunks"][0]["metadata"]["page_assets"] == [ + { + "page_num": 1, + "artifact_ref": "page_citation_assets/page-1.png", + "content_type": "image/png", + "source": "knowhere-rendered-page-citation-source", + "width": 1200, + "height": 1800, + } + ] + + def test_parse_task_should_charge_user_when_billing_is_enabled( worker_contract_environment: None, monkeypatch: pytest.MonkeyPatch, diff --git a/packages/shared-python/shared/services/storage/zip_package_writer.py b/packages/shared-python/shared/services/storage/zip_package_writer.py index f3ec2337b..34ec35477 100644 --- a/packages/shared-python/shared/services/storage/zip_package_writer.py +++ b/packages/shared-python/shared/services/storage/zip_package_writer.py @@ -22,6 +22,7 @@ class ZipPackageWriteRequest: formatted_chunks: list[dict[str, Any]] image_files: tuple[ZipResourceFileInfo, ...] table_files: tuple[ZipResourceFileInfo, ...] + page_citation_files: tuple[ZipResourceFileInfo, ...] doc_nav: dict[str, Any] | None manifest: dict[str, Any] temp_dir: str | None @@ -75,6 +76,11 @@ def write(self, request: ZipPackageWriteRequest) -> ZipPackageArtifact: self._write_resource_files(zip_file, request.image_files, label="Image") self._write_resource_files(zip_file, request.table_files, label="Table") + self._write_resource_files( + zip_file, + request.page_citation_files, + label="Page citation asset", + ) if request.doc_nav is not None: doc_nav_json = json.dumps(request.doc_nav, ensure_ascii=False, indent=2) diff --git a/packages/shared-python/shared/services/storage/zip_result_resources.py b/packages/shared-python/shared/services/storage/zip_result_resources.py index 0a2a88473..797b36eed 100644 --- a/packages/shared-python/shared/services/storage/zip_result_resources.py +++ b/packages/shared-python/shared/services/storage/zip_result_resources.py @@ -18,6 +18,7 @@ class ZipPackageResources: image_files: tuple[ZipResourceFileInfo, ...] table_files: tuple[ZipResourceFileInfo, ...] + page_citation_files: tuple[ZipResourceFileInfo, ...] @property def image_files_map(self) -> dict[str, ZipResourceFileInfo]: @@ -39,9 +40,16 @@ def collect( ) -> ZipPackageResources: images_dir = os.path.join(add_dir, "images") tables_dir = os.path.join(add_dir, "tables") + page_citation_assets_dir = os.path.join(add_dir, "page_citation_assets") return ZipPackageResources( image_files=tuple(self._collect_image_files(chunks, images_dir)), table_files=tuple(self._collect_table_files(chunks, tables_dir)), + page_citation_files=tuple( + self._collect_page_citation_files( + chunks, + page_citation_assets_dir, + ) + ), ) def _collect_image_files( @@ -213,6 +221,51 @@ def _collect_table_files( return table_files + def _collect_page_citation_files( + self, + chunks: list[dict[str, Any]], + page_citation_assets_dir: str, + ) -> list[ZipResourceFileInfo]: + page_citation_refs = _collect_page_citation_refs(chunks) + if not page_citation_refs: + return [] + if not os.path.exists(page_citation_assets_dir): + raise StorageServiceException( + internal_message=( + "Page citation asset directory not found for ZIP packaging: " + f"page_citation_assets_dir={page_citation_assets_dir}" + ), + operation="collect_page_citation_files", + ) + + page_citation_files: list[ZipResourceFileInfo] = [] + for artifact_ref in page_citation_refs: + source_path = os.path.join( + page_citation_assets_dir, + os.path.basename(artifact_ref), + ) + if not os.path.isfile(source_path): + raise StorageServiceException( + internal_message=( + "Cannot resolve page citation asset for ZIP packaging: " + f"artifact_ref={artifact_ref}, source_path={source_path}" + ), + operation="collect_page_citation_files", + ) + page_citation_files.append( + { + "id": artifact_ref, + "file_path": artifact_ref, + "original_name": os.path.basename(artifact_ref), + "size_bytes": os.path.getsize(source_path), + "format": os.path.splitext(artifact_ref)[1].lstrip("."), + "source_path": source_path, + "zip_path": artifact_ref, + } + ) + + return page_citation_files + def _collect_files_by_name(directory_path: str) -> dict[str, str]: files: dict[str, str] = {} for filename in os.listdir(directory_path): @@ -222,6 +275,45 @@ def _collect_files_by_name(directory_path: str) -> dict[str, str]: return files +def _collect_page_citation_refs(chunks: list[dict[str, Any]]) -> list[str]: + refs: list[str] = [] + seen_refs: set[str] = set() + for chunk in chunks: + chunk_type = str(chunk.get("type") or "").strip().lower() + if chunk_type != "page": + continue + metadata = chunk.get("metadata") + if not isinstance(metadata, dict): + continue + page_assets = metadata.get("page_assets") + if not isinstance(page_assets, list): + continue + for page_asset in page_assets: + if not isinstance(page_asset, dict): + continue + artifact_ref = _normalize_page_citation_ref( + page_asset.get("artifact_ref"), + ) + if artifact_ref and artifact_ref not in seen_refs: + seen_refs.add(artifact_ref) + refs.append(artifact_ref) + return refs + + +def _normalize_page_citation_ref(value: Any) -> str | None: + if value is None: + return None + normalized = str(value).strip().replace("\\", "/").lstrip("/") + parts = [ + part + for part in normalized.split("/") + if part and part not in {".", ".."} + ] + if len(parts) != 2 or parts[0] != "page_citation_assets": + return None + return "/".join(parts) + + def _add_candidate(candidates: list[str], value: str | None) -> None: if not value: return diff --git a/packages/shared-python/shared/services/storage/zip_result_service.py b/packages/shared-python/shared/services/storage/zip_result_service.py index b2fee6e17..fe1f0c6dc 100644 --- a/packages/shared-python/shared/services/storage/zip_result_service.py +++ b/packages/shared-python/shared/services/storage/zip_result_service.py @@ -87,6 +87,7 @@ def generate_zip_package( formatted_chunks=formatted_chunks, image_files=resources.image_files, table_files=resources.table_files, + page_citation_files=resources.page_citation_files, doc_nav=doc_nav, manifest=manifest, temp_dir=temp_dir,