Compare commits

...

1 commit

Author SHA1 Message Date
Yiorgis Gozadinos
d97aa15af9
Keep footnotes and matched items in expanded context
_build_result applied the noise-label filter to every item in the range,
including the ones the result matched on. A hit on a footnote or index
entry returned its section with the matched text removed, the clip anchor
could not find the evidence and fell back to a prefix window, and the
un-merge path rebuilt through the same filter.

Noise is now a set of positions computed once per group by
_noise_positions: noise-labelled items minus the matched ones.
_expand_outward and _build_result take that set instead of a flag, so the
matched item is kept in content and counted toward the budget.

Footnotes leave the noise set. They carry sources, cross-references and
clarifications, and docling attaches table and figure footnotes to the
table itself, so the filter was dropping part of the table. The noise set
is page_header, page_footer and document_index.

Refs #609
2026-09-07 13:52:14 +03:00
6 changed files with 200 additions and 58 deletions

View file

@ -51,6 +51,8 @@
- `haiku-rag mcp` covers the configured `lancedb.databases` set. `sources` on
`search_documents`, `search_documents_by_image` and `execute_code`; `source`
on `get_document`; an unknown name is a tool error. `DocumentInfo.source`.
- `footnote` items are no longer filtered from expanded context. The noise
labels are `page_header`, `page_footer` and `document_index`.
### Fixed
@ -59,6 +61,8 @@
- Past `analysis.code_timeout` a sandbox program starts no further host call.
Files served from memory and in-code `search()` / `list_documents()` were
not checked against the deadline.
- Context expansion keeps the item a result matched on when it carries a
noise label, and counts it toward the character budget.
### Removed

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: 5
- **max_context_chars**: Hard limit on total characters in expanded content. Default: 5000.
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.
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 (page headers, page footers, table of contents). 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

@ -367,7 +367,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. 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.
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 (page headers, page footers, table of contents). 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

@ -28,7 +28,7 @@ When configured, a cross-encoder reranker re-scores 10x the requested candidates
`limit` controls how many results reach the LLM. More candidates improve recall but increase token usage. See [Search Settings](configuration/qa.md#search-settings).
Context expansion is automatic and section-aware. Search results are expanded to include surrounding content from the same document section. For structured documents, expansion stays within section boundaries and filters noise (footnotes, page headers). For unstructured documents, expansion grows outward until the character budget is filled. `max_context_chars` caps expansion to prevent context bloat.
Context expansion is automatic and section-aware. Search results are expanded to include surrounding content from the same document section. For structured documents, expansion stays within section boundaries and filters noise (page headers, page footers and the table of contents). For unstructured documents, expansion grows outward until the character budget is filled. `max_context_chars` caps expansion to prevent context bloat.
## Tuning Generation
@ -89,7 +89,7 @@ Press `c` on a chunk to see the expanded context that would be fed to the RAG ca
- The expanded text. Section-aware expansion stays within section boundaries on structured documents and fills `max_context_chars` outward on unstructured ones.
- Source document, content type, and relevance score.
- Filtered noise. Footnotes, page headers and footers are excluded from structured documents.
- Filtered noise. Page headers, page footers and the table of contents are excluded from structured documents.
If `qa.model.vision = true` is set, the modal also renders the picture bytes attached to that chunk, so you see exactly what the vision model would receive.

View file

@ -27,17 +27,19 @@ For UNSTRUCTURED documents (no section headers):
In both cases:
- max_context_chars caps total characters per expanded result
- Noise labels (footnote, page_header, page_footer, document_index) are
excluded from content AND budget counting in structured documents
- Noise labels (page_header, page_footer, document_index) are
excluded from content AND budget counting in structured documents,
except the items a result matched on
- Results without doc_item_refs pass through unexpanded
"""
from collections.abc import Set as AbstractSet
from typing import Any
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.store.models.document_item import DocumentItem
_NOISE_LABELS = {"footnote", "page_header", "page_footer", "document_index"}
_NOISE_LABELS = {"page_header", "page_footer", "document_index"}
_SECTION_BOUNDARY_LABELS = {"section_header", "title"}
# Labels whose pertinent unit is the item plus its own section: expansion
@ -166,37 +168,43 @@ def _span_in_window(
return start < win_end and end > win_start
def _noise_positions(
items: list[DocumentItem], matched_positions: set[int]
) -> set[int]:
"""Positions skipped as noise; an item a result matched on never is."""
return {
item.position for item in items if item.label in _NOISE_LABELS
} - matched_positions
def _expand_outward(
items: list[DocumentItem],
center_idx: int,
max_chars: int,
skip_noise: bool = False,
noise: AbstractSet[int] = frozenset(),
lo_bound: int = 0,
hi_bound: int | None = None,
) -> tuple[int, int]:
"""Expand item-by-item outward from center until char budget is filled.
When skip_noise is True, noise labels are excluded from char counting
(used in structured documents so footnotes don't consume budget).
Items at ``noise`` positions are excluded from char counting.
lo_bound and hi_bound constrain expansion (e.g., to section edges).
"""
if hi_bound is None:
hi_bound = len(items) - 1
lo = hi = center_idx
center_is_noise = skip_noise and items[center_idx].label in _NOISE_LABELS
char_count = 0 if center_is_noise else len(items[center_idx].text)
char_count = len(items[center_idx].text)
while char_count < max_chars:
grew = False
if lo > lo_bound:
lo -= 1
if not (skip_noise and items[lo].label in _NOISE_LABELS):
if items[lo].position not in noise:
char_count += len(items[lo].text)
grew = True
if hi < hi_bound and char_count < max_chars:
hi += 1
if not (skip_noise and items[hi].label in _NOISE_LABELS):
if items[hi].position not in noise:
char_count += len(items[hi].text)
grew = True
if not grew:
@ -219,6 +227,8 @@ def _find_expansion_range(
if not has_sections:
return _expand_outward(items, center_idx, max_chars)
noise = _noise_positions(items, matched_positions)
# Build section spans: [(start_idx, end_idx), ...]
headers = [
i for i, item in enumerate(items) if item.label in _SECTION_BOUNDARY_LABELS
@ -241,7 +251,7 @@ def _find_expansion_range(
sec_chars = sum(
len(items[i].text)
for i in range(sec_start, sec_end + 1)
if items[i].label not in _NOISE_LABELS
if items[i].position not in noise
)
min_useful = int(max_chars * _MIN_SECTION_BUDGET_RATIO)
@ -253,12 +263,7 @@ def _find_expansion_range(
if sec_chars > max_chars:
# Section too large — expand outward bounded by section edges
return _expand_outward(
items,
center_idx,
max_chars,
skip_noise=True,
lo_bound=sec_start,
hi_bound=sec_end,
items, center_idx, max_chars, noise, lo_bound=sec_start, hi_bound=sec_end
)
# Picture/table hits stay section-bounded: their pertinent unit is the
@ -267,7 +272,7 @@ def _find_expansion_range(
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)
return _expand_outward(items, center_idx, max_chars, noise)
def _group_lost_constituent(built: SearchResult, group: list[SearchResult]) -> bool:
@ -300,7 +305,7 @@ def _build_result(
range_end: int,
original_results: list[SearchResult],
pos_to_item: dict[int, DocumentItem],
has_sections: bool,
noise: AbstractSet[int],
max_chars: int,
) -> SearchResult:
"""Build one expanded result from the items in ``[range_start, range_end]``."""
@ -313,9 +318,7 @@ def _build_result(
for pos in range(range_start, range_end + 1):
item = pos_to_item.get(pos)
if item is None:
continue
if has_sections and item.label in _NOISE_LABELS:
if item is None or pos in noise:
continue
if item.text:
if content_parts:
@ -457,8 +460,21 @@ def expand_with_items(
ranges: list[tuple[int, int, SearchResult]] = []
passthrough: list[SearchResult] = []
def matched_positions(group: list[SearchResult]) -> set[int]:
return {
ref_positions[ref]
for result in group
for ref in result.doc_item_refs
if ref in ref_positions
}
def noise_for(group: list[SearchResult]) -> set[int]:
if not has_sections:
return set()
return _noise_positions(window_items, matched_positions(group))
for result in results:
matched = {ref_positions[r] for r in result.doc_item_refs if r in ref_positions}
matched = matched_positions([result])
if not matched:
passthrough.append(result)
continue
@ -473,7 +489,7 @@ def expand_with_items(
final_results: list[SearchResult] = []
for range_start, range_end, group in merged:
built = _build_result(
range_start, range_end, group, pos_to_item, has_sections, max_chars
range_start, range_end, group, pos_to_item, noise_for(group), max_chars
)
if len(group) > 1 and _group_lost_constituent(built, group):
# The merged window cannot afford every constituent's evidence:
@ -483,7 +499,7 @@ def expand_with_items(
lo, hi = constituent_range[id(result)]
final_results.append(
_build_result(
lo, hi, [result], pos_to_item, has_sections, max_chars
lo, hi, [result], pos_to_item, noise_for([result]), max_chars
)
)
continue

View file

@ -103,32 +103,20 @@ class TestExpandOutward:
lo, hi = _expand_outward(items, 9, max_chars=999999)
assert hi == 9
def test_skip_noise_excludes_from_char_count(self):
def test_noise_positions_excluded_from_char_count(self):
items = [
_item(0, text="a" * 100),
_item(1, label="footnote", text="f" * 5000),
_item(1, label="page_header", text="f" * 5000),
_item(2, text="b" * 100),
_item(3, text="c" * 100),
_item(4, label="footnote", text="f" * 5000),
_item(4, label="page_header", text="f" * 5000),
_item(5, text="d" * 100),
]
lo, hi = _expand_outward(items, 2, max_chars=500, skip_noise=True)
# Footnotes (5000 chars each) should NOT count toward budget
# So we should expand past them
lo, hi = _expand_outward(items, 2, max_chars=500, noise={1, 4})
# Noise (5000 chars each) does not count, so expansion passes it
assert lo <= 0
assert hi >= 5
def test_noise_center_gets_zero_chars(self):
items = [
_item(0, text="a" * 200),
_item(1, label="document_index", text="x" * 10000),
_item(2, text="b" * 200),
]
lo, hi = _expand_outward(items, 1, max_chars=500, skip_noise=True)
# Center is noise, should start at 0 chars and expand outward
assert lo == 0
assert hi == 2
def test_respects_bounds(self):
items = [_item(i, text="x" * 100) for i in range(20)]
lo, hi = _expand_outward(items, 10, max_chars=999999, lo_bound=8, hi_bound=12)
@ -224,16 +212,28 @@ class TestFindExpansionRange:
items = [
_item(0, label="section_header", text="Section"),
_item(1, text="Real content." * 10),
_item(2, label="footnote", text="x" * 10000),
_item(2, label="page_header", text="x" * 10000),
_item(3, text="More content." * 10),
]
# Section non-noise chars: ~260 chars (items 0,1,3). Under 5000 budget.
# The footnote's 10000 chars should NOT count.
# The header's 10000 chars do not count.
lo, hi = _find_expansion_range(items, {1}, has_sections=True, max_chars=5000)
# Should return full section (it fits in budget excluding noise)
assert lo == 0
assert hi == 3
def test_matched_noise_item_counts_toward_section_chars(self):
items = [
_item(0, label="section_header", text="Contents"),
_item(1, text="Real content." * 10),
_item(2, label="document_index", text="x" * 10000),
_item(3, text="More content." * 10),
]
lo, hi = _find_expansion_range(items, {2}, has_sections=True, max_chars=5000)
# The matched index is 10000 chars: the section is over budget and
# expansion stays on the match.
assert (lo, hi) == (2, 2)
def test_items_before_first_header_form_section(self):
items = [
_item(0, text="Preamble text."),
@ -373,12 +373,11 @@ class TestExpandWithItems:
assert len(expanded) == 1
assert expanded[0].content == "original"
async def test_noise_only_range_preserves_original(self, temp_db_path):
"""When noise filtering removes all content, original chunk is preserved."""
async def test_matched_noise_item_survives_expansion(self, temp_db_path):
"""A result that matched on a noise-labelled item keeps that item."""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(temp_db_path, create=True) as rag:
# Structured document where the matched item's section has only noise
items = [
DocumentItem(
document_id="doc-1",
@ -421,11 +420,134 @@ class TestExpandWithItems:
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
# The TOC section's only non-header item is document_index (noise).
# The section_header "Table of Contents" has text but _expand_outward
# with skip_noise crosses into the Introduction section which has
# real content — so we get expanded content, not the fallback.
assert len(expanded[0].content) > 0
assert "x" * 2000 in expanded[0].content
assert "#/texts/1" in expanded[0].doc_item_refs
async def test_footnotes_are_kept_in_expanded_content(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",
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/texts/1",
label="text",
text="Body paragraph. " * 80,
),
DocumentItem(
document_id="doc-1",
position=2,
self_ref="#/texts/2",
label="footnote",
text="1 See Smith v Jones, para 12.",
),
]
await rag.document_item_repository.create_items("doc-1", items)
result = SearchResult(
content="Body paragraph.",
score=0.9,
document_id="doc-1",
doc_item_refs=["#/texts/1"],
)
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
assert "See Smith v Jones" in expanded[0].content
assert "#/texts/2" in expanded[0].doc_item_refs
async def test_unmatched_noise_excluded_in_structured_document(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",
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/texts/1",
label="text",
text="Body paragraph. " * 80,
),
DocumentItem(
document_id="doc-1",
position=2,
self_ref="#/texts/2",
label="page_header",
text="RUNNING HEADER",
),
DocumentItem(
document_id="doc-1",
position=3,
self_ref="#/texts/3",
label="text",
text="Closing paragraph.",
),
]
await rag.document_item_repository.create_items("doc-1", items)
result = SearchResult(
content="Body paragraph.",
score=0.9,
document_id="doc-1",
doc_item_refs=["#/texts/1"],
)
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
assert "Closing paragraph." in expanded[0].content
assert "RUNNING HEADER" not in expanded[0].content
assert "#/texts/2" not in expanded[0].doc_item_refs
async def test_unstructured_document_keeps_noise_labels(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="page_header",
text="RUNNING HEADER",
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/texts/1",
label="text",
text="Body paragraph.",
),
]
await rag.document_item_repository.create_items("doc-1", items)
result = SearchResult(
content="Body paragraph.",
score=0.9,
document_id="doc-1",
doc_item_refs=["#/texts/1"],
)
expanded = await _fetch_and_expand(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
assert "RUNNING HEADER" in expanded[0].content
async def test_picture_expansion_stays_within_section_pages(self, temp_db_path):
from haiku.rag.client import HaikuRAG
@ -1559,6 +1681,6 @@ def test_build_result_skips_positions_with_no_item():
),
}
built = _build_result(0, 3, [original], pos_to_item, False, 5000)
built = _build_result(0, 3, [original], pos_to_item, set(), 5000)
assert built.content == "first\n\nlast"