Coverage for src/lilbee/retrieval/query/formatting.py: 100%
156 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-04 17:08 +0000
1"""Source formatting, context templating, and LLM citation extraction."""
3from __future__ import annotations
5import re
6from pathlib import Path
8from lilbee.data.store import ChunkType, CitationRecord, SearchChunk
10CONTEXT_TEMPLATE = """Context:
11{context}
13Question: {question}"""
16# Bracketed citation groups: [1], [1, 2], [1-3], [1, 3-5]. Models mix all of
17# these despite being asked for single [n] markers; matching only [n] made
18# cited_sources under-count and fed JSON consumers false-negative grounding.
19_CITE_GROUP_RE = re.compile(r"\[(\d+(?:\s*[-,]\s*\d+)*)\]")
20_CITE_RANGE_RE = re.compile(r"(\d+)\s*-\s*(\d+)")
21# Ranges wider than this are page spans or line numbers, not citation lists.
22_MAX_CITE_RANGE = 32
24# Heading shapes that open a model-authored citation block: anchored to the
25# start of the text as well as a preceding newline (an answer that is nothing
26# but a fabricated block still strips), optionally markdown-decorated (ATX
27# hashes, emphasis around the word or the colon: "**Sources:**", "*References:*").
28_HEADING_WORDS = r"(?:(?:Key\s+)?Sources|References|Bibliography|Citations)"
29_CITE_HEADING = (
30 r"(?:\n{1,3}|\A)[ \t]*(?:#+\s*)?[*_]{0,3}" + _HEADING_WORDS + r"[*_]{0,3}\s*:?\s*[*_]{0,3}"
31)
32# One list line: a bullet, arrow, "[1]" or "1." marker and the rest of its line.
33_CITE_LIST_LINE = r"[ \t]*(?:[-*•→\[]|\d+[.)])[^\n]*"
35# A heading line followed by a list. The list is required so prose discussing
36# such a heading is not clipped; the match ends with the list, not end-of-text,
37# so an answer resuming after its citations keeps the continuation. Items may be
38# blank-line separated, which markdown does routinely; stopping at the first
39# blank line would leave the rest of a fabricated list in the answer.
40_LLM_CITATION_BLOCK_RE = re.compile(
41 _CITE_HEADING + r"\n\s*" + _CITE_LIST_LINE + r"(?:\n+" + _CITE_LIST_LINE + r")*",
42 re.IGNORECASE,
43)
45# A citation-style heading at the very end of the text, nothing after it yet.
46# Mid-stream this is ambiguous (the next line decides list vs prose), so it is
47# held back rather than shown; at end of stream it is a dangling artifact of a
48# citation block the model never finished, and is dropped either way.
49_TRAILING_HEADING_RE = re.compile(_CITE_HEADING + r"\s*$", re.IGNORECASE)
52def display_source_path(source: str) -> str:
53 """Render a chunk's source as an absolute path with ``~`` expansion.
55 Source values in the store are stored relative to ``documents_dir`` so the
56 database is portable across machines. For display we resolve back to the
57 user's filesystem and substitute ``~`` for the home directory so the path
58 is unambiguous without being noisy.
60 Falls back to the raw source string only if resolution itself fails (an
61 exotic OSError such as a symlink loop). A missing file is not a failure:
62 ``resolve(strict=False)`` still returns the absolute path to where the
63 file would be, so a moved documents directory renders its old location.
64 """
65 from lilbee.data.ingest.discovery import resolve_source_path
67 candidate = resolve_source_path(source)
68 try:
69 resolved = candidate.resolve(strict=False)
70 except OSError:
71 return source
72 home = Path.home()
73 try:
74 return f"~/{resolved.relative_to(home)}"
75 except ValueError:
76 return str(resolved)
79_WEB_PREFIX = "_web"
80_SOURCE_SEP = " · "
83def _source_label(source: str) -> str:
84 """Readable name for a source. Web-ingested docs collapse to ``host · slug``;
85 local files keep their documents-dir-relative path."""
86 prefix = f"{_WEB_PREFIX}/"
87 if not source.startswith(prefix):
88 return source
89 segments = [p for p in source.removeprefix(prefix).split("/") if p != "index.md"]
90 if not segments:
91 return source
92 host = segments[0].removeprefix("www.")
93 slug = segments[-1].removesuffix(".md")
94 return host if slug == host else f"{host}{_SOURCE_SEP}{slug}"
97def _source_file_url(source: str) -> str | None:
98 """A ``file://`` URL to the source on disk, so a reader can click it open, or
99 None when the path can't be resolved to an absolute location."""
100 from lilbee.data.ingest.discovery import resolve_source_path
102 try:
103 return resolve_source_path(source).resolve(strict=False).as_uri()
104 except (OSError, ValueError):
105 return None
108def _source_locator(result: SearchChunk) -> str:
109 """The ``, page N`` / ``, lines A-B`` suffix for a source line, or ''."""
110 location = _location_suffix(result)
111 return f", {location}" if location else ""
114def _format_citation(citation: CitationRecord) -> str:
115 """Format a single wiki transitive citation as an indented attribution line."""
116 source_display = display_source_path(citation["source_filename"])
117 if citation["page_start"] or citation["page_end"]:
118 ps, pe = citation["page_start"], citation["page_end"]
119 pages = f"page {ps}" if ps == pe else f"pages {ps}-{pe}"
120 return f" → {source_display}, {pages}"
121 if citation["line_start"] or citation["line_end"]:
122 ls, le = citation["line_start"], citation["line_end"]
123 lines = f"line {ls}" if ls == le else f"lines {ls}-{le}"
124 return f" → {source_display}, {lines}"
125 return f" → {source_display}"
128def _location_suffix(result: SearchChunk) -> str:
129 """The page or line span of a chunk, or empty when neither applies.
131 Zero means "no location": PDF chunks whose page metadata was missing are
132 stored with page 0, so a locator is only rendered when at least one end
133 of the span is set. Sole owner of the rule -- ``_source_locator`` adds
134 the leading separator and nothing else.
135 """
136 if result.content_type == "pdf" and (result.page_start or result.page_end):
137 ps, pe = result.page_start, result.page_end
138 return f"page {ps}" if ps == pe else f"pages {ps}-{pe}"
139 if result.content_type == "code" and (result.line_start or result.line_end):
140 ls, le = result.line_start, result.line_end
141 return f"line {ls}" if ls == le else f"lines {ls}-{le}"
142 return ""
145def source_markdown_link(source: str) -> str:
146 """A bare source name as the same clickable ``[label](file-url)`` markdown a
147 live answer's Sources block uses; the plain label when no path resolves.
148 Public so restored transcripts render sources identically to live ones."""
149 label = _source_label(source)
150 url = _source_file_url(source)
151 return f"[{label}]({url})" if url else label
154def format_source(result: SearchChunk, citations: list[CitationRecord] | None = None) -> str:
155 """Format a source as a clickable, readable citation: a ``[label](file-url)``
156 markdown link plus any page/line locator. Web docs render as ``host · slug``;
157 wiki chunks append their indented transitive citations.
158 """
159 head = source_markdown_link(result.source)
160 if result.chunk_type is ChunkType.WIKI and citations:
161 return "\n".join([head, *(_format_citation(c) for c in citations)])
162 return f"{head}{_source_locator(result)}"
165def unique_sources(results: list[SearchChunk]) -> list[SearchChunk]:
166 """The first chunk of each distinct source, in retrieval order. A source's
167 1-based position here is the citation number the model emits (see
168 ``build_context``) and the number shown in the Sources block, so the answer's
169 ``[n]`` markers and the source list always agree."""
170 seen: set[str] = set()
171 out: list[SearchChunk] = []
172 for r in results:
173 if r.source not in seen:
174 seen.add(r.source)
175 out.append(r)
176 return out
179def _context_header(result: SearchChunk) -> str:
180 """One-line provenance for a context block: source name plus location.
182 Without it the answering model sees bare numbered text: it cannot
183 attribute a claim to a named document, notice two chunks share a source,
184 or confirm it is reading the document the user asked about.
185 """
186 location = _location_suffix(result)
187 if location:
188 return f"{result.source}, {location}"
189 return result.source
192def build_context(results: list[SearchChunk]) -> str:
193 """Number each passage by its source file, not its position, so citation
194 numbers are stable while streaming and map 1:1 to the Sources block. Passages
195 from the same file share a number.
197 Each block carries a provenance header: without it the answering model sees
198 bare numbered text and cannot attribute a claim to a named document, notice
199 two passages share a source, or confirm it is reading the document asked for.
200 """
201 order: dict[str, int] = {}
202 for r in results:
203 order.setdefault(r.source, len(order) + 1)
204 return "\n\n".join(f"[{order[r.source]}] ({_context_header(r)})\n{r.chunk}" for r in results)
207# Grepped by consumers to know an answer carries its own Sources list (the
208# no-results toast; the pill row, which must not stack a second list).
209SOURCES_BLOCK_MARKER = "\n\nSources:\n"
212def format_sources_block(
213 results: list[SearchChunk],
214 citations_map: dict[str, list[CitationRecord]] | None = None,
215) -> str:
216 """The authoritative numbered ``Sources:`` block appended to an answer. Each
217 unique source is numbered to match the ``[n]`` markers the model emitted, so
218 every inline citation resolves to a line here. '' when there are no sources."""
219 sources = unique_sources(results)
220 if not sources:
221 return ""
222 # A markdown ordered list, so a Markdown renderer puts each source on its own
223 # line (plain " → path" lines collapse into one soft-wrapped paragraph) and
224 # the list number is the citation number the model cited inline.
225 lines = [
226 f"{i}. {format_source(r, citations=(citations_map or {}).get(r.source))}"
227 for i, r in enumerate(sources, 1)
228 ]
229 return SOURCES_BLOCK_MARKER + "\n" + "\n".join(lines)
232def _extract_cited_indices(text: str) -> set[int]:
233 """Extract citation references from LLM answer text: [1], [1, 2], [1-3]."""
234 indices: set[int] = set()
235 for m in _CITE_GROUP_RE.finditer(text):
236 group = m.group(1)
237 remainder = _CITE_RANGE_RE.sub("", group)
238 for start, end in _CITE_RANGE_RE.findall(group):
239 lo, hi = int(start), int(end)
240 if lo <= hi <= lo + _MAX_CITE_RANGE:
241 indices.update(range(lo, hi + 1))
242 indices.update(int(n) for n in re.findall(r"\d+", remainder))
243 return indices
246def _identifier_shaped(stem: str) -> bool:
247 """Whether a filename stem is distinctive enough to match in prose.
249 A stem carrying a digit or a separator ("survey_report", "ARC-00000482")
250 only appears in an answer when the model names the document; a bare word
251 stem ("notes") collides with ordinary prose and cannot be trusted.
252 """
253 return any(c.isdigit() or c in "_-" for c in stem)
256def cited_subset(answer: str, sources: list[SearchChunk]) -> list[SearchChunk]:
257 """The sources the answer actually referenced, in order (empty if none).
259 ``[n]`` markers are the primary signal, and ``n`` indexes the unique sources
260 (matching ``build_context``/``format_sources_block``). Name mentions count
261 too: context blocks show the model each source's name, and models often
262 attribute by name ("according to survey_report.pdf") instead of by marker,
263 which otherwise reads as an ungrounded answer to JSON consumers.
264 """
265 uniq = unique_sources(sources)
266 cited = _extract_cited_indices(answer)
267 picked = {i - 1 for i in cited if 1 <= i <= len(uniq)}
268 lowered = answer.lower()
269 for i, source in enumerate(uniq):
270 if i in picked:
271 continue
272 name = Path(source.source).name.lower()
273 stem = Path(source.source).stem
274 if _mentions(lowered, name) or (
275 _identifier_shaped(stem) and _mentions(lowered, stem.lower())
276 ):
277 picked.add(i)
278 return [uniq[i] for i in sorted(picked)]
281def _mentions(answer_lower: str, needle: str) -> bool:
282 """Whether *answer_lower* names *needle* as a whole token.
284 Plain containment marks a source cited whenever its name embeds in a
285 longer one ("log-1" inside "catalog-10", "notes.md" inside
286 "footnotes.md"), which inflates the grounding signal. Filename
287 characters are what must not abut the match; surrounding punctuation
288 and whitespace still count as a mention.
289 """
290 return re.search(rf"(?<![\w.-]){re.escape(needle)}(?![\w-])", answer_lower) is not None
293def _stream_safe_prefix(text: str) -> str:
294 """The prefix of *text* safe to show: a model-authored citation block
295 (heading plus list) is removed, and a bare trailing citation heading is
296 withheld until whatever follows disambiguates it.
298 Kept rstrip-free so the streaming filter can track emitted length exactly;
299 ``strip_llm_citations`` layers the final rstrip on top for one-shot answers.
300 """
301 cleaned = _LLM_CITATION_BLOCK_RE.sub("", text)
302 return _TRAILING_HEADING_RE.sub("", cleaned)
305def strip_llm_citations(text: str) -> str:
306 """Remove an LLM-generated trailing citation block (or dangling citation
307 heading) from answer text. Prose that merely mentions such a heading stays."""
308 return _stream_safe_prefix(text).rstrip()
311class StreamingCitationFilter:
312 """Suppress a model-generated ``Sources:``/``References:`` citation block as
313 it streams, so only lilbee's authoritative source list reaches the reader.
315 A one-shot answer can be stripped after the fact, but streamed tokens are
316 already on screen by the time the block is recognizable. This feeds the
317 answer in incrementally and only releases text once it is certain not to be
318 the start of a citation block: everything up to the last newline is safe to
319 show, the final (possibly partial) line is held back until more text
320 arrives, and a bare citation heading is held until the next line shows
321 whether a list (a citation block, dropped) or prose (a legitimate mention,
322 shown) follows. A heading left dangling when the stream ends is dropped.
323 """
325 def __init__(self) -> None:
326 self._buffer = ""
327 self._emitted = 0
329 def feed(self, text: str) -> str:
330 """Accept the next answer chunk; return the portion safe to show now."""
331 self._buffer += text
332 stripped = _stream_safe_prefix(self._buffer)
333 cut = stripped.rfind("\n")
334 committed = stripped if cut == -1 else stripped[:cut]
335 if len(committed) > self._emitted:
336 out = committed[self._emitted :]
337 self._emitted = len(committed)
338 return out
339 return ""
341 def flush(self) -> str:
342 """Release any remaining safe text once the stream has ended."""
343 final = _stream_safe_prefix(self._buffer)
344 if len(final) > self._emitted:
345 out = final[self._emitted :]
346 self._emitted = len(final)
347 return out
348 return ""
350 @property
351 def answer(self) -> str:
352 """The full answer shown so far, with any citation block removed."""
353 return strip_llm_citations(self._buffer)