Coverage for src/lilbee/retrieval/query/formatting.py: 100%
160 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-17 10:02 +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, is_memory_source
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 if is_memory_source(source):
151 return label
152 url = _source_file_url(source)
153 return f"[{label}]({url})" if url else label
156def format_source(result: SearchChunk, citations: list[CitationRecord] | None = None) -> str:
157 """Format a source as a clickable, readable citation: a ``[label](file-url)``
158 markdown link plus any page/line locator. Web docs render as ``host · slug``;
159 wiki chunks append their indented transitive citations. Memory rows render
160 as their plain ``memory:<id>`` label: they have no file to link.
161 """
162 if result.memory_id is not None:
163 return _source_label(result.source)
164 head = source_markdown_link(result.source)
165 if result.chunk_type is ChunkType.WIKI and citations:
166 return "\n".join([head, *(_format_citation(c) for c in citations)])
167 return f"{head}{_source_locator(result)}"
170def unique_sources(results: list[SearchChunk]) -> list[SearchChunk]:
171 """The first chunk of each distinct source, in retrieval order. A source's
172 1-based position here is the citation number the model emits (see
173 ``build_context``) and the number shown in the Sources block, so the answer's
174 ``[n]`` markers and the source list always agree."""
175 seen: set[str] = set()
176 out: list[SearchChunk] = []
177 for r in results:
178 if r.source not in seen:
179 seen.add(r.source)
180 out.append(r)
181 return out
184def _context_header(result: SearchChunk) -> str:
185 """One-line provenance for a context block: source name plus location.
187 Without it the answering model sees bare numbered text: it cannot
188 attribute a claim to a named document, notice two chunks share a source,
189 or confirm it is reading the document the user asked about.
190 """
191 location = _location_suffix(result)
192 if location:
193 return f"{result.source}, {location}"
194 return result.source
197def build_context(results: list[SearchChunk]) -> str:
198 """Number each passage by its source file, not its position, so citation
199 numbers are stable while streaming and map 1:1 to the Sources block. Passages
200 from the same file share a number.
202 Each block carries a provenance header: without it the answering model sees
203 bare numbered text and cannot attribute a claim to a named document, notice
204 two passages share a source, or confirm it is reading the document asked for.
205 """
206 order: dict[str, int] = {}
207 for r in results:
208 order.setdefault(r.source, len(order) + 1)
209 return "\n\n".join(f"[{order[r.source]}] ({_context_header(r)})\n{r.chunk}" for r in results)
212# Grepped by consumers to know an answer carries its own Sources list (the
213# no-results toast; the pill row, which must not stack a second list).
214SOURCES_BLOCK_MARKER = "\n\nSources:\n"
217def format_sources_block(
218 results: list[SearchChunk],
219 citations_map: dict[str, list[CitationRecord]] | None = None,
220) -> str:
221 """The authoritative numbered ``Sources:`` block appended to an answer. Each
222 unique source is numbered to match the ``[n]`` markers the model emitted, so
223 every inline citation resolves to a line here. '' when there are no sources."""
224 sources = unique_sources(results)
225 if not sources:
226 return ""
227 # A markdown ordered list, so a Markdown renderer puts each source on its own
228 # line (plain " → path" lines collapse into one soft-wrapped paragraph) and
229 # the list number is the citation number the model cited inline.
230 lines = [
231 f"{i}. {format_source(r, citations=(citations_map or {}).get(r.source))}"
232 for i, r in enumerate(sources, 1)
233 ]
234 return SOURCES_BLOCK_MARKER + "\n" + "\n".join(lines)
237def _extract_cited_indices(text: str) -> set[int]:
238 """Extract citation references from LLM answer text: [1], [1, 2], [1-3]."""
239 indices: set[int] = set()
240 for m in _CITE_GROUP_RE.finditer(text):
241 group = m.group(1)
242 remainder = _CITE_RANGE_RE.sub("", group)
243 for start, end in _CITE_RANGE_RE.findall(group):
244 lo, hi = int(start), int(end)
245 if lo <= hi <= lo + _MAX_CITE_RANGE:
246 indices.update(range(lo, hi + 1))
247 indices.update(int(n) for n in re.findall(r"\d+", remainder))
248 return indices
251def _identifier_shaped(stem: str) -> bool:
252 """Whether a filename stem is distinctive enough to match in prose.
254 A stem carrying a digit or a separator ("survey_report", "ARC-00000482")
255 only appears in an answer when the model names the document; a bare word
256 stem ("notes") collides with ordinary prose and cannot be trusted.
257 """
258 return any(c.isdigit() or c in "_-" for c in stem)
261def cited_subset(answer: str, sources: list[SearchChunk]) -> list[SearchChunk]:
262 """The sources the answer actually referenced, in order (empty if none).
264 ``[n]`` markers are the primary signal, and ``n`` indexes the unique sources
265 (matching ``build_context``/``format_sources_block``). Name mentions count
266 too: context blocks show the model each source's name, and models often
267 attribute by name ("according to survey_report.pdf") instead of by marker,
268 which otherwise reads as an ungrounded answer to JSON consumers.
269 """
270 uniq = unique_sources(sources)
271 cited = _extract_cited_indices(answer)
272 picked = {i - 1 for i in cited if 1 <= i <= len(uniq)}
273 lowered = answer.lower()
274 for i, source in enumerate(uniq):
275 if i in picked:
276 continue
277 name = Path(source.source).name.lower()
278 stem = Path(source.source).stem
279 if _mentions(lowered, name) or (
280 _identifier_shaped(stem) and _mentions(lowered, stem.lower())
281 ):
282 picked.add(i)
283 return [uniq[i] for i in sorted(picked)]
286def _mentions(answer_lower: str, needle: str) -> bool:
287 """Whether *answer_lower* names *needle* as a whole token.
289 Plain containment marks a source cited whenever its name embeds in a
290 longer one ("log-1" inside "catalog-10", "notes.md" inside
291 "footnotes.md"), which inflates the grounding signal. Filename
292 characters are what must not abut the match; surrounding punctuation
293 and whitespace still count as a mention.
294 """
295 return re.search(rf"(?<![\w.-]){re.escape(needle)}(?![\w-])", answer_lower) is not None
298def _stream_safe_prefix(text: str) -> str:
299 """The prefix of *text* safe to show: a model-authored citation block
300 (heading plus list) is removed, and a bare trailing citation heading is
301 withheld until whatever follows disambiguates it.
303 Kept rstrip-free so the streaming filter can track emitted length exactly;
304 ``strip_llm_citations`` layers the final rstrip on top for one-shot answers.
305 """
306 cleaned = _LLM_CITATION_BLOCK_RE.sub("", text)
307 return _TRAILING_HEADING_RE.sub("", cleaned)
310def strip_llm_citations(text: str) -> str:
311 """Remove an LLM-generated trailing citation block (or dangling citation
312 heading) from answer text. Prose that merely mentions such a heading stays."""
313 return _stream_safe_prefix(text).rstrip()
316class StreamingCitationFilter:
317 """Suppress a model-generated ``Sources:``/``References:`` citation block as
318 it streams, so only lilbee's authoritative source list reaches the reader.
320 A one-shot answer can be stripped after the fact, but streamed tokens are
321 already on screen by the time the block is recognizable. This feeds the
322 answer in incrementally and only releases text once it is certain not to be
323 the start of a citation block: everything up to the last newline is safe to
324 show, the final (possibly partial) line is held back until more text
325 arrives, and a bare citation heading is held until the next line shows
326 whether a list (a citation block, dropped) or prose (a legitimate mention,
327 shown) follows. A heading left dangling when the stream ends is dropped.
328 """
330 def __init__(self) -> None:
331 self._buffer = ""
332 self._emitted = 0
334 def feed(self, text: str) -> str:
335 """Accept the next answer chunk; return the portion safe to show now."""
336 self._buffer += text
337 stripped = _stream_safe_prefix(self._buffer)
338 cut = stripped.rfind("\n")
339 committed = stripped if cut == -1 else stripped[:cut]
340 if len(committed) > self._emitted:
341 out = committed[self._emitted :]
342 self._emitted = len(committed)
343 return out
344 return ""
346 def flush(self) -> str:
347 """Release any remaining safe text once the stream has ended."""
348 final = _stream_safe_prefix(self._buffer)
349 if len(final) > self._emitted:
350 out = final[self._emitted :]
351 self._emitted = len(final)
352 return out
353 return ""
355 @property
356 def answer(self) -> str:
357 """The full answer shown so far, with any citation block removed."""
358 return strip_llm_citations(self._buffer)