From cf0d425077a763b81698967934f573eb3a91c737 Mon Sep 17 00:00:00 2001 From: Gordon YUEN Date: Sun, 9 Aug 2026 01:30:25 +0800 Subject: [PATCH] feat: include all query-hit snippets for page chunks in retrieval responses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Page chunks can be very large (whole scanned directory pages) and their LLM summary rarely contains the exact queried term, so retrieval results were unhelpful for name lookups. Extract every occurrence of the query terms from the full page content as bounded snippets (±100 chars, capped at 20, deduplicated) and append them after the summary. - New shared/services/retrieval/hydration/page_snippets.py - assemble_retrieval_results gains a query param; page chunks set content_source='content_snippets' when hits exist, else 'summary' - Thread query=context.query through all assembly call sites in routes.py - Unit tests (6) + worker contract tests (2) + API contract test (1) --- .../tests/contract/test_retrieval_contract.py | 48 +++++++ .../test_page_memory_retrieval_contract.py | 59 +++++++++ .../services/retrieval/execution/routes.py | 3 + .../retrieval/hydration/page_snippets.py | 122 ++++++++++++++++++ .../retrieval/hydration/result_assembly.py | 36 +++++- .../shared/tests/test_page_snippets.py | 78 +++++++++++ 6 files changed, 344 insertions(+), 2 deletions(-) create mode 100644 packages/shared-python/shared/services/retrieval/hydration/page_snippets.py create mode 100644 packages/shared-python/shared/tests/test_page_snippets.py diff --git a/apps/api/tests/contract/test_retrieval_contract.py b/apps/api/tests/contract/test_retrieval_contract.py index 58de969c7..f42c70caf 100644 --- a/apps/api/tests/contract/test_retrieval_contract.py +++ b/apps/api/tests/contract/test_retrieval_contract.py @@ -301,6 +301,54 @@ async def test_should_return_seeded_retrieval_results_for_the_authenticated_user } +@pytest.mark.asyncio +async def test_page_chunk_result_includes_all_query_snippets( + developer_api_client_factory: Callable[ + [], AbstractAsyncContextManager[AsyncClient] + ], +) -> None: + page_content = ( + "Name: Mr. HUI Kim\nPost: Deputy Commissioner\n" + + "X" * 300 + + "\nCHEUNG Hon-lam Gordon\n2835 2147\n" + + "Y" * 300 + + "\nYUEN Chun-cheung Gordon\n2835 2154\n" + ) + async with developer_api_client_factory() as api_client: + await _seed_retrieval_document( + user_id="local-dev-user", + namespace="contract-retrieval", + source_file_name="contract-directory.pdf", + section_path="directory/root", + content=page_content, + chunk_type="page", + chunk_metadata={"summary": "Directory contact summary"}, + ) + + response = await api_client.post( + "/api/v1/retrieval/query", + json={ + "namespace": "contract-retrieval", + "query": "Gordon", + "top_k": 10, + }, + ) + + assert response.status_code == 200 + + response_json = cast(dict[str, object], response.json()) + results = cast(list[dict[str, object]], response_json["results"]) + + assert len(results) == 1 + assert results[0]["chunk_type"] == "page" + assert results[0]["content_source"] == "content_snippets" + content = str(results[0]["content"]) + assert content.startswith("Directory contact summary") + assert "CHEUNG Hon-lam Gordon" in content + assert "YUEN Chun-cheung Gordon" in content + assert "contract-directory.pdf" in str(response_json["evidence_text"]) + + @pytest.mark.asyncio async def test_should_default_the_namespace_to_default_when_it_is_omitted( developer_api_client_factory: Callable[ 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 922aff4c5..bf8d43883 100644 --- a/apps/worker/tests/contract/test_page_memory_retrieval_contract.py +++ b/apps/worker/tests/contract/test_page_memory_retrieval_contract.py @@ -101,6 +101,65 @@ async def test_page_result_assembly_uses_summary_not_raw_content() -> None: assert assembled[0]["content"] == "制度标准总则摘要" +@pytest.mark.asyncio +async def test_page_result_assembly_appends_query_snippets_when_query_matches() -> None: + rows = [ + { + "chunk_id": "page-node-1", + "chunk_type": "page", + "content": ( + "Name: Mr. HUI Kim\nPost: Deputy Commissioner\n" + + "X" * 300 + + "\nCHEUNG Hon-lam Gordon\n2835 2147\nALO II(YE3)6\n" + + "Y" * 300 + + "\nYUEN Chun-cheung Gordon\n2835 2154\n" + ), + "chunk_metadata": { + "summary": "联络资料摘要", + "page_nums": [1], + }, + } + ] + + assembled = await assemble_retrieval_results( + rows=rows, + exclude_document_ids=[], + exclude_sections=[], + query="Gordon", + ) + + assert assembled[0]["content_source"] == "content_snippets" + content = assembled[0]["content"] + assert content.startswith("联络资料摘要") + assert "CHEUNG Hon-lam Gordon" in content + assert "YUEN Chun-cheung Gordon" in content + + +@pytest.mark.asyncio +async def test_page_result_assembly_falls_back_to_summary_without_query_hits() -> None: + rows = [ + { + "chunk_id": "page-node-1", + "chunk_type": "page", + "content": "CHEUNG Hon-lam Gordon\n2835 2147\n", + "chunk_metadata": { + "summary": "联络资料摘要", + "page_nums": [1], + }, + } + ] + + assembled = await assemble_retrieval_results( + rows=rows, + exclude_document_ids=[], + exclude_sections=[], + query="NoSuchName", + ) + + assert assembled[0]["content_source"] == "summary" + assert assembled[0]["content"] == "联络资料摘要" + + @pytest.mark.asyncio async def test_table_result_assembly_uses_summary_not_html() -> None: rows = [ diff --git a/packages/shared-python/shared/services/retrieval/execution/routes.py b/packages/shared-python/shared/services/retrieval/execution/routes.py index 08b33a2a4..25688b730 100644 --- a/packages/shared-python/shared/services/retrieval/execution/routes.py +++ b/packages/shared-python/shared/services/retrieval/execution/routes.py @@ -77,6 +77,7 @@ async def _try_run_small_corpus_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, allowed_chunk_types=context.allowed_chunk_types, + query=context.query, ) results = [attach_citation(row) for row in assembled_rows] response = { @@ -136,6 +137,7 @@ async def _run_classic_topk_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, allowed_chunk_types=context.allowed_chunk_types, + query=context.query, ) results = [attach_citation(row) for row in assembled_rows] response = { @@ -222,6 +224,7 @@ async def _run_agentic_route( exclude_document_ids=context.exclude_document_ids, exclude_sections=context.exclude_sections, allowed_chunk_types=context.allowed_chunk_types, + query=context.query, ) response = workflow_result.to_api_response() response["answer_text"] = "" diff --git a/packages/shared-python/shared/services/retrieval/hydration/page_snippets.py b/packages/shared-python/shared/services/retrieval/hydration/page_snippets.py new file mode 100644 index 000000000..34a7607bd --- /dev/null +++ b/packages/shared-python/shared/services/retrieval/hydration/page_snippets.py @@ -0,0 +1,122 @@ +""" +Query-hit snippet extraction for page chunks in retrieval responses. + +Page chunks can be very large (a whole scanned page of a directory, form, or +manual). Retrieval responses surface only the chunk's LLM ``summary``, which +rarely contains the exact queried term. These helpers extract every +occurrence of the query terms from the full page content so the response can +show the actual matching lines. +""" + +from __future__ import annotations + +import re + +_PAGE_SNIPPET_CONTEXT_CHARS = 100 +_PAGE_SNIPPET_MAX = 20 +_ELLIPSIS = "…" + + +def extract_page_snippets( + content: str, + query_tokens: list[str], + *, + context_chars: int = _PAGE_SNIPPET_CONTEXT_CHARS, + max_snippets: int = _PAGE_SNIPPET_MAX, +) -> list[str]: + """Extract every query-term occurrence from ``content`` as a snippet. + + Each snippet centers one occurrence of any query token with + ``context_chars`` characters of surrounding context on each side, trimmed + at word boundaries and marked with an ellipsis at truncated edges. + Occurrences are returned in document order, deduplicated, and capped at + ``max_snippets`` so pathological pages (hundreds of hits for a common + name) stay bounded. + """ + if not content or not query_tokens: + return [] + + lower_content = content.lower() + patterns = [ + re.compile(rf"(?= max_snippets: + break + return snippets + + +def _iter_occurrences(lower_content: str, patterns: list[re.Pattern[str]]): + """Yield non-overlapping occurrence matches across all patterns in order.""" + matches: list[re.Match[str]] = [] + for pattern in patterns: + matches.extend(pattern.finditer(lower_content)) + matches.sort(key=lambda m: (m.start(), m.end())) + return _dedupe_overlapping(matches) + + +def _dedupe_overlapping(matches: list[re.Match[str]]): + last_end = -1 + for match in matches: + if match.start() < last_end: + continue + last_end = match.end() + yield match + + +def _build_snippet( + content: str, + start: int, + end: int, + *, + context_chars: int, +) -> str: + snippet_start = max(0, start - context_chars) + snippet_end = min(len(content), end + context_chars) + + if snippet_start > 0: + snippet_start = _next_word_boundary(content, snippet_start, direction=-1) + if snippet_end < len(content): + snippet_end = _next_word_boundary(content, snippet_end, direction=1) + + prefix = _ELLIPSIS if snippet_start > 0 else "" + suffix = _ELLIPSIS if snippet_end < len(content) else "" + return f"{prefix}{content[snippet_start:snippet_end]}{suffix}" + + +def _next_word_boundary(text: str, index: int, *, direction: int) -> int: + """Move ``index`` to a nearby word boundary in the given direction. + + Only walks a few characters (``_BOUNDARY_WALK_LIMIT``) so long unbroken + runs (e.g. filler lines) do not stretch the snippet across the whole + page. Returns the original index when no boundary is nearby. + """ + length = len(text) + cursor = index + for _ in range(_BOUNDARY_WALK_LIMIT): + if not 0 < cursor < length: + break + if not text[cursor - 1].isalnum() and not text[cursor].isalnum(): + return cursor + cursor += direction + return index + + +_BOUNDARY_WALK_LIMIT = 16 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..6628b8607 100644 --- a/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py +++ b/packages/shared-python/shared/services/retrieval/hydration/result_assembly.py @@ -5,12 +5,14 @@ from sqlalchemy.ext.asyncio import AsyncSession from shared.services.retrieval.hydration.connected import hydrate_connected_target_rows +from shared.services.retrieval.hydration.page_snippets import extract_page_snippets from shared.services.retrieval.hydration.row_utils import ( clean_content, filter_excluded_rows, iter_connected_target_ids, normalize_chunk_type, ) +from shared.utils.text_utils import tokenize_for_retrieval async def assemble_retrieval_results( @@ -20,6 +22,7 @@ async def assemble_retrieval_results( exclude_document_ids: list[str], exclude_sections: list[dict[str, str]], allowed_chunk_types: set[str] | None = None, + query: str | None = None, ) -> list[dict[str, Any]]: filtered_rows = filter_excluded_rows( rows, @@ -50,6 +53,7 @@ async def assemble_retrieval_results( embedded_targets.add(target_id) assembled: list[dict[str, Any]] = [] + query_tokens = tokenize_for_retrieval(query or "", dedupe=True) for row in filtered_rows: if row.get('chunk_id') in embedded_targets: continue @@ -57,8 +61,15 @@ async def assemble_retrieval_results( base_content = str(row.get('content') or '') chunk_type = normalize_chunk_type(row.get('chunk_type')) if chunk_type == 'page': - assembled_row['content'] = _page_summary(row) - assembled_row['content_source'] = 'summary' + page_content = _page_content_with_snippets( + row, + query_tokens, + base_content=base_content, + ) + assembled_row['content'] = page_content + assembled_row['content_source'] = ( + 'content_snippets' if page_content != _page_summary(row) else 'summary' + ) elif chunk_type == 'table': assembled_row['content'] = _compose_table_content(row, rows_by_chunk_id) assembled_row['content_source'] = 'summary' @@ -86,6 +97,27 @@ def _page_summary(row: dict[str, Any]) -> str: return str(metadata.get('summary') or '').strip() +def _page_content_with_snippets( + row: dict[str, Any], + query_tokens: list[str], + *, + base_content: str, +) -> str: + """Compose page-chunk content as the summary plus query-hit snippets. + + Page chunks can be very large and their LLM summary rarely contains the + exact queried term. When the query matches the full page text, append + every occurrence snippet so the response surfaces the actual matching + lines. Falls back to the summary alone when there are no hits. + """ + summary = _page_summary(row) + snippets = extract_page_snippets(base_content, query_tokens) + if not snippets: + return summary + parts = [part for part in [summary, *snippets] if part] + return '\n\n'.join(parts) + + def _compose_table_content( row: dict[str, Any], rows_by_chunk_id: dict[str, dict[str, Any]], diff --git a/packages/shared-python/shared/tests/test_page_snippets.py b/packages/shared-python/shared/tests/test_page_snippets.py new file mode 100644 index 000000000..f7f37c7ee --- /dev/null +++ b/packages/shared-python/shared/tests/test_page_snippets.py @@ -0,0 +1,78 @@ +"""Tests for retrieval page-snippet extraction.""" + +from __future__ import annotations + +import os + +os.environ.setdefault("DATABASE_URL", "postgresql+asyncpg://test:test@localhost/test") +os.environ.setdefault("TMP_PATH", "/tmp/knowhere-test") +os.environ.setdefault("S3_BUCKET_NAME", "test-uploads") +os.environ.setdefault("S3_ACCESS_KEY_ID", "test") +os.environ.setdefault("S3_SECRET_ACCESS_KEY", "test") +os.environ.setdefault("S3_TEMP_PATH", "/tmp") + +from shared.services.retrieval.hydration.page_snippets import extract_page_snippets + + +def test_extract_page_snippets_returns_all_occurrences() -> None: + content = ( + "Name: Mr. HUI Kim\nPost: Deputy Commissioner\n" + + "X" * 300 + + "\nCHEUNG Hon-lam Gordon\n2835 2147\n" + + "Y" * 300 + + "\nYUEN Chun-cheung Gordon\n2835 2154\n" + ) + + snippets = extract_page_snippets(content, ["gordon"]) + + assert len(snippets) == 2 + assert "CHEUNG Hon-lam Gordon" in snippets[0] + assert "YUEN Chun-cheung Gordon" in snippets[1] + + +def test_extract_page_snippets_respects_max_snippets_cap() -> None: + content = "\n".join(f"Row {i}: Gordon {i}" for i in range(50)) + + snippets = extract_page_snippets(content, ["gordon"], max_snippets=3) + + assert len(snippets) == 3 + + +def test_extract_page_snippets_is_case_insensitive_and_word_boundary_aware() -> None: + content = ( + "Gordon here." + + "X" * 300 + + " Gordonstoun is not a match." + + "Y" * 300 + + " gordon again." + ) + + snippets = extract_page_snippets(content, ["gordon"]) + + assert len(snippets) == 2 + assert "Gordonstoun" not in "".join(snippets) + + +def test_extract_page_snippets_returns_empty_without_hits() -> None: + assert extract_page_snippets("Gordon only once", ["missing"]) == [] + assert extract_page_snippets("", ["gordon"]) == [] + assert extract_page_snippets("Gordon here", []) == [] + + +def test_extract_page_snippets_deduplicates_identical_snippets() -> None: + content = "Mr A Gordon" + "X" * 300 + "\nMr A Gordon" + + snippets = extract_page_snippets(content, ["gordon"]) + + assert len(snippets) == 1 + + +def test_extract_page_snippets_keeps_snippets_bounded() -> None: + content = "A" * 1000 + " Gordon " + "B" * 1000 + + snippets = extract_page_snippets(content, ["gordon"], context_chars=100) + + assert len(snippets) == 1 + assert len(snippets[0]) <= 220 + assert snippets[0].startswith("…") + assert snippets[0].endswith("…")