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("…")