From 14695aea4af7b4dd6772b38cf78b784dc284c837 Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Thu, 20 Aug 2026 22:46:19 +0800 Subject: [PATCH] feat: expose page_nums on retrieval source and clean dead citation path Make page-memory node page lists a declared source field for results and referenced_chunks, drop unused attach_citation / source-passthrough, and fix related test isolation/mocks so contract suites stay green. Co-authored-by: Cursor --- AGENTS.md | 2 +- .../test_retrieval_mapnav_session_contract.py | 2 +- .../worker/tests/unit/test_summary_builder.py | 79 ++++++++++--------- .../services/retrieval/execution/plan.py | 5 +- .../retrieval/execution/reference_resolver.py | 12 ++- .../execution/response_projection.py | 17 +--- .../services/retrieval/execution/routes.py | 9 +-- .../retrieval/hydration/result_assembly.py | 4 + .../services/retrieval/hydration/row_utils.py | 10 ++- .../tests/test_page_memory_vlm_limiter.py | 4 +- 10 files changed, 73 insertions(+), 71 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index c9afaa80c..f3f2a91d3 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -586,7 +586,7 @@ Default agentic path is checklist map-nav (`nav/`): PLANNER (`plan_query`) → H 2. Filters by `allowed_chunk_types` (data_type parameter) 3. Hydrates `connect_to` targets (related table chunks inlined into text) 4. Cleans asset path references from content -5. Attaches citation: `{document_id, chunk_id, source_file_name, section_path}` +5. Public projection builds `source`: `{document_id, source_file_name, section_path}` plus `page_nums` for `chunk_type=page` when present ### Small Corpus Optimization diff --git a/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py b/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py index ca271f66f..a277c4c03 100644 --- a/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py +++ b/apps/api/tests/contract/test_retrieval_mapnav_session_contract.py @@ -220,5 +220,5 @@ async def fake_assemble_retrieval_results( "fresh_db_close", ] assert outcome.response["router_used"] == "mapnav" - assert outcome.response["results"][0]["citation"]["document_id"] == "doc_contract" + assert outcome.response["results"][0]["document_id"] == "doc_contract" assert outcome.completion_label == "MAPNAV RETRIEVAL" diff --git a/apps/worker/tests/unit/test_summary_builder.py b/apps/worker/tests/unit/test_summary_builder.py index a0eaf2c73..a7f0d3084 100644 --- a/apps/worker/tests/unit/test_summary_builder.py +++ b/apps/worker/tests/unit/test_summary_builder.py @@ -2,17 +2,20 @@ from __future__ import annotations +import importlib +from types import ModuleType from typing import Any, Dict, List import pytest -from app.services.connect_builder.summary_builder import ( - SUMMARY_MAX_LEN, - _deterministic_section_summary, - _llm_summarize, - _recursive_summarize_nav, - build_self_only_lookup, -) + +def _summary_builder() -> ModuleType: + """Resolve the live module. + + Worker contract tests may evict/reimport ``app.*``. Calling functions bound at + collection time would miss later monkeypatches on the new module object. + """ + return importlib.import_module("app.services.connect_builder.summary_builder") def _leaf(title: str, summary: str = "", path: str = "") -> Dict[str, Any]: @@ -41,7 +44,8 @@ def _parent( class TestDeterministicAssembly: def test_order_covers_self_only_then_titles(self) -> None: - text = _deterministic_section_summary( + sb = _summary_builder() + text = sb._deterministic_section_summary( is_top_level=False, self_only="intro paragraph here", child_titles=["Alpha", "Beta"], @@ -53,6 +57,7 @@ def test_order_covers_self_only_then_titles(self) -> None: assert text.index("intro paragraph here") < text.index("Alpha, Beta") def test_all_child_titles_even_when_summary_empty(self) -> None: + sb = _summary_builder() parent = _parent( "Parent", [ @@ -61,7 +66,7 @@ def test_all_child_titles_even_when_summary_empty(self) -> None: ], path="doc.pdf/Parent", ) - result = _recursive_summarize_nav( + result = sb._recursive_summarize_nav( parent, use_llm=False, source_file_name="doc.pdf", @@ -73,6 +78,7 @@ def test_all_child_titles_even_when_summary_empty(self) -> None: class TestSelfOnlyLookup: def test_exact_path_only_excludes_descendants(self) -> None: + sb = _summary_builder() chunks = [ { "path": "doc.pdf/2.4.4 隐患治理", @@ -83,12 +89,13 @@ def test_exact_path_only_excludes_descendants(self) -> None: "content": "CHILD_BODY_SHOULD_NOT_APPEAR", }, ] - lookup = build_self_only_lookup(chunks, source_file_name="doc.pdf") + lookup = sb.build_self_only_lookup(chunks, source_file_name="doc.pdf") assert lookup["2.4.4 隐患治理"] == "PARENT_INTRO_ONLY" assert "CHILD_BODY_SHOULD_NOT_APPEAR" not in lookup["2.4.4 隐患治理"] assert "2.4.4 隐患治理 / 清单项A" in lookup def test_nonleaf_includes_self_only_in_deterministic(self) -> None: + sb = _summary_builder() parent = _parent( "2.4.4 隐患治理", [ @@ -97,11 +104,11 @@ def test_nonleaf_includes_self_only_in_deterministic(self) -> None: ], path="doc.pdf/2.4.4 隐患治理", ) - lookup = build_self_only_lookup( + lookup = sb.build_self_only_lookup( [{"path": "doc.pdf/2.4.4 隐患治理", "content": "方案包括以下内容："}], source_file_name="doc.pdf", ) - result = _recursive_summarize_nav( + result = sb._recursive_summarize_nav( parent, use_llm=False, self_only_lookup=lookup, @@ -115,22 +122,22 @@ def test_nonleaf_includes_self_only_in_deterministic(self) -> None: class TestLlmTrigger: def test_short_contrib_skips_llm(self, monkeypatch: pytest.MonkeyPatch) -> None: + sb = _summary_builder() called = {"n": 0} def _boom(**kwargs: Any) -> str: called["n"] += 1 return "SHOULD_NOT_USE" - monkeypatch.setattr( - "app.services.connect_builder.summary_builder._llm_summarize", - _boom, - ) + monkeypatch.setattr(sb, "_llm_summarize", _boom) parent = _parent( "P", [_leaf("A", summary="x"), _leaf("B", summary="y")], path="doc.pdf/P", ) - result = _recursive_summarize_nav(parent, use_llm=True, source_file_name="doc.pdf") + result = sb._recursive_summarize_nav( + parent, use_llm=True, source_file_name="doc.pdf" + ) assert called["n"] == 0 assert result.startswith("This section covers: ") assert "A" in result and "B" in result @@ -138,17 +145,15 @@ def _boom(**kwargs: Any) -> str: def test_long_contrib_calls_llm_with_title_for_empty_summary( self, monkeypatch: pytest.MonkeyPatch ) -> None: + sb = _summary_builder() captured: Dict[str, Any] = {} def _fake_llm(**kwargs: Any) -> str: captured.update(kwargs) return "LLM_SUMMARY" - monkeypatch.setattr( - "app.services.connect_builder.summary_builder._llm_summarize", - _fake_llm, - ) - long_a = "A" * (SUMMARY_MAX_LEN + 5) + monkeypatch.setattr(sb, "_llm_summarize", _fake_llm) + long_a = "A" * (sb.SUMMARY_MAX_LEN + 5) parent = _parent( "P", [ @@ -159,7 +164,7 @@ def _fake_llm(**kwargs: Any) -> str: ) lookup = {"P": "SELF_ONLY_INTRO"} # section path from doc.pdf/P is "P" - result = _recursive_summarize_nav( + result = sb._recursive_summarize_nav( parent, use_llm=True, self_only_lookup=lookup, @@ -176,17 +181,19 @@ def _fake_llm(**kwargs: Any) -> str: def test_single_child_with_self_only_does_not_copy_child( self, monkeypatch: pytest.MonkeyPatch ) -> None: + sb = _summary_builder() monkeypatch.setattr( - "app.services.connect_builder.summary_builder._llm_summarize", + sb, + "_llm_summarize", lambda **kwargs: "MERGED", ) - long_child = "C" * (SUMMARY_MAX_LEN + 1) + long_child = "C" * (sb.SUMMARY_MAX_LEN + 1) parent = _parent( "P", [_leaf("OnlyChild", summary=long_child)], path="doc.pdf/P", ) - result = _recursive_summarize_nav( + result = sb._recursive_summarize_nav( parent, use_llm=True, self_only_lookup={"P": "intro"}, @@ -200,6 +207,7 @@ class TestPromptPayload: def test_file_summary_prompt_contains_scope_blocks( self, monkeypatch: pytest.MonkeyPatch ) -> None: + sb = _summary_builder() captured: Dict[str, Any] = {} def _fake_client(**_kwargs: Any) -> Any: @@ -215,7 +223,7 @@ def chat_completion(self, **kwargs: Any) -> str: _fake_client, ) # Ensure build_prompt path works - out = _llm_summarize( + out = sb._llm_summarize( node_name="Parent", self_only="intro text", child_rows=[("ChildA", "summary A"), ("ChildB", "ChildB")], @@ -239,11 +247,7 @@ def test_enrich_persists_top_summary_and_defaults_top_llm( ) -> None: import json - from app.services.connect_builder.summary_builder import ( - enrich_doc_nav_summaries, - load_nav_top_summary, - ) - + sb = _summary_builder() captured: Dict[str, Any] = {} def _fake_llm(**kwargs: Any) -> str: @@ -252,14 +256,11 @@ def _fake_llm(**kwargs: Any) -> str: captured["calls"] = int(captured.get("calls") or 0) + 1 return "LLM document overview" - monkeypatch.setattr( - "app.services.connect_builder.summary_builder._llm_summarize", - _fake_llm, - ) + monkeypatch.setattr(sb, "_llm_summarize", _fake_llm) file_dir = tmp_path / "report.pdf" file_dir.mkdir() - long_leaf = "L" * (SUMMARY_MAX_LEN + 5) + long_leaf = "L" * (sb.SUMMARY_MAX_LEN + 5) doc_nav = { "version": "1.0", "file_name": "report.pdf", @@ -287,7 +288,7 @@ def _fake_llm(**kwargs: Any) -> str: encoding="utf-8", ) - results = enrich_doc_nav_summaries( + results = sb.enrich_doc_nav_summaries( str(tmp_path), source_file="report.pdf", use_llm=False, @@ -301,6 +302,6 @@ def _fake_llm(**kwargs: Any) -> str: assert saved["top_summary"] == "LLM document overview" # Section leaves keep original summaries; top LLM must not rewrite them. assert saved["sections"][0]["summary"] == long_leaf - assert load_nav_top_summary(str(file_dir), "report.pdf") == ( + assert sb.load_nav_top_summary(str(file_dir), "report.pdf") == ( "LLM document overview" ) diff --git a/packages/shared-python/shared/services/retrieval/execution/plan.py b/packages/shared-python/shared/services/retrieval/execution/plan.py index ed9f9f591..8d05753ab 100644 --- a/packages/shared-python/shared/services/retrieval/execution/plan.py +++ b/packages/shared-python/shared/services/retrieval/execution/plan.py @@ -286,12 +286,11 @@ def _log_retrieval_complete( results = outcome.get("results", []) if isinstance(results, list): for index, result in enumerate(results[:10]): - source = result.get("source", {}) logger.info( f" [{index + 1}] type={result.get('chunk_type', '?')} " f"score={result.get('score') or 0.0:.4f}" - f" path={source.get('section_path', '')}" - f" file={source.get('source_file_name', '')}" + f" path={result.get('section_path') or ''}" + f" file={result.get('source_file_name') or ''}" ) if len(results) > 10: logger.info(f" ... and {len(results) - 10} more") 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 d5e2c3586..802fa2786 100644 --- a/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py +++ b/packages/shared-python/shared/services/retrieval/execution/reference_resolver.py @@ -9,7 +9,10 @@ from shared.services.retrieval.execution.response_projection import ( enrich_referenced_chunks_with_asset_url, ) -from shared.services.retrieval.hydration.row_utils import build_reference_lookup_key +from shared.services.retrieval.hydration.row_utils import ( + build_reference_lookup_key, + extract_page_nums, +) @dataclass(frozen=True) @@ -36,7 +39,7 @@ async def resolve_workflow_references( resolved = _select_matching_references(refs, hydrated_rows) enriched_rows = await enrich_referenced_chunks_with_asset_url(resolved.rows) return ResolvedWorkflowReferences( - refs=_merge_reference_asset_url(resolved.refs, enriched_rows), + refs=_merge_reference_projection(resolved.refs, enriched_rows), rows=resolved.rows, ) @@ -95,7 +98,7 @@ def _row_key(row: dict[str, Any]) -> tuple[str, str, str, str]: ) -def _merge_reference_asset_url( +def _merge_reference_projection( refs: list[dict[str, Any]], rows: list[dict[str, Any]], ) -> list[dict[str, Any]]: @@ -116,6 +119,9 @@ def _merge_reference_asset_url( if row is not None: if row.get("asset_url"): merged["asset_url"] = row["asset_url"] + page_nums = extract_page_nums(row) + if page_nums is not None: + merged["page_nums"] = page_nums merged_refs.append(merged) return merged_refs diff --git a/packages/shared-python/shared/services/retrieval/execution/response_projection.py b/packages/shared-python/shared/services/retrieval/execution/response_projection.py index 313eb8ece..656533c65 100644 --- a/packages/shared-python/shared/services/retrieval/execution/response_projection.py +++ b/packages/shared-python/shared/services/retrieval/execution/response_projection.py @@ -9,18 +9,8 @@ ) -def attach_citation(row: dict[str, Any]) -> dict[str, Any]: - citation = { - 'document_id': row.get('document_id'), - 'chunk_id': row.get('chunk_id'), - 'source_file_name': row.get('source_file_name'), - 'section_path': row.get('section_path'), - } - return {**row, 'citation': citation} - - def to_public_source(row: dict[str, Any]) -> dict[str, Any]: - return {field: row.get(field) for field in PUBLIC_SOURCE_FIELDS} + return {field: row[field] for field in PUBLIC_SOURCE_FIELDS if field in row} async def enrich_referenced_chunks_with_asset_url(refs: list[dict[str, Any]]) -> list[dict[str, Any]]: @@ -63,10 +53,7 @@ async def project_public_retrieval_response(response: dict[str, Any]) -> dict[st metadata = row.get('chunk_metadata') if isinstance(metadata, dict): public_row['metadata'] = metadata - if 'source' in row: - public_row['source'] = row['source'] - else: - public_row['source'] = to_public_source(row) + public_row['source'] = to_public_source(row) public_results.append(public_row) public_response['results'] = public_results diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index df521e4d6..e67a1ead8 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -9,9 +9,6 @@ from shared.services.retrieval.search.discovery import bottom_discovery from shared.services.retrieval.execution.reference_resolver import resolve_workflow_references from shared.services.retrieval.hydration.result_assembly import assemble_retrieval_results -from shared.services.retrieval.execution.response_projection import ( - attach_citation, -) from shared.services.retrieval.hydration.legacy_evidence import render_legacy_evidence_text from shared.services.retrieval.execution.route_types import ( RetrievalRouteContext, @@ -84,7 +81,7 @@ async def _try_run_small_corpus_route( exclude_sections=context.exclude_sections, allowed_chunk_types=context.allowed_chunk_types, ) - results = [attach_citation(row) for row in assembled_rows] + results = assembled_rows response = { "namespace": context.namespace, "query": context.query, @@ -143,7 +140,7 @@ async def _run_classic_topk_route( exclude_sections=context.exclude_sections, allowed_chunk_types=context.allowed_chunk_types, ) - results = [attach_citation(row) for row in assembled_rows] + results = assembled_rows response = { "namespace": context.namespace, "query": context.query, @@ -269,7 +266,7 @@ async def _run_mapnav_route( "evidence_text": evidence_text, "answer_text": "", "referenced_chunks": resolved.refs, - "results": [attach_citation(row) for row in assembled_rows], + "results": assembled_rows, "stop_reason": stop_reason, "decision_trace": decision_trace, } diff --git a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py index 96534bce3..ed4487046 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py +++ b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py @@ -7,6 +7,7 @@ from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows from shared.services.retrieval.hydration.row_utils import ( clean_content, + extract_page_nums, filter_excluded_rows, iter_connected_target_ids, normalize_chunk_type, @@ -59,6 +60,9 @@ async def assemble_retrieval_results( if chunk_type == 'page': assembled_row['content'] = _page_summary(row) assembled_row['content_source'] = 'summary' + page_nums = extract_page_nums(row) + if page_nums is not None: + assembled_row['page_nums'] = page_nums elif chunk_type == 'table': assembled_row['content'] = _compose_table_content(row, rows_by_chunk_id) assembled_row['content_source'] = 'summary' diff --git a/packages/shared-python/shared/services/retrieval/hydration/row_utils.py b/packages/shared-python/shared/services/retrieval/hydration/row_utils.py index fe8ab4f0e..32bcb2ba2 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/row_utils.py +++ b/packages/shared-python/shared/services/retrieval/hydration/row_utils.py @@ -17,7 +17,7 @@ 'file_path', } PUBLIC_SOURCE_FIELDS = { - 'document_id', 'source_file_name', 'section_path', + 'document_id', 'source_file_name', 'section_path', 'page_nums', } ReferenceLookupKey = tuple[str, str, str, str] @@ -40,6 +40,14 @@ def is_media_chunk(row: dict[str, Any]) -> bool: return normalize_chunk_type(row.get('chunk_type')) in MEDIA_CHUNK_TYPES +def extract_page_nums(row: dict[str, Any]) -> list[int] | None: + metadata = row.get('chunk_metadata') or row.get('metadata') or {} + if not isinstance(metadata, dict): + return None + page_nums = metadata.get('page_nums') + return page_nums if isinstance(page_nums, list) else None + + def build_reference_lookup_key( *, document_id: object, diff --git a/packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py b/packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py index c69854395..619133b9f 100644 --- a/packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py +++ b/packages/shared-python/shared/tests/test_page_memory_vlm_limiter.py @@ -179,7 +179,7 @@ def chat_completion_with_usage(self, **kwargs: Any) -> Any: monkeypatch.setattr( summary_engine._client_mod, "get_openai_client", - lambda model=None: _FakeClient(), + lambda model=None, api_key=None, api_url=None, **_kwargs: _FakeClient(), ) with pytest.raises(UnavailableException): @@ -208,7 +208,7 @@ def chat_completion_with_usage(self, **kwargs: Any) -> Any: monkeypatch.setattr( summary_engine._client_mod, "get_openai_client", - lambda model=None: _FakeClient(), + lambda model=None, api_key=None, api_url=None, **_kwargs: _FakeClient(), ) with pytest.raises(UnavailableException):