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:
parent
cd1a73b389
commit
b94e13083f
3 changed files with 290 additions and 0 deletions
|
|
@ -1,6 +1,10 @@
|
||||||
# Changelog
|
# Changelog
|
||||||
## [Unreleased]
|
## [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
|
## [0.62.0] - 2026-06-26
|
||||||
|
|
||||||
### Added
|
### Added
|
||||||
|
|
|
||||||
|
|
@ -69,6 +69,64 @@ def _merge_ranges(
|
||||||
return merged
|
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(
|
def _expand_outward(
|
||||||
items: list[DocumentItem],
|
items: list[DocumentItem],
|
||||||
center_idx: int,
|
center_idx: int,
|
||||||
|
|
@ -273,6 +331,10 @@ async def expand_with_items(
|
||||||
expanded_content = "\n\n".join(content_parts)
|
expanded_content = "\n\n".join(content_parts)
|
||||||
if len(expanded_content) < len(first.content):
|
if len(expanded_content) < len(first.content):
|
||||||
expanded_content = 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(
|
final_results.append(
|
||||||
SearchResult(
|
SearchResult(
|
||||||
|
|
|
||||||
|
|
@ -2,6 +2,8 @@ import pytest
|
||||||
|
|
||||||
from haiku.rag.client.documents import _store_document_with_chunks
|
from haiku.rag.client.documents import _store_document_with_chunks
|
||||||
from haiku.rag.context import (
|
from haiku.rag.context import (
|
||||||
|
_clip_to_budget,
|
||||||
|
_evidence_anchors,
|
||||||
_expand_outward,
|
_expand_outward,
|
||||||
_find_expansion_range,
|
_find_expansion_range,
|
||||||
_merge_ranges,
|
_merge_ranges,
|
||||||
|
|
@ -243,6 +245,48 @@ class TestFindExpansionRange:
|
||||||
assert hi >= 2
|
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
|
@pytest.mark.asyncio
|
||||||
class TestExpandWithItems:
|
class TestExpandWithItems:
|
||||||
async def test_unresolvable_refs_returns_original(self, temp_db_path):
|
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
|
# which is less than the chunk's 46 chars — fallback preserves the chunk
|
||||||
assert expanded[0].content == result.content
|
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
|
@pytest.mark.asyncio
|
||||||
class TestExpandWithItemsPictureBytes:
|
class TestExpandWithItemsPictureBytes:
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue