diff --git a/CHANGELOG.md b/CHANGELOG.md index 7e8786ad..7f328381 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ ### Fixed - vLLM embedding and vLLM/Jina reranking reuse one HTTP client across requests instead of opening one per request. +- Picture and table search hits expand within their section, never across section boundaries. ## [0.64.0] - 2026-07-08 diff --git a/docs/configuration/qa.md b/docs/configuration/qa.md index ec527a1e..ac7e61d9 100644 --- a/docs/configuration/qa.md +++ b/docs/configuration/qa.md @@ -13,7 +13,7 @@ search: - **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, and QA. Default: 10 - **max_context_chars**: Hard limit on total characters in expanded content. Default: 10000. -Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded. +Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. Picture and table matches are exempt: they return their enclosing section as-is and never cross section boundaries. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded. !!! note "Reranking behavior" When a reranker is configured, search automatically retrieves 10x the requested limit, then reranks to return the final count. This improves result quality without requiring you to adjust `limit`. diff --git a/docs/python.md b/docs/python.md index 51322ad8..38308424 100644 --- a/docs/python.md +++ b/docs/python.md @@ -298,7 +298,7 @@ for result in expanded_results: print(f"Expanded content: {result.content}") ``` -Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded. +Context expansion is automatic and section-aware. For structured documents (with section headers), expansion includes the entire section containing the match. For sections that exceed the budget or are too small (e.g., a title+authors area), expansion grows outward item-by-item from the match center, skipping noise labels (footnotes, page headers). This naturally crosses into adjacent sections until the budget is filled. Picture and table matches are exempt: they return their enclosing section as-is and never cross section boundaries. For unstructured documents, expansion grows outward item-by-item. Results without `doc_item_refs` (e.g., custom chunks passed to `import_document`) pass through unexpanded. Configuration: diff --git a/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py index ee1ea311..fbea5c85 100644 --- a/haiku_rag_slim/haiku/rag/context.py +++ b/haiku_rag_slim/haiku/rag/context.py @@ -12,7 +12,8 @@ For STRUCTURED documents (containing section_header or title labels): 5. If the section is too small (under 20% of max_context_chars), expand item-by-item crossing into adjacent sections until the budget is filled. This lets small sections (e.g., title+authors) grow into neighboring - content. + content. Picture and table matches are exempt: they return their + enclosing section as-is, never crossing section boundaries. 6. Merge overlapping ranges from multiple results in the same document. Adjacent but non-overlapping ranges stay separate to preserve section independence. @@ -36,6 +37,10 @@ from haiku.rag.store.repositories.document_item import DocumentItemRepository _NOISE_LABELS = {"footnote", "page_header", "page_footer", "document_index"} _SECTION_BOUNDARY_LABELS = {"section_header", "title"} +# Labels whose pertinent unit is the item plus its own section: expansion +# never crosses section boundaries for these matches. +_TIGHT_LABELS = {"picture", "table"} + # Sections with fewer chars than this fraction of max_context_chars are # considered too small — expansion falls through to item-by-item outward # growth, which naturally crosses into adjacent sections. @@ -222,6 +227,11 @@ def _find_expansion_range( hi_bound=sec_end, ) + # Picture/table hits stay section-bounded: their pertinent unit is the + # figure or table plus its section, never neighboring sections. + if any(items[i].label in _TIGHT_LABELS for i in matched_indices): + return (items[sec_start].position, items[sec_end].position) + # Section too small (e.g., title+authors) — expand across boundaries return _expand_outward(items, center_idx, max_chars, skip_noise=True) diff --git a/tests/test_context.py b/tests/test_context.py index 4c6a8239..0bbcef3c 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -244,6 +244,46 @@ class TestFindExpansionRange: # Should expand outward into the first section assert hi >= 2 + def test_picture_in_small_section_stays_section_bounded(self): + items = [ + _item(0, label="section_header", text="Chapter 1"), + _item(1, text="Chapter 1 prose. " * 100), + _item(2, label="section_header", text="Figure heading"), + _item(3, label="picture", text="Diagram description."), + _item(4, label="caption", text="Figure 2-3. Balance arm."), + _item(5, label="section_header", text="Chapter 3"), + _item(6, text="Chapter 3 prose. " * 100), + ] + # Figure section (items 2-4) is far under 20% of 5000 chars. + lo, hi = _find_expansion_range(items, {3}, has_sections=True, max_chars=5000) + # Never crosses either header + assert (lo, hi) == (2, 4) + + def test_table_in_small_section_stays_section_bounded(self): + items = [ + _item(0, label="section_header", text="Chapter 1"), + _item(1, text="Chapter 1 prose. " * 100), + _item(2, label="section_header", text="Table heading"), + _item(3, label="table", text="Header | Value"), + _item(4, label="section_header", text="Chapter 3"), + _item(5, text="Chapter 3 prose. " * 100), + ] + lo, hi = _find_expansion_range(items, {3}, has_sections=True, max_chars=5000) + assert (lo, hi) == (2, 3) + + def test_text_in_small_section_still_expands_outward(self): + items = [ + _item(0, label="section_header", text="Chapter 1"), + _item(1, text="Chapter 1 prose. " * 100), + _item(2, label="section_header", text="Short note"), + _item(3, text="A brief remark."), + _item(4, label="section_header", text="Chapter 3"), + _item(5, text="Chapter 3 prose. " * 100), + ] + lo, hi = _find_expansion_range(items, {3}, has_sections=True, max_chars=5000) + # Text hits keep growing across section boundaries + assert lo < 2 or hi > 3 + class TestEvidenceAnchors: def test_empty_content(self): @@ -369,6 +409,85 @@ class TestExpandWithItems: # real content — so we get expanded content, not the fallback. assert len(expanded[0].content) > 0 + async def test_picture_expansion_stays_within_section_pages(self, temp_db_path): + from haiku.rag.client import HaikuRAG + + async with HaikuRAG(temp_db_path, create=True) as rag: + items = [ + DocumentItem( + document_id="doc-1", + position=0, + self_ref="#/texts/0", + label="section_header", + text="Chapter 1", + page_numbers=[10], + ), + DocumentItem( + document_id="doc-1", + position=1, + self_ref="#/texts/1", + label="text", + text="Chapter 1 prose. " * 200, + page_numbers=[10], + ), + DocumentItem( + document_id="doc-1", + position=2, + self_ref="#/texts/2", + label="section_header", + text="Figure heading", + page_numbers=[13], + ), + DocumentItem( + document_id="doc-1", + position=3, + self_ref="#/pictures/0", + label="picture", + text="Diagram description.", + page_numbers=[13], + ), + DocumentItem( + document_id="doc-1", + position=4, + self_ref="#/texts/3", + label="caption", + text="Figure 2-3. Balance arm.", + page_numbers=[13], + ), + DocumentItem( + document_id="doc-1", + position=5, + self_ref="#/texts/4", + label="section_header", + text="Chapter 3", + page_numbers=[15], + ), + DocumentItem( + document_id="doc-1", + position=6, + self_ref="#/texts/5", + label="text", + text="Chapter 3 prose. " * 200, + page_numbers=[15], + ), + ] + await rag.document_item_repository.create_items("doc-1", items) + + result = SearchResult( + content="Diagram description.", + score=0.9, + document_id="doc-1", + doc_item_refs=["#/pictures/0"], + page_numbers=[13], + ) + expanded = await expand_with_items( + rag.document_item_repository, "doc-1", [result], 5000 + ) + assert len(expanded) == 1 + assert expanded[0].page_numbers == [13] + assert "Chapter 1 prose" not in expanded[0].content + assert "Chapter 3 prose" not in expanded[0].content + async def test_fragmented_items_preserve_chunk(self, temp_db_path): """When items are fragmented (e.g., list_item children), the original chunk content is preserved if expansion produces less text."""