Keep picture and table hits section-bounded during context expansion

This commit is contained in:
Yiorgis Gozadinos 2026-07-08 12:55:28 +03:00
parent 777ce193f8
commit b14152a45b
No known key found for this signature in database
5 changed files with 133 additions and 3 deletions

View file

@ -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

View file

@ -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`.

View file

@ -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:

View file

@ -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)

View file

@ -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."""