From 21e09514b7549d062ee6764c6286e34a1cd38d1c Mon Sep 17 00:00:00 2001 From: chengke <404835780@qq.com> Date: Sat, 12 Sep 2026 00:46:59 +0800 Subject: [PATCH] fix(retrieval): retain connected body chunks for asset filters --- AGENTS.md | 4 +- .../retrieval/hydration/result_assembly.py | 37 +++++++++--- .../shared/tests/test_asset_inline.py | 60 +++++++++++++++++++ 3 files changed, 91 insertions(+), 10 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index fa13a28f..4dc21c38 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -589,8 +589,8 @@ are resolved to chunks, then assembled like classic. Map-nav is archived at `hydration.result_assembly.assemble_retrieval_results()`: 1. Filters by `exclude_document_ids` and `exclude_sections` -2. Filters by `allowed_chunk_types` (data_type parameter) -3. Hydrates `connect_to` targets (related table chunks inlined into text) +2. Hydrates `connect_to` targets +3. Filters by `allowed_chunk_types` while retaining body chunks connected to a requested image/table 4. Cleans asset path references from content 5. Public projection builds `source`: `{document_id, source_file_name, section_path}` plus `page_nums` for `chunk_type=page` when present 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 610bf99f..96a38fd6 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py +++ b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py @@ -30,20 +30,15 @@ async def assemble_retrieval_results( allowed_chunk_types: set[str] | None = None, revision_pins: Mapping[str, str] | None = None, ) -> list[dict[str, Any]]: - filtered_rows = filter_excluded_rows( + scoped_rows = filter_excluded_rows( rows, exclude_document_ids=exclude_document_ids, exclude_sections=exclude_sections, document_scope=document_scope, ) - if allowed_chunk_types is not None: - filtered_rows = [ - row for row in filtered_rows - if normalize_chunk_type(row.get('chunk_type')) in allowed_chunk_types - ] hydrated_rows = await hydrate_connected_target_rows( db=db, - rows=filtered_rows, + rows=scoped_rows, exclude_document_ids=exclude_document_ids, exclude_sections=exclude_sections, document_scope=document_scope, @@ -51,9 +46,14 @@ async def assemble_retrieval_results( ) rows_by_chunk_id = { str(row.get('chunk_id') or ''): row - for row in [*filtered_rows, *hydrated_rows] + for row in [*scoped_rows, *hydrated_rows] if row.get('chunk_id') } + filtered_rows = _filter_rows_by_allowed_chunk_types( + scoped_rows, + allowed_chunk_types=allowed_chunk_types, + rows_by_chunk_id=rows_by_chunk_id, + ) embedded_targets: set[str] = set() for row in filtered_rows: @@ -88,6 +88,27 @@ async def assemble_retrieval_results( return assembled +def _filter_rows_by_allowed_chunk_types( + rows: list[dict[str, Any]], + *, + allowed_chunk_types: set[str] | None, + rows_by_chunk_id: Mapping[str, dict[str, Any]], +) -> list[dict[str, Any]]: + if allowed_chunk_types is None: + return rows + + return [ + row + for row in rows + if normalize_chunk_type(row.get('chunk_type')) in allowed_chunk_types + or any( + normalize_chunk_type(rows_by_chunk_id.get(target_id, {}).get('chunk_type')) + in allowed_chunk_types + for target_id in iter_connected_target_ids(row) + ) + ] + + def _page_summary(row: dict[str, Any]) -> str: metadata = row.get('chunk_metadata') or row.get('metadata') or {} if not isinstance(metadata, dict): diff --git a/packages/shared-python/shared/tests/test_asset_inline.py b/packages/shared-python/shared/tests/test_asset_inline.py index d8345c24..75b09c04 100644 --- a/packages/shared-python/shared/tests/test_asset_inline.py +++ b/packages/shared-python/shared/tests/test_asset_inline.py @@ -101,6 +101,66 @@ async def test_assemble_inserts_table_at_placeholder() -> None: assert "SHOULD NOT LEAK" not in content +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("asset_type", "file_path", "placeholder", "display_marker"), + [ + ("image", "images/a.png", "[images/a.png]", "[Image: images/a.png]"), + ("table", "tables/a.html", "[tables/a.html]", "[Table: tables/a.html]"), + ], +) +async def test_asset_type_filter_keeps_body_that_connects_to_requested_asset( + monkeypatch, + asset_type: str, + file_path: str, + placeholder: str, + display_marker: str, +) -> None: + body_row = { + "chunk_id": "text-1", + "chunk_type": "text", + "content": f"查看 {placeholder}", + "chunk_metadata": { + "connect_to": [ + { + "target": "asset-1", + "relation": "embeds", + "ref": placeholder, + } + ] + }, + } + asset_row = { + "chunk_id": "asset-1", + "chunk_type": asset_type, + "content": "资产说明" if asset_type == "image" else "
", + "file_path": file_path, + "chunk_metadata": {"summary": "资产说明"}, + } + + async def hydrate_connected_rows(**_kwargs: object) -> list[dict[str, object]]: + return [asset_row] + + monkeypatch.setattr( + "shared.services.retrieval.hydration.result_assembly.hydrate_connected_target_rows", + hydrate_connected_rows, + ) + + assembled = await assemble_retrieval_results( + rows=[ + body_row, + {"chunk_id": "text-2", "chunk_type": "text", "content": "无图"}, + ], + exclude_document_ids=[], + exclude_sections=[], + allowed_chunk_types={asset_type}, + ) + + assert [row["chunk_id"] for row in assembled] == ["text-1"] + assert display_marker in assembled[0]["content"] + assert "资产说明" in assembled[0]["content"] + + def test_node_unit_span_inlines_section_assets() -> None: provider = KnowhereProvider( doc_id="doc-1",