Hard-cap context expansion at max_context_chars

A single oversized document_items row (e.g. a spreadsheet converted to one
table) expanded far past search.max_context_chars and could overflow the
model context window. _expand_outward only used the budget as a soft
accumulation threshold and expand_with_items never capped the joined result.

Add _clip_to_budget to clip each expanded result to max_context_chars,
returning a window centered on the matched chunk (via _evidence_anchors) so
the retrieved evidence survives the cut.
This commit is contained in:
Yiorgis Gozadinos 2026-06-27 09:46:36 +03:00
parent cd1a73b389
commit b94e13083f
No known key found for this signature in database
3 changed files with 290 additions and 0 deletions

View file

@ -1,6 +1,10 @@
# Changelog
## [Unreleased]
### Fixed
- Context expansion now hard-caps each expanded search result at `search.max_context_chars`, anchored on the matched chunk. A single oversized `document_items` row (e.g. a spreadsheet converted to one table) no longer pushes a result far past the budget and overflows the model context window.
## [0.62.0] - 2026-06-26
### Added

View file

@ -69,6 +69,64 @@ def _merge_ranges(
return merged
_MIN_ANCHOR = 128
def _evidence_anchors(content: str, max_chars: int) -> list[str]:
"""Substrings of a matched chunk used to locate it inside expanded content.
Tries the exact chunk text first, then a substantial central slice that
tolerates edge formatting drift between the chunk and the joined item text.
Every anchor is bounded by ``max_chars`` so it always fits inside the window
the caller returns, and is never shorter than ``_MIN_ANCHOR`` (unless the
chunk itself is) so a short, common substring can't anchor by accident.
"""
if not content:
return []
anchors: list[str] = []
if len(content) <= max_chars:
anchors.append(content)
if max_chars > 0:
min_anchor = min(_MIN_ANCHOR, max_chars, len(content))
target = min(max_chars // 2, len(content) // 2)
target = min(max(target, min_anchor), max_chars, len(content))
if target < len(content):
mid = len(content) // 2
start = max(0, mid - target // 2)
anchors.append(content[start : start + target])
if not anchors:
anchors.append(content[:max_chars] if max_chars > 0 else content)
return anchors
def _clip_to_budget(content: str, results: list[SearchResult], max_chars: int) -> str:
"""Clip expanded content to ``max_chars``, keeping the matched evidence.
Anchors on the first locatable result in ``results`` order (the primary
chunk that supplies the expanded result's identity) and returns a
``max_chars``-wide window centered on it. Falls back to a prefix cut only
when no anchor is locatable (heavy drift).
"""
if max_chars <= 0:
return ""
evidence_start, evidence_len = -1, 0
for result in results:
for anchor in _evidence_anchors(result.content, max_chars):
idx = content.find(anchor)
if idx != -1:
evidence_start, evidence_len = idx, len(anchor)
break
if evidence_start != -1:
break
if evidence_start == -1:
return content[:max_chars]
center = evidence_start + evidence_len // 2
start = max(0, center - max_chars // 2)
end = min(len(content), start + max_chars)
start = max(0, end - max_chars)
return content[start:end]
def _expand_outward(
items: list[DocumentItem],
center_idx: int,
@ -273,6 +331,10 @@ async def expand_with_items(
expanded_content = "\n\n".join(content_parts)
if len(expanded_content) < len(first.content):
expanded_content = first.content
if len(expanded_content) > max_chars:
expanded_content = _clip_to_budget(
expanded_content, original_results, max_chars
)
final_results.append(
SearchResult(

View file

@ -2,6 +2,8 @@ import pytest
from haiku.rag.client.documents import _store_document_with_chunks
from haiku.rag.context import (
_clip_to_budget,
_evidence_anchors,
_expand_outward,
_find_expansion_range,
_merge_ranges,
@ -243,6 +245,48 @@ class TestFindExpansionRange:
assert hi >= 2
class TestEvidenceAnchors:
def test_empty_content(self):
assert _evidence_anchors("", 5000) == []
def test_anchors_never_exceed_budget(self):
# Even when max_chars is smaller than the minimum anchor length, no
# anchor may exceed the window it has to fit inside.
anchors = _evidence_anchors("z" * 1000, 50)
assert anchors
assert all(len(a) <= 50 for a in anchors)
def test_tiny_content_uses_full_text(self):
anchors = _evidence_anchors("short", 5000)
assert anchors == ["short"]
def test_long_content_offers_full_and_central_slice(self):
content = "L" * 500 + "M" * 500
anchors = _evidence_anchors(content, 5000)
assert content in anchors
# A strictly shorter central slice is also offered for drift tolerance.
assert any(a != content and a in content for a in anchors)
class TestClipToBudget:
def test_zero_budget_returns_empty(self):
result = SearchResult(content="anything", score=0.9, document_id="d")
assert _clip_to_budget("some long content", [result], 0) == ""
def test_no_anchor_found_falls_back_to_prefix(self):
content = "xyz" * 1000
result = SearchResult(content="NOTPRESENT", score=0.9, document_id="d")
assert _clip_to_budget(content, [result], 100) == content[:100]
def test_centers_window_on_evidence(self):
marker = "UNIQUE_MATCH_TEXT"
content = "A" * 300_000 + marker + "B" * 300_000
result = SearchResult(content=marker, score=0.9, document_id="d")
clipped = _clip_to_budget(content, [result], 10_000)
assert len(clipped) <= 10_000
assert marker in clipped
@pytest.mark.asyncio
class TestExpandWithItems:
async def test_unresolvable_refs_returns_original(self, temp_db_path):
@ -387,6 +431,186 @@ class TestExpandWithItems:
# which is less than the chunk's 46 chars — fallback preserves the chunk
assert expanded[0].content == result.content
async def test_oversized_item_clipped_to_budget(self, temp_db_path):
"""A single huge item (e.g. a whole spreadsheet as one table) is clipped
to the budget, centered on the matched text."""
from haiku.rag.client import HaikuRAG
marker = "UNIQUE_MATCH_TEXT"
async with HaikuRAG(temp_db_path, create=True) as rag:
items = [
DocumentItem(
document_id="doc-1",
position=0,
self_ref="#/tables/0",
label="table",
text="A" * 300_000 + marker + "B" * 300_000,
),
]
await rag.document_item_repository.create_items("doc-1", items)
result = SearchResult(
content=marker,
score=0.9,
document_id="doc-1",
doc_item_refs=["#/tables/0"],
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [result], 10_000
)
assert len(expanded) == 1
assert len(expanded[0].content) <= 10_000
assert marker in expanded[0].content
async def test_giant_neighbor_does_not_blow_budget(self, temp_db_path):
"""An adjacent giant item swept in by expansion cannot blow the budget,
and the matched item's text survives."""
from haiku.rag.client import HaikuRAG
matched_text = "matched small text here"
async with HaikuRAG(temp_db_path, create=True) as rag:
items = [
DocumentItem(
document_id="doc-1",
position=0,
self_ref="#/texts/0",
label="text",
text=matched_text,
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/tables/1",
label="table",
text="C" * 600_000,
),
]
await rag.document_item_repository.create_items("doc-1", items)
result = SearchResult(
content=matched_text,
score=0.9,
document_id="doc-1",
doc_item_refs=["#/texts/0"],
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [result], 10_000
)
assert len(expanded) == 1
assert len(expanded[0].content) <= 10_000
assert matched_text in expanded[0].content
async def test_small_items_under_budget_not_clipped(self, temp_db_path):
"""Normal small-item expansion is untouched — clipping never triggers."""
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="text",
text="First paragraph here.",
),
DocumentItem(
document_id="doc-1",
position=1,
self_ref="#/texts/1",
label="text",
text="Second matched paragraph.",
),
DocumentItem(
document_id="doc-1",
position=2,
self_ref="#/texts/2",
label="text",
text="Third paragraph here.",
),
]
await rag.document_item_repository.create_items("doc-1", items)
result = SearchResult(
content="Second matched paragraph.",
score=0.9,
document_id="doc-1",
doc_item_refs=["#/texts/1"],
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
assert len(expanded[0].content) < 5000
# All three items expanded in, nothing truncated.
assert "First paragraph here." in expanded[0].content
assert "Second matched paragraph." in expanded[0].content
assert "Third paragraph here." in expanded[0].content
async def test_original_chunk_larger_than_budget_is_capped(self, temp_db_path):
"""When the floor restores an original chunk bigger than the budget, the
hard cap still wins (returning less than the original chunk)."""
from haiku.rag.client import HaikuRAG
marker = "CENTRAL_MARKER_" + "Z" * 200
big_chunk = "A" * 10_000 + marker + "A" * 10_000
async with HaikuRAG(temp_db_path, create=True) as rag:
items = [
DocumentItem(
document_id="doc-1",
position=0,
self_ref="#/texts/0",
label="text",
text="tiny",
),
]
await rag.document_item_repository.create_items("doc-1", items)
result = SearchResult(
content=big_chunk,
score=0.9,
document_id="doc-1",
doc_item_refs=["#/texts/0"],
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [result], 5000
)
assert len(expanded) == 1
assert len(expanded[0].content) <= 5000
assert marker in expanded[0].content
async def test_fuzzy_match_preserves_central_marker(self, temp_db_path):
"""The chunk's text need not be verbatim in the joined item text: a clean
central marker is still located via the central-slice anchor."""
from haiku.rag.client import HaikuRAG
marker = "M" * 2000
async with HaikuRAG(temp_db_path, create=True) as rag:
items = [
DocumentItem(
document_id="doc-1",
position=0,
self_ref="#/tables/0",
label="table",
text="A" * 300_000 + marker + "B" * 300_000,
),
]
await rag.document_item_repository.create_items("doc-1", items)
# Chunk content has edge formatting that is NOT verbatim in the item
# text, but the central marker is identical.
result = SearchResult(
content="<edge>" + marker + "</edge>",
score=0.9,
document_id="doc-1",
doc_item_refs=["#/tables/0"],
)
expanded = await expand_with_items(
rag.document_item_repository, "doc-1", [result], 10_000
)
assert len(expanded) == 1
assert len(expanded[0].content) <= 10_000
assert "M" * 500 in expanded[0].content
@pytest.mark.asyncio
class TestExpandWithItemsPictureBytes: