Anchor and clip merged citations on the highest-scoring chunk
A merged search result took its chunk_id from whichever constituent sorted earliest in the document, while its score was the max across the group — so the citation's identity could point at a different, less relevant chunk. Anchor chunk_id and the content/refs fallbacks on the max-score constituent, clip the budget window around that same chunk so its evidence is never trimmed away, and narrow page_numbers, doc_item_refs, and attached image bytes to the items that survive the clip.
This commit is contained in:
Yiorgis Gozadinos
parent
6d86237dd6
commit
c424f51056
2 changed files with 276 additions and 38 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
Loading…
Reference in a new issue