Coverage for src/lilbee/retrieval/query/neighbors.py: 100%
96 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""Neighbor-window expansion: widen retrieved passages with adjacent chunks.
3A hit that lands in the middle of an argument loses the sentences before and
4after it. After context selection, each selected chunk pulls up to N adjacent
5chunks per side from its own source and merges them into one contiguous
6passage, deduplicating the overlap text adjacent chunks share from chunking.
7The widened passage keeps the original chunk's score and citation identity;
8only its text and page/line span change.
9"""
11from __future__ import annotations
13import re
14from typing import TYPE_CHECKING
16if TYPE_CHECKING:
17 from collections.abc import Callable
19 from lilbee.data.store import SearchChunk, Store
21# Below this a suffix-prefix match is coincidence, not chunker overlap; deduping
22# it would delete real text at an overlap-free seam.
23_MIN_OVERLAP = 16
25# xberg markdown breadcrumb ("# Guide > ## Install\n\n"); heading-context
26# content lines never start with "#", so only the breadcrumb matches.
27_BREADCRUMB_RE = re.compile(r"\A#{1,6} [^\n]*\n\n")
30def _overlap_chars(left: str, right: str) -> int:
31 """Length of the longest suffix of *left* that is a prefix of *right*.
33 Adjacent chunks share the chunker's overlap verbatim, so the longest match
34 is the shared region. A fully contained text matches whole, which keeps a
35 re-merge of an already-widened passage idempotent. Longest length first,
36 so the first match wins.
37 """
38 for k in range(min(len(left), len(right)), 0, -1):
39 if left.endswith(right[:k]):
40 return k
41 return 0
44def _seam(left: str, right: str) -> tuple[int, str]:
45 """(overlap length, effective right text) for one adjacent seam.
47 A heading breadcrumb on the right hides the chunker overlap, so the
48 stripped form is tried too and wins only when it reveals one. A match
49 shorter than ``_MIN_OVERLAP`` counts only when one side is fully
50 contained in the other.
51 """
52 k = _overlap_chars(left, right)
53 if k >= _MIN_OVERLAP or k == len(right) or (k and k == len(left)):
54 return k, right
55 stripped = _BREADCRUMB_RE.sub("", right)
56 if stripped != right:
57 ks = _overlap_chars(left, stripped)
58 if ks >= _MIN_OVERLAP or (ks and ks in (len(stripped), len(left))):
59 return ks, stripped
60 return 0, right
63def _merge_span(
64 span: list[int],
65 texts: dict[int, str],
66 seams: dict[tuple[int, int], tuple[int, str]],
67) -> str:
68 """Merge the span's texts in index order, deduplicating seam overlaps.
70 *seams* caches per-pair seam scans so a caller re-merging shrinking spans
71 (the budget shed loop) pays for each scan once.
72 """
73 merged = texts[span[0]]
74 prev = span[0]
75 for index in span[1:]:
76 key = (prev, index)
77 if key not in seams:
78 seams[key] = _seam(texts[prev], texts[index])
79 k, effective = seams[key]
80 tail = effective[k:]
81 if not tail:
82 continue
83 merged = merged + tail if k else f"{merged}\n{tail}"
84 prev = index
85 return merged
88def merge_adjacent_texts(texts: list[str]) -> str:
89 """Concatenate adjacent chunk texts, deduplicating their shared overlap.
91 Adjacent texts with no real overlap (a chunk_overlap=0 build, or only a
92 coincidental short match) join with a newline seam rather than gluing two
93 words together or deleting text.
94 """
95 return _merge_span(list(range(len(texts))), dict(enumerate(texts)), {})
98def expand_neighbors(
99 results: list[SearchChunk],
100 store: Store,
101 radius: int,
102 budget: int,
103 cost: Callable[[str], int],
104 *,
105 exclude: Callable[[str], bool] | None = None,
106) -> list[SearchChunk]:
107 """Widen each result with up to *radius* adjacent same-source chunks.
109 Results are processed in rank order and only spend *budget*, the tokens
110 left over after the originals were fitted: a widened passage whose extra
111 cost does not fit sheds its farthest neighbors first and falls back to
112 the original text, so expansion is always trimmed before any original
113 chunk. An index that is itself selected, or already claimed by a
114 higher-ranked expansion, is never pulled again, so no passage text is
115 duplicated (a document routed whole expands to nothing). A neighbor whose
116 text matches *exclude* is treated as absent, ending the run at its side,
117 so expansion cannot re-import text an upstream filter dropped.
118 """
119 if budget <= 0:
120 # Nothing can be spent, so skip the per-source store fetches entirely:
121 # on a tight window every widen attempt would fail against a zero
122 # budget after paying for the reads.
123 return results
124 centers: dict[str, set[int]] = {}
125 for r in results:
126 centers.setdefault(r.source, set()).add(r.chunk_index)
127 rows = _fetch_neighbor_rows(store, centers, radius)
128 if exclude is not None:
129 rows = {
130 key: row
131 for key, row in rows.items()
132 if row.chunk_index in centers.get(row.source, ()) or not exclude(row.chunk)
133 }
134 if not rows:
135 return results
136 claimed = {source: set(indices) for source, indices in centers.items()}
137 remaining = budget
138 expanded: list[SearchChunk] = []
139 for r in results:
140 widened, spent = _widen(r, rows, claimed[r.source], radius, remaining, cost)
141 remaining -= spent
142 expanded.append(widened)
143 return expanded
146def _fetch_neighbor_rows(
147 store: Store, centers: dict[str, set[int]], radius: int
148) -> dict[tuple[str, int], SearchChunk]:
149 """Every candidate neighbor row, fetched with one store call per source.
151 The centers are fetched alongside their neighbors so the caller can tell
152 whether the document was re-ingested since the search snapshot; indices past
153 the end of a document are simply absent from the reply.
154 """
155 rows: dict[tuple[str, int], SearchChunk] = {}
156 for source, owned in centers.items():
157 wanted = sorted(
158 {
159 index
160 for center in owned
161 for index in range(center - radius, center + radius + 1)
162 if index >= 0
163 }
164 )
165 for row in store.get_chunks_by_indices(source, wanted):
166 rows[(source, row.chunk_index)] = row
167 return rows
170def _neighbor_run(
171 result: SearchChunk,
172 rows: dict[tuple[str, int], SearchChunk],
173 claimed: set[int],
174 step: int,
175 radius: int,
176) -> list[int]:
177 """Contiguous free neighbor indices on one side of the center, nearest first."""
178 indices: list[int] = []
179 for offset in range(1, radius + 1):
180 index = result.chunk_index + step * offset
181 if index in claimed or (result.source, index) not in rows:
182 break
183 indices.append(index)
184 return indices
187def _widen(
188 result: SearchChunk,
189 rows: dict[tuple[str, int], SearchChunk],
190 claimed: set[int],
191 radius: int,
192 remaining: int,
193 cost: Callable[[str], int],
194) -> tuple[SearchChunk, int]:
195 """One result widened within *remaining* tokens: (chunk, tokens spent).
197 Neighbors extend from the center until a missing, selected, or already
198 claimed index stops each side. While the widened text over-spends, the
199 farthest neighbor is shed first (a tie sheds the trailing side, keeping
200 the text that leads up to the hit); shedding everything keeps the
201 original chunk untouched.
202 """
203 center = result.chunk_index
204 current = rows.get((result.source, center))
205 if current is not None and current.chunk != result.chunk:
206 # Re-ingested since the search: the neighbor rows are a different
207 # chunking, so splicing them would invent text and a page span.
208 return result, 0
209 left = _neighbor_run(result, rows, claimed, -1, radius)
210 right = _neighbor_run(result, rows, claimed, +1, radius)
211 texts = {center: result.chunk}
212 for index in [*left, *right]:
213 texts[index] = rows[(result.source, index)].chunk
214 seams: dict[tuple[int, int], tuple[int, str]] = {}
215 while left or right:
216 span = sorted([*left, center, *right])
217 merged = _merge_span(span, texts, seams)
218 extra = cost(merged) - cost(result.chunk)
219 if extra <= remaining:
220 claimed.update(index for index in span if index != center)
221 neighbors = [rows[(result.source, index)] for index in span if index != center]
222 return _widened_copy(result, merged, neighbors), extra
223 if right and (not left or right[-1] - center >= center - left[-1]):
224 right.pop()
225 else:
226 left.pop()
227 return result, 0
230def _widened_copy(result: SearchChunk, merged: str, neighbors: list[SearchChunk]) -> SearchChunk:
231 """The result with widened text and a truthfully recomputed page/line span."""
232 spans = [result, *neighbors]
233 return result.model_copy(
234 update={
235 "chunk": merged,
236 "page_start": min(s.page_start for s in spans),
237 "page_end": max(s.page_end for s in spans),
238 "line_start": min(s.line_start for s in spans),
239 "line_end": max(s.line_end for s in spans),
240 }
241 )