From b94e13083f609bdaf94dc388d0eb6b9dc5731cf8 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 27 Jun 2026 09:46:36 +0300 Subject: [PATCH] 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. --- CHANGELOG.md | 4 + haiku_rag_slim/haiku/rag/context.py | 62 ++++++++ tests/test_context.py | 224 ++++++++++++++++++++++++++++ 3 files changed, 290 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 84a694cc..d235b1a2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py index 38eb531c..ee1ea311 100644 --- a/haiku_rag_slim/haiku/rag/context.py +++ b/haiku_rag_slim/haiku/rag/context.py @@ -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( diff --git a/tests/test_context.py b/tests/test_context.py index 2c807a53..4c6a8239 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -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="" + 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 "M" * 500 in expanded[0].content + @pytest.mark.asyncio class TestExpandWithItemsPictureBytes: