danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix context expansion: respect section boundaries, remove max_context_items

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-20 14:53:36 +03:00
parent af7731f4e1
commit ff8ad0879c
No known key found for this signature in database
11 changed files with 114 additions and 117 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
docs/configuration/index.md
View file

@ -99,7 +99,6 @@ research:
search:
limit: 10 # Default number of results to return
max_context_items: 10 # Maximum items in expanded context
max_context_chars: 10000 # Maximum characters in expanded context
vector_index_metric: cosine # cosine, l2, or dot
vector_refine_factor: 30

2
docs/configuration/qa-research.md
View file

@ -7,12 +7,10 @@ Configure search behavior and context expansion:
```yaml
search:
limit: 10 # Default number of results to return
max_context_items: 10 # Maximum items in expanded context
max_context_chars: 10000 # Maximum characters in expanded context
```
- **limit**: Default number of search results to return when no limit is specified. Used by CLI, MCP server, QA, and research workflows. Default: 10
- **max_context_items**: Limits how many document items (paragraphs, list items, etc.) can be included in expanded context. 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.

2
docs/configuration/storage.md
View file

@ -112,7 +112,7 @@ search:
vector_refine_factor: 30 # Re-ranking factor for accuracy
```
For search behavior settings (`limit`, `max_context_items`, `max_context_chars`), see [QA and Research](qa-research.md#search-settings).
For search behavior settings (`limit`, `max_context_chars`), see [QA and Research](qa-research.md#search-settings).
- **vector_index_metric**: Distance metric for vector similarity:
- `cosine`: Cosine similarity (default, best for most embeddings)

1
docs/python.md
View file

@ -376,7 +376,6 @@ Context expansion is automatic and section-aware. For structured documents (with
Configuration:
- **search.max_context_items**: Maximum items in expanded context. Default: 10.
- **search.max_context_chars**: Maximum characters in expanded context. Default: 10000.
**Smart Merging**: When expanded results overlap within the same document, they are automatically merged into a single result with continuous content and the highest relevance score.

2
docs/tuning.md
View file

@ -26,7 +26,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-research.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_items` and `max_context_chars` cap 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 (footnotes, page headers). For unstructured documents, expansion grows outward until the character budget is filled. `max_context_chars` caps expansion to prevent context bloat.
## Tuning Generation

1
evaluations/evaluations/benchmark.py
View file

@ -49,7 +49,6 @@ def build_experiment_metadata(
"embedder_dim": config.embeddings.model.vector_dim,
"chunk_size": config.processing.chunk_size,
"search_limit": config.search.limit,
"max_context_items": config.search.max_context_items,
"max_context_chars": config.search.max_context_chars,
"rerank_provider": config.reranking.model.provider
if config.reranking.model

28
haiku_rag_slim/haiku/rag/client.py
View file

@ -1105,7 +1105,6 @@ class HaikuRAG:
"""
from haiku.rag.context import expand_with_items
max_items = self._config.search.max_context_items
max_chars = self._config.search.max_context_chars
# Group by document_id for efficient processing
@ -1132,7 +1131,6 @@ class HaikuRAG:
self.document_item_repository,
doc_id,
doc_results,
max_items,
max_chars,
)
expanded_results.extend(expanded)
@ -1272,9 +1270,9 @@ class HaikuRAG:
async def visualize_chunk(self, chunk: Chunk) -> list:
"""Render page images with bounding box highlights for a chunk.
Gets the DoclingDocument from the chunk's document, resolves bounding boxes
from chunk metadata, and renders all pages that contain bounding boxes with
yellow/orange highlight overlays.
Expands the chunk's context to find the full section, then resolves
bounding boxes from all items in the expanded range. This ensures
visualization covers all pages the expanded content spans.
Args:
chunk: The chunk to visualize.
@ -1287,6 +1285,8 @@ class HaikuRAG:
from PIL import ImageDraw
from haiku.rag.store.models.chunk import ChunkMetadata
# Get the document structure (from cache if available)
if not chunk.document_id:
return []
@ -1299,9 +1299,23 @@ class HaikuRAG:
if not docling_doc:
return []
# Resolve bounding boxes from chunk metadata
# Expand context to get all doc_item_refs in the section
chunk_meta = chunk.get_chunk_metadata()
bounding_boxes = chunk_meta.resolve_bounding_boxes(docling_doc)
if chunk_meta.doc_item_refs:
search_result = SearchResult(
content=chunk.content,
score=1.0,
chunk_id=chunk.id,
document_id=chunk.document_id,
doc_item_refs=chunk_meta.doc_item_refs,
page_numbers=chunk_meta.page_numbers,
)
expanded = await self.expand_context([search_result])
refs = expanded[0].doc_item_refs if expanded else chunk_meta.doc_item_refs
meta = ChunkMetadata(doc_item_refs=refs)
else:
meta = chunk_meta
bounding_boxes = meta.resolve_bounding_boxes(docling_doc)
if not bounding_boxes:
return []

1
haiku_rag_slim/haiku/rag/config/models.py
View file

@ -174,7 +174,6 @@ class ProcessingConfig(BaseModel):
class SearchConfig(BaseModel):
limit: int = 10
max_context_items: int = 10
max_context_chars: int = 10000
vector_index_metric: Literal["cosine", "l2", "dot"] = "cosine"
vector_refine_factor: int = 30

72
haiku_rag_slim/haiku/rag/context.py
View file

@ -6,12 +6,16 @@ the document_items table. The algorithm adapts to document structure:
For STRUCTURED documents (containing section_header or title labels):
1. Resolve matched doc_item_refs to positions in the items table
2. Find section boundaries around each match (section_header/title labels)
3. If the section fits within the budget, include it entirely
4. If the section exceeds the budget, OR the section is too small (under
20% of max_context_chars), expand item-by-item from the match center
outward, skipping noise labels. This lets small sections (e.g., a
title+authors area) grow into the next section's content.
5. Merge overlapping ranges from multiple results in the same document
3. If the section fits within the char budget, include it entirely
4. If the section exceeds the char budget, expand item-by-item from the
match center outward, bounded by section edges
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.
6. Merge overlapping ranges from multiple results in the same document.
Adjacent but non-overlapping ranges stay separate to preserve section
independence.
For UNSTRUCTURED documents (no section headers):
Expand outward item-by-item from the match center until the character
@ -20,7 +24,6 @@ For UNSTRUCTURED documents (no section headers):
In both cases:
- max_context_chars caps total characters per expanded result
- max_context_items caps total items per expanded result
- Noise labels (footnote, page_header, page_footer, document_index) are
excluded from content AND budget counting in structured documents
- Results without doc_item_refs pass through unexpanded
@ -42,7 +45,7 @@ _MIN_SECTION_BUDGET_RATIO = 0.2
def _merge_ranges(
ranges: list[tuple[int, int, SearchResult]],
) -> list[tuple[int, int, list[SearchResult]]]:
"""Merge overlapping or adjacent ranges."""
"""Merge overlapping ranges. Adjacent but non-overlapping ranges stay separate."""
if not ranges:
return []
@ -55,7 +58,7 @@ def _merge_ranges(
)
for min_idx, max_idx, result in sorted_ranges[1:]:
if cur_max >= min_idx - 1: # Overlapping or adjacent
if cur_max >= min_idx: # Truly overlapping
cur_max = max(cur_max, max_idx)
cur_results.append(result)
else:
@ -69,27 +72,32 @@ def _merge_ranges(
def _expand_outward(
items: list[DocumentItem],
center_idx: int,
max_items: int,
max_chars: int,
skip_noise: bool = False,
lo_bound: int = 0,
hi_bound: int | None = None,
) -> tuple[int, int]:
"""Expand item-by-item outward from center until budget is filled.
"""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).
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)
while char_count < max_chars and hi - lo + 1 < max_items:
while char_count < max_chars:
grew = False
if lo > 0:
if lo > lo_bound:
lo -= 1
if not (skip_noise and items[lo].label in _NOISE_LABELS):
char_count += len(items[lo].text)
grew = True
if hi < len(items) - 1 and char_count < max_chars:
if hi < hi_bound and char_count < max_chars:
hi += 1
if not (skip_noise and items[hi].label in _NOISE_LABELS):
char_count += len(items[hi].text)
@ -104,7 +112,6 @@ def _find_expansion_range(
items: list[DocumentItem],
matched_positions: set[int],
has_sections: bool,
max_items: int,
max_chars: int,
) -> tuple[int, int]:
"""Find the expansion range for matched positions within a window of items."""
@ -113,7 +120,7 @@ def _find_expansion_range(
center_idx = matched_indices[len(matched_indices) // 2]
if not has_sections:
return _expand_outward(items, center_idx, max_items, max_chars)
return _expand_outward(items, center_idx, max_chars)
# Build section spans: [(start_idx, end_idx), ...]
headers = [
@ -140,23 +147,34 @@ def _find_expansion_range(
if items[i].label not in _NOISE_LABELS
)
# Section fits nicely in the budget — return it as-is
min_useful = int(max_chars * _MIN_SECTION_BUDGET_RATIO)
if min_useful <= sec_chars <= max_chars and sec_end - sec_start + 1 <= max_items:
if sec_chars <= max_chars and sec_chars >= min_useful:
# Section fits in char budget — return it regardless of item count
return (items[sec_start].position, items[sec_end].position)
# Section is too large or too small — expand item-by-item from center.
# For too-large sections this stays within budget.
# For too-small sections (e.g., title+authors) this naturally grows
# into adjacent sections until the budget is filled.
return _expand_outward(items, center_idx, max_items, max_chars, skip_noise=True)
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,
)
# Section too small (e.g., title+authors) — expand across boundaries
return _expand_outward(items, center_idx, max_chars, skip_noise=True)
_WINDOW_MARGIN = 100
async def expand_with_items(
document_item_repository: DocumentItemRepository,
document_id: str,
results: list[SearchResult],
max_items: int,
max_chars: int,
) -> list[SearchResult]:
"""Expand results using the document_items table."""
@ -172,7 +190,7 @@ async def expand_with_items(
# wide enough to find section boundaries (the nearest section_header/title
# above and below the match).
all_positions = sorted(ref_positions.values())
window_margin = max_items * 10
window_margin = _WINDOW_MARGIN
window_start = max(0, min(all_positions) - window_margin)
window_end = max(all_positions) + window_margin
window_items = await document_item_repository.get_items_in_range(
@ -194,9 +212,7 @@ async def expand_with_items(
passthrough.append(result)
continue
lo, hi = _find_expansion_range(
window_items, matched, has_sections, max_items, max_chars
)
lo, hi = _find_expansion_range(window_items, matched, has_sections, max_chars)
ranges.append((lo, hi, result))
merged = _merge_ranges(ranges)

93
tests/test_context.py
View file

@ -49,11 +49,12 @@ class TestMergeRanges:
assert len(merged) == 1
assert merged[0] == (0, 15, [r1, r2])
def test_adjacent(self):
def test_adjacent_stay_separate(self):
r1, r2 = _result(), _result()
merged = _merge_ranges([(0, 5, r1), (6, 10, r2)])
assert len(merged) == 1
assert merged[0] == (0, 10, [r1, r2])
assert len(merged) == 2
assert merged[0] == (0, 5, [r1])
assert merged[1] == (6, 10, [r2])
def test_sorts_by_position(self):
r1, r2 = _result(), _result()
@ -65,7 +66,7 @@ class TestMergeRanges:
class TestExpandOutward:
def test_basic_expansion(self):
items = [_item(i, text=f"{'x' * 100}") for i in range(10)]
lo, hi = _expand_outward(items, 5, max_items=10, max_chars=500)
lo, hi = _expand_outward(items, 5, max_chars=500)
assert lo <= 5
assert hi >= 5
total = sum(
@ -76,16 +77,9 @@ class TestExpandOutward:
# Should be around 500 chars (may overshoot by one item)
assert total >= 400
def test_respects_max_items(self):
items = [_item(i, text="x") for i in range(100)]
lo, hi = _expand_outward(items, 50, max_items=5, max_chars=999999)
count = hi - lo + 1
# May overshoot by 1-2 items due to alternating expansion
assert count <= 7
def test_respects_max_chars(self):
items = [_item(i, text=f"{'x' * 200}") for i in range(20)]
lo, hi = _expand_outward(items, 10, max_items=999, max_chars=500)
lo, hi = _expand_outward(items, 10, max_chars=500)
total = sum(
len(items[i].text)
for i in range(lo, hi + 1)
@ -96,12 +90,12 @@ class TestExpandOutward:
def test_center_at_start(self):
items = [_item(i) for i in range(10)]
lo, hi = _expand_outward(items, 0, max_items=5, max_chars=999999)
lo, hi = _expand_outward(items, 0, max_chars=999999)
assert lo == 0
def test_center_at_end(self):
items = [_item(i) for i in range(10)]
lo, hi = _expand_outward(items, 9, max_items=5, max_chars=999999)
lo, hi = _expand_outward(items, 9, max_chars=999999)
assert hi == 9
def test_skip_noise_excludes_from_char_count(self):
@ -113,7 +107,7 @@ class TestExpandOutward:
_item(4, label="footnote", text="f" * 5000),
_item(5, text="d" * 100),
]
lo, hi = _expand_outward(items, 2, max_items=10, max_chars=500, skip_noise=True)
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
assert lo <= 0
@ -125,11 +119,17 @@ class TestExpandOutward:
_item(1, label="document_index", text="x" * 10000),
_item(2, text="b" * 200),
]
lo, hi = _expand_outward(items, 1, max_items=10, max_chars=500, skip_noise=True)
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)
assert lo == 8
assert hi == 12
class TestFindExpansionRange:
def _structured_items(self):
@ -146,36 +146,47 @@ class TestFindExpansionRange:
def test_structured_returns_section(self):
items = self._structured_items()
lo, hi = _find_expansion_range(
items, {1}, has_sections=True, max_items=20, max_chars=5000
)
lo, hi = _find_expansion_range(items, {1}, has_sections=True, max_chars=5000)
# Should return the Introduction section (items 0-3)
assert lo == 0
assert hi == 3
def test_structured_different_section(self):
items = self._structured_items()
lo, hi = _find_expansion_range(
items, {5}, has_sections=True, max_items=20, max_chars=5000
)
lo, hi = _find_expansion_range(items, {5}, has_sections=True, max_chars=5000)
# Should return the Methods section (items 4-6)
assert lo == 4
assert hi == 6
def test_structured_large_section_falls_back_to_outward(self):
def test_structured_large_section_bounded_by_section(self):
items = [
_item(0, label="section_header", text="Big Section"),
] + [_item(i, text="x" * 1000) for i in range(1, 20)]
# Section has 19 * 1000 = 19000 chars, way over 5000 budget
lo, hi = _find_expansion_range(
items, {10}, has_sections=True, max_items=50, max_chars=5000
)
# Should NOT return the full section
lo, hi = _find_expansion_range(items, {10}, has_sections=True, max_chars=5000)
# Should NOT return the full section, but should stay within it
total = sum(
len(items[i].text) for i in range(lo, hi + 1) if items[i].position >= lo
)
assert total < 10000
def test_structured_section_with_many_items_returned_whole(self):
"""A section that fits in char budget is returned even with many items."""
items = (
[
_item(0, label="section_header", text="Section"),
]
+ [_item(i, text="x" * 200) for i in range(1, 20)]
+ [
_item(20, label="section_header", text="Next"),
]
)
# Section has 19 * 200 = 3800 chars + header, under 5000 and over min_useful
lo, hi = _find_expansion_range(items, {10}, has_sections=True, max_chars=5000)
# Should return entire section despite 20 items
assert lo == 0
assert hi == 19
def test_structured_small_section_expands_outward(self):
items = [
_item(0, label="title", text="Paper Title"),
@ -186,26 +197,20 @@ class TestFindExpansionRange:
_item(5, text="Intro content. " * 50),
]
# Title section (items 0-1) is tiny (~25 chars) < 20% of 5000
lo, hi = _find_expansion_range(
items, {0}, has_sections=True, max_items=20, max_chars=5000
)
lo, hi = _find_expansion_range(items, {0}, has_sections=True, max_chars=5000)
# Should expand past the title section into the abstract
assert hi >= 3
def test_unstructured_expands_outward(self):
items = [_item(i, text=f"Paragraph {i}. " * 10) for i in range(10)]
lo, hi = _find_expansion_range(
items, {5}, has_sections=False, max_items=20, max_chars=5000
)
lo, hi = _find_expansion_range(items, {5}, has_sections=False, max_chars=5000)
assert lo < 5
assert hi > 5
def test_multiple_matched_positions_uses_center(self):
items = [_item(i, text="x" * 100) for i in range(20)]
# Match at positions 3 and 7, center should be index for position 5 (median)
lo, hi = _find_expansion_range(
items, {3, 7}, has_sections=False, max_items=5, max_chars=999999
)
# Use a char budget that forces partial expansion so center matters
lo, hi = _find_expansion_range(items, {3, 7}, has_sections=False, max_chars=500)
center = (lo + hi) // 2
# Center should be around position 5
assert 3 <= center <= 7
@ -219,9 +224,7 @@ class TestFindExpansionRange:
]
# Section non-noise chars: ~260 chars (items 0,1,3). Under 5000 budget.
# The footnote's 10000 chars should NOT count.
lo, hi = _find_expansion_range(
items, {1}, has_sections=True, max_items=20, max_chars=5000
)
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
@ -233,9 +236,7 @@ class TestFindExpansionRange:
_item(2, label="section_header", text="First Section"),
_item(3, text="Section content."),
]
lo, hi = _find_expansion_range(
items, {0}, has_sections=True, max_items=20, max_chars=5000
)
lo, hi = _find_expansion_range(items, {0}, has_sections=True, max_chars=5000)
# Match is in preamble section (items 0-1), which is small
# Should expand outward into the first section
assert hi >= 2
@ -262,7 +263,7 @@ class TestExpandWithItems:
doc_item_refs=["#/texts/999999"],
)
expanded = await expand_with_items(
rag.document_item_repository, doc.id, [result], 10, 5000
rag.document_item_repository, doc.id, [result], 5000
)
assert len(expanded) == 1
assert expanded[0].content == "original"
@ -312,7 +313,7 @@ class TestExpandWithItems:
doc_item_refs=["#/texts/1"],
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [result], 10, 5000
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
# The TOC section's only non-header item is document_index (noise).
@ -376,7 +377,7 @@ class TestExpandWithItems:
doc_item_refs=["#/texts/1", "#/texts/2", "#/texts/3", "#/texts/4"],
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [result], 10, 5000
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
# Expansion produces "Steps\n\nClick\n\n+\n\nAdd a New Service" = 38 chars

28
tests/test_context_enhancement.py
View file

@ -80,7 +80,6 @@ def small_chunk_config() -> AppConfig:
"""Config with small chunk size to force splitting."""
config = AppConfig()
config.processing.chunk_size = 32
config.search.max_context_items = 25
config.search.max_context_chars = 10000
return config
@ -305,33 +304,6 @@ async def test_format_for_agent_output(temp_db_path, small_chunk_config):
assert "Content:" in formatted
@pytest.mark.vcr()
async def test_max_items_limit_caps_expansion(temp_db_path):
"""Expansion should respect max_context_items limit."""
config = AppConfig()
config.processing.chunk_size = 32
config.search.max_context_items = 2 # Very restrictive
docling_doc = create_list_document()
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
doc = await create_document_with_docling(client, docling_doc, "Limit Test")
assert doc.id is not None
results = await client.search("grapes", limit=1)
assert len(results) > 0
expanded = await client.expand_context(results)
# With max_items=2, expansion should be limited
content = expanded[0].content.lower()
item_count = sum(
1 for item in ["apples", "bananas", "oranges", "grapes"] if item in content
)
# Should have at most 2 items (the limit)
assert item_count <= 2, f"Expected at most 2 items, got {item_count}"
async def test_expand_context_single_item_document(temp_db_path):
"""Test expand_context with a single-item document."""
from haiku.rag.store.models.document import Document