b94e13083f
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.
356 lines
13 KiB
Python
356 lines
13 KiB
Python
"""Section-bounded context expansion for search results.
|
|
|
|
Expands search results with surrounding content from the document using
|
|
the document_items table. The algorithm adapts to document structure:
|
|
|
|
For STRUCTURED documents (containing section_header or title labels):
|
|
1. Resolve matched doc_item_refs to positions in the items table
|
|
2. Find section boundaries around each match (section_header/title labels)
|
|
3. If the section fits within the char budget, include it entirely
|
|
4. If the section exceeds the char budget, expand item-by-item from the
|
|
match center outward, bounded by section edges
|
|
5. If the section is too small (under 20% of max_context_chars), expand
|
|
item-by-item crossing into adjacent sections until the budget is filled.
|
|
This lets small sections (e.g., title+authors) grow into neighboring
|
|
content.
|
|
6. Merge overlapping ranges from multiple results in the same document.
|
|
Adjacent but non-overlapping ranges stay separate to preserve section
|
|
independence.
|
|
|
|
For UNSTRUCTURED documents (no section headers):
|
|
Expand outward item-by-item from the match center until the character
|
|
budget is filled. No noise filtering (unstructured docs typically only
|
|
have text items).
|
|
|
|
In both cases:
|
|
- max_context_chars caps total characters per expanded result
|
|
- Noise labels (footnote, page_header, page_footer, document_index) are
|
|
excluded from content AND budget counting in structured documents
|
|
- Results without doc_item_refs pass through unexpanded
|
|
"""
|
|
|
|
from haiku.rag.store.models.chunk import SearchResult
|
|
from haiku.rag.store.models.document_item import DocumentItem
|
|
from haiku.rag.store.repositories.document_item import DocumentItemRepository
|
|
|
|
_NOISE_LABELS = {"footnote", "page_header", "page_footer", "document_index"}
|
|
_SECTION_BOUNDARY_LABELS = {"section_header", "title"}
|
|
|
|
# Sections with fewer chars than this fraction of max_context_chars are
|
|
# considered too small — expansion falls through to item-by-item outward
|
|
# growth, which naturally crosses into adjacent sections.
|
|
_MIN_SECTION_BUDGET_RATIO = 0.2
|
|
|
|
|
|
def _merge_ranges(
|
|
ranges: list[tuple[int, int, SearchResult]],
|
|
) -> list[tuple[int, int, list[SearchResult]]]:
|
|
"""Merge overlapping ranges. Adjacent but non-overlapping ranges stay separate."""
|
|
if not ranges:
|
|
return []
|
|
|
|
sorted_ranges = sorted(ranges, key=lambda x: x[0])
|
|
merged: list[tuple[int, int, list[SearchResult]]] = []
|
|
cur_min, cur_max, cur_results = (
|
|
sorted_ranges[0][0],
|
|
sorted_ranges[0][1],
|
|
[sorted_ranges[0][2]],
|
|
)
|
|
|
|
for min_idx, max_idx, result in sorted_ranges[1:]:
|
|
if cur_max >= min_idx: # Truly overlapping
|
|
cur_max = max(cur_max, max_idx)
|
|
cur_results.append(result)
|
|
else:
|
|
merged.append((cur_min, cur_max, cur_results))
|
|
cur_min, cur_max, cur_results = min_idx, max_idx, [result]
|
|
|
|
merged.append((cur_min, cur_max, cur_results))
|
|
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,
|
|
max_chars: int,
|
|
skip_noise: bool = False,
|
|
lo_bound: int = 0,
|
|
hi_bound: int | None = None,
|
|
) -> tuple[int, int]:
|
|
"""Expand item-by-item outward from center until char budget is filled.
|
|
|
|
When skip_noise is True, noise labels are excluded from char counting
|
|
(used in structured documents so footnotes don't consume budget).
|
|
|
|
lo_bound and hi_bound constrain expansion (e.g., to section edges).
|
|
"""
|
|
if hi_bound is None:
|
|
hi_bound = len(items) - 1
|
|
lo = hi = center_idx
|
|
center_is_noise = skip_noise and items[center_idx].label in _NOISE_LABELS
|
|
char_count = 0 if center_is_noise else len(items[center_idx].text)
|
|
|
|
while char_count < max_chars:
|
|
grew = False
|
|
if lo > lo_bound:
|
|
lo -= 1
|
|
if not (skip_noise and items[lo].label in _NOISE_LABELS):
|
|
char_count += len(items[lo].text)
|
|
grew = True
|
|
if hi < hi_bound and char_count < max_chars:
|
|
hi += 1
|
|
if not (skip_noise and items[hi].label in _NOISE_LABELS):
|
|
char_count += len(items[hi].text)
|
|
grew = True
|
|
if not grew:
|
|
break
|
|
|
|
return (items[lo].position, items[hi].position)
|
|
|
|
|
|
def _find_expansion_range(
|
|
items: list[DocumentItem],
|
|
matched_positions: set[int],
|
|
has_sections: bool,
|
|
max_chars: int,
|
|
) -> tuple[int, int]:
|
|
"""Find the expansion range for matched positions within a window of items."""
|
|
pos_to_idx = {item.position: i for i, item in enumerate(items)}
|
|
matched_indices = sorted(pos_to_idx[p] for p in matched_positions)
|
|
center_idx = matched_indices[len(matched_indices) // 2]
|
|
|
|
if not has_sections:
|
|
return _expand_outward(items, center_idx, max_chars)
|
|
|
|
# Build section spans: [(start_idx, end_idx), ...]
|
|
headers = [
|
|
i for i, item in enumerate(items) if item.label in _SECTION_BOUNDARY_LABELS
|
|
]
|
|
sections: list[tuple[int, int]] = []
|
|
if headers[0] > 0:
|
|
sections.append((0, headers[0] - 1))
|
|
for j, h in enumerate(headers):
|
|
end = headers[j + 1] - 1 if j + 1 < len(headers) else len(items) - 1
|
|
sections.append((h, end))
|
|
|
|
# Find which section contains the center match
|
|
current = 0
|
|
for j, (start, end) in enumerate(sections):
|
|
if start <= center_idx <= end:
|
|
current = j
|
|
break
|
|
|
|
sec_start, sec_end = sections[current]
|
|
sec_chars = sum(
|
|
len(items[i].text)
|
|
for i in range(sec_start, sec_end + 1)
|
|
if items[i].label not in _NOISE_LABELS
|
|
)
|
|
|
|
min_useful = int(max_chars * _MIN_SECTION_BUDGET_RATIO)
|
|
|
|
if sec_chars <= max_chars and sec_chars >= min_useful:
|
|
# Section fits in char budget — return it regardless of item count
|
|
return (items[sec_start].position, items[sec_end].position)
|
|
|
|
if sec_chars > max_chars:
|
|
# Section too large — expand outward bounded by section edges
|
|
return _expand_outward(
|
|
items,
|
|
center_idx,
|
|
max_chars,
|
|
skip_noise=True,
|
|
lo_bound=sec_start,
|
|
hi_bound=sec_end,
|
|
)
|
|
|
|
# Section too small (e.g., title+authors) — expand across boundaries
|
|
return _expand_outward(items, center_idx, max_chars, skip_noise=True)
|
|
|
|
|
|
_WINDOW_MARGIN = 100
|
|
|
|
|
|
async def expand_with_items(
|
|
document_item_repository: DocumentItemRepository,
|
|
document_id: str,
|
|
results: list[SearchResult],
|
|
max_chars: int,
|
|
) -> list[SearchResult]:
|
|
"""Expand results using the document_items table."""
|
|
all_refs = []
|
|
for result in results:
|
|
all_refs.extend(result.doc_item_refs)
|
|
|
|
ref_positions = await document_item_repository.resolve_refs(document_id, all_refs)
|
|
if not ref_positions:
|
|
return results
|
|
|
|
# Fetch a window of items around matched positions. The margin must be
|
|
# wide enough to find section boundaries (the nearest section_header/title
|
|
# above and below the match).
|
|
all_positions = sorted(ref_positions.values())
|
|
window_margin = _WINDOW_MARGIN
|
|
window_start = max(0, min(all_positions) - window_margin)
|
|
window_end = max(all_positions) + window_margin
|
|
window_items = await document_item_repository.get_items_in_range(
|
|
document_id, window_start, window_end
|
|
)
|
|
|
|
if not window_items:
|
|
return results
|
|
|
|
has_sections = any(item.label in _SECTION_BOUNDARY_LABELS for item in window_items)
|
|
|
|
# Compute expansion ranges per result
|
|
ranges: list[tuple[int, int, SearchResult]] = []
|
|
passthrough: list[SearchResult] = []
|
|
|
|
for result in results:
|
|
matched = {ref_positions[r] for r in result.doc_item_refs if r in ref_positions}
|
|
if not matched:
|
|
passthrough.append(result)
|
|
continue
|
|
|
|
lo, hi = _find_expansion_range(window_items, matched, has_sections, max_chars)
|
|
ranges.append((lo, hi, result))
|
|
|
|
merged = _merge_ranges(ranges)
|
|
|
|
# Build results from the window items
|
|
pos_to_item = {item.position: item for item in window_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()
|
|
|
|
for pos in range(range_start, range_end + 1):
|
|
item = pos_to_item.get(pos)
|
|
if item is None:
|
|
continue
|
|
if has_sections and item.label in _NOISE_LABELS:
|
|
continue
|
|
if item.text:
|
|
content_parts.append(item.text)
|
|
refs.append(item.self_ref)
|
|
if item.label:
|
|
labels.add(item.label)
|
|
pages.update(item.page_numbers)
|
|
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)
|
|
|
|
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]
|
|
|
|
# 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
|
|
)
|
|
|
|
final_results.append(
|
|
SearchResult(
|
|
content=expanded_content,
|
|
score=max(r.score for r in original_results),
|
|
chunk_id=first.chunk_id,
|
|
document_id=first.document_id,
|
|
document_uri=first.document_uri,
|
|
document_title=first.document_title,
|
|
doc_item_refs=refs or first.doc_item_refs,
|
|
page_numbers=sorted(pages) or first.page_numbers,
|
|
headings=all_headings or None,
|
|
labels=sorted(labels) or first.labels,
|
|
image_data=merged_image_data or None,
|
|
picture_captions=merged_captions,
|
|
)
|
|
)
|
|
|
|
return final_results + passthrough
|