feat: include all query-hit snippets for page chunks in retrieval responses

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[

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 = [

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"] = ""

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"(?<![a-z0-9]){re.escape(token.lower())}(?![a-z0-9])")
+        for token in query_tokens
+        if token
+    ]
+    if not patterns:
+        return []
+
+    seen: set[str] = set()
+    snippets: list[str] = []
+    for match in _iter_occurrences(lower_content, patterns):
+        snippet = _build_snippet(
+            content,
+            match.start(),
+            match.end(),
+            context_chars=context_chars,
+        )
+        if snippet in seen:
+            continue
+        seen.add(snippet)
+        snippets.append(snippet)
+        if len(snippets) >= 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

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,15 +53,23 @@ 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
         assembled_row = dict(row)
         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]],