diff --git a/haiku_rag_slim/haiku/rag/context.py b/haiku_rag_slim/haiku/rag/context.py index c7bab5f8..0b6c37ef 100644 --- a/haiku_rag_slim/haiku/rag/context.py +++ b/haiku_rag_slim/haiku/rag/context.py @@ -104,16 +104,18 @@ def _evidence_anchors(content: str, max_chars: int) -> list[str]: 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. +def _clip_window( + content: str, results: list[SearchResult], max_chars: int +) -> tuple[int, int]: + """Return the ``[start, end)`` char window to keep when clipping ``content``. 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 + ``max_chars``-wide window centered on it. Falls back to a prefix window only when no anchor is locatable (heavy drift). """ if max_chars <= 0: - return "" + return (0, 0) evidence_start, evidence_len = -1, 0 for result in results: for anchor in _evidence_anchors(result.content, max_chars): @@ -124,14 +126,45 @@ def _clip_to_budget(content: str, results: list[SearchResult], max_chars: int) - if evidence_start != -1: break if evidence_start == -1: - return content[:max_chars] + return (0, min(len(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 (start, end) + + +def _clip_to_budget(content: str, results: list[SearchResult], max_chars: int) -> str: + """Clip expanded content to ``max_chars``, keeping the matched evidence.""" + start, end = _clip_window(content, results, max_chars) return content[start:end] +def _collect_meta( + spans: list[tuple[int, int, DocumentItem]], +) -> tuple[set[int], list[str], set[str]]: + """Union the page numbers, refs, and labels of the given item spans.""" + pages: set[int] = set() + refs: list[str] = [] + labels: set[str] = set() + for _start, _end, item in spans: + refs.append(item.self_ref) + if item.label: + labels.add(item.label) + pages.update(item.page_numbers) + return pages, refs, labels + + +def _span_in_window( + span: tuple[int, int, DocumentItem], win_start: int, win_end: int +) -> bool: + """Whether an item's char span overlaps the ``[win_start, win_end]`` clip window.""" + start, end, _item = span + if start == end: # zero-width picture position + return win_start <= start <= win_end + return start < win_end and end > win_start + + def _expand_outward( items: list[DocumentItem], center_idx: int, @@ -290,9 +323,11 @@ async def expand_with_items( final_results: list[SearchResult] = [] for range_start, range_end, original_results in merged: content_parts: list[str] = [] - refs: list[str] = [] - pages: set[int] = set() - labels: set[str] = set() + # Char span of each contributing item within the joined content, so + # metadata can be narrowed to whatever survives a budget clip. + item_spans: list[tuple[int, int, DocumentItem]] = [] + cursor = 0 + separator = "\n\n" for pos in range(range_start, range_end + 1): item = pos_to_item.get(pos) @@ -301,55 +336,89 @@ async def expand_with_items( if has_sections and item.label in _NOISE_LABELS: continue if item.text: + if content_parts: + cursor += len(separator) + start = cursor content_parts.append(item.text) - refs.append(item.self_ref) - if item.label: - labels.add(item.label) - pages.update(item.page_numbers) + cursor += len(item.text) + item_spans.append((start, cursor, item)) elif item.label == "picture": # Pictures may legitimately have empty text (no VLM # description configured). Keep their self_ref so the - # downstream image_data lookup can still attach bytes. - refs.append(item.self_ref) - labels.add(item.label) - pages.update(item.page_numbers) + # downstream image_data lookup can still attach bytes. They + # occupy a zero-width position in reading order. + item_spans.append((cursor, cursor, item)) all_headings: list[str] = [] for r in original_results: if r.headings: all_headings.extend(h for h in r.headings if h not in all_headings) - # Carry image_data and picture_captions through expansion so that - # only pictures from the originally retrieved chunks get attached. - # Pictures swept in by section expansion are referenced in `refs` - # for cross-referencing but their bytes are not re-fetched — - # otherwise a single search can balloon the response with adjacent - # figures the model did not actually retrieve. - merged_image_data: dict[str, str] = {} - merged_captions: dict[str, str] = {} - for r in original_results: - if r.image_data: - merged_image_data.update(r.image_data) - if r.picture_captions: - merged_captions.update(r.picture_captions) - - first = original_results[0] + # Anchor identity (chunk_id, content/refs fallbacks) on the + # best-scoring constituent — the chunk that earned the result its + # rank — rather than whichever sits earliest in the document. + first = max(original_results, key=lambda r: r.score) chunk_ids: list[str] = [] for r in original_results: if r.chunk_id and r.chunk_id not in chunk_ids: chunk_ids.append(r.chunk_id) + joined = separator.join(content_parts) + # Expansion should never return less content than the original chunk. # This can happen when item texts are fragmented (e.g., docling splits - # formatted HTML list items into many small text nodes). - 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 + # formatted HTML list items into many small text nodes); fall back to + # the chunk's own content, described by the chunk's own metadata. + if len(joined) < len(first.content): + base_content, base_spans = first.content, None + else: + base_content, base_spans = joined, item_spans + + if len(base_content) > max_chars: + # Clip to the budget, anchored on the primary (highest-scoring) + # chunk, and narrow metadata to whatever survives the window. + win_start, win_end = _clip_window( + base_content, [first, *original_results], max_chars ) + expanded_content = base_content[win_start:win_end] + if base_spans is None: + pages, refs, labels = ( + set(first.page_numbers), + list(first.doc_item_refs), + set(first.labels), + ) + else: + pages, refs, labels = _collect_meta( + [s for s in base_spans if _span_in_window(s, win_start, win_end)] + ) + else: + expanded_content = base_content + if base_spans is None: + pages, refs, labels = ( + set(first.page_numbers), + list(first.doc_item_refs), + set(first.labels), + ) + else: + pages, refs, labels = _collect_meta(base_spans) + + # Carry image_data and picture_captions from the originally retrieved + # chunks, but only for constituents whose refs survive the window — a + # chunk clipped out of the budget must not still ship its image to the + # model. Pictures swept in by section expansion are referenced in + # ``refs`` but their bytes are never re-fetched, so the multimodal + # payload stays bounded to what was actually retrieved and shown. + surviving_refs = set(refs) + merged_image_data: dict[str, str] = {} + merged_captions: dict[str, str] = {} + for r in original_results: + if r.doc_item_refs and not surviving_refs.intersection(r.doc_item_refs): + continue + if r.image_data: + merged_image_data.update(r.image_data) + if r.picture_captions: + merged_captions.update(r.picture_captions) final_results.append( SearchResult( diff --git a/tests/test_context.py b/tests/test_context.py index 3533110a..5bc6dcf3 100644 --- a/tests/test_context.py +++ b/tests/test_context.py @@ -767,6 +767,117 @@ class TestExpandWithItems: # to the model. assert "c2" not in expanded[0].format_for_agent() + async def test_merged_anchor_is_highest_scoring_constituent(self, temp_db_path): + """A merged result's chunk_id anchors on the best-scoring constituent, + not whichever chunk sits earliest in the document.""" + from haiku.rag.client import HaikuRAG + + async with HaikuRAG(temp_db_path, create=True) as rag: + items = [ + DocumentItem( + document_id="doc-1", + position=i, + self_ref=f"#/texts/{i}", + label="text", + text=f"Paragraph {i}. " * 10, + ) + for i in range(5) + ] + await rag.document_item_repository.create_items("doc-1", items) + + # earlier in the document, lower score + r_early = SearchResult( + content="Paragraph 1.", + score=0.40, + chunk_id="c-early", + document_id="doc-1", + doc_item_refs=["#/texts/1"], + ) + # later in the document, higher score — the real hit + r_best = SearchResult( + content="Paragraph 3.", + score=0.95, + chunk_id="c-best", + document_id="doc-1", + doc_item_refs=["#/texts/3"], + ) + expanded = await expand_with_items( + rag.document_item_repository, "doc-1", [r_early, r_best], 5000 + ) + assert len(expanded) == 1 + assert expanded[0].chunk_id == "c-best" + assert expanded[0].score == 0.95 + # provenance still lists both + assert set(expanded[0].chunk_ids) == {"c-early", "c-best"} + + async def test_clipped_merged_result_keeps_anchor_evidence_and_pages( + self, temp_db_path + ): + """When a merged result is clipped to budget, the surviving window is + centered on the anchor (highest-scoring) chunk, and page_numbers reflect + only the content that survived — not the full merged range.""" + 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="LOWMARK " + "a" * 400, + page_numbers=[1], + ), + DocumentItem( + document_id="doc-1", + position=1, + self_ref="#/texts/1", + label="text", + text="b" * 400, + page_numbers=[2], + ), + DocumentItem( + document_id="doc-1", + position=2, + self_ref="#/texts/2", + label="text", + text="c" * 400 + " HIGHMARK", + page_numbers=[3], + ), + ] + await rag.document_item_repository.create_items("doc-1", items) + + r_low = SearchResult( + content="LOWMARK " + "a" * 400, + score=0.4, + chunk_id="c-low", + document_id="doc-1", + doc_item_refs=["#/texts/0"], + page_numbers=[1], + ) + r_high = SearchResult( + content="c" * 400 + " HIGHMARK", + score=0.9, + chunk_id="c-high", + document_id="doc-1", + doc_item_refs=["#/texts/2"], + page_numbers=[3], + ) + expanded = await expand_with_items( + rag.document_item_repository, "doc-1", [r_low, r_high], 500 + ) + assert len(expanded) == 1 + e = expanded[0] + # anchor is the high-scoring chunk, and its evidence survives clipping + assert e.chunk_id == "c-high" + assert "HIGHMARK" in e.content + assert "LOWMARK" not in e.content + # page_numbers reflect only the surviving window, not the full range + assert 3 in e.page_numbers + assert 1 not in e.page_numbers + # refs likewise exclude the clipped-out item + assert "#/texts/0" not in e.doc_item_refs + 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.""" @@ -904,3 +1015,61 @@ class TestExpandWithItemsPictureBytes: "#/pictures/1": "A", "#/pictures/3": "B", } + + async def test_clipped_out_picture_bytes_dropped(self, temp_db_path): + """A lower-scoring picture chunk clipped out of the budget window no + longer contributes its image bytes — the model must not receive an + image the citation and visualization omit.""" + 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="#/pictures/0", + label="picture", + text="LOWPIC " + "a" * 400, + page_numbers=[1], + ), + DocumentItem( + document_id="doc-1", + position=1, + self_ref="#/texts/1", + label="text", + text="b" * 400, + page_numbers=[2], + ), + DocumentItem( + document_id="doc-1", + position=2, + self_ref="#/pictures/1", + label="picture", + text="c" * 400 + " HIGHPIC", + page_numbers=[3], + ), + ] + await rag.document_item_repository.create_items("doc-1", items) + + r_low = SearchResult( + content="LOWPIC " + "a" * 400, + score=0.4, + chunk_id="c-low", + document_id="doc-1", + doc_item_refs=["#/pictures/0"], + image_data={"#/pictures/0": "LOWBYTES"}, + ) + r_high = SearchResult( + content="c" * 400 + " HIGHPIC", + score=0.9, + chunk_id="c-high", + document_id="doc-1", + doc_item_refs=["#/pictures/1"], + image_data={"#/pictures/1": "HIGHBYTES"}, + ) + expanded = await expand_with_items( + rag.document_item_repository, "doc-1", [r_low, r_high], 500 + ) + assert len(expanded) == 1 + assert "#/pictures/0" not in expanded[0].doc_item_refs + assert expanded[0].image_data == {"#/pictures/1": "HIGHBYTES"}