Coverage for src/lilbee/retrieval/query/formatting.py: 100%
211 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-11-06 08:17 +0000
1"""Source formatting, context templating, and LLM citation extraction."""
3from __future__ import annotations
5import re
6import string
7from collections.abc import Callable, Mapping, Sequence
8from dataclasses import dataclass
9from functools import cache
10from pathlib import Path
12from markdown_it import MarkdownIt
13from markdown_it.token import Token
15from lilbee.core.text import collapse_whitespace, strip_control_characters
16from lilbee.core.web_names import INDEX_STEM, WEB_PREFIX, is_web_source, page_stem
17from lilbee.data.store import SearchChunk, is_memory_source
18from lilbee.data.types import CODE_CONTENT_TYPE
20CONTEXT_TEMPLATE = """Context:
21{context}
23Question: {question}"""
26# Bracketed citation groups: [1], [1, 2], [1-3], [1, 3-5]. Models mix all of
27# these despite being asked for single [n] markers; matching only [n] made
28# cited_sources under-count and fed JSON consumers false-negative grounding.
29_CITE_GROUP_RE = re.compile(r"\[(\d+(?:\s*[-,]\s*\d+)*)\]")
30_CITE_RANGE_RE = re.compile(r"(\d+)\s*-\s*(\d+)")
31# Ranges wider than this are page spans or line numbers, not citation lists.
32_MAX_CITE_RANGE = 32
34# Heading shapes that open a model-authored citation block: anchored to the
35# start of the text as well as a preceding newline (an answer that is nothing
36# but a fabricated block still strips), optionally markdown-decorated (ATX
37# hashes, emphasis around the word or the colon: "**Sources:**", "*References:*").
38_HEADING_WORDS = r"(?:(?:Key\s+)?Sources|References|Bibliography|Citations)"
39_CITE_HEADING = (
40 r"(?:\n{1,3}|\A)[ \t]*(?:#+\s*)?[*_]{0,3}" + _HEADING_WORDS + r"[*_]{0,3}\s*:?\s*[*_]{0,3}"
41)
42# One list line: a bullet, arrow, "[1]" or "1." marker and the rest of its line.
43_CITE_LIST_LINE = r"[ \t]*(?:[-*•→\[]|\d+[.)])[^\n]*"
45# A heading line followed by a list. The list is required so prose discussing
46# such a heading is not clipped; the match ends with the list, not end-of-text,
47# so an answer resuming after its citations keeps the continuation. Items may be
48# blank-line separated, which markdown does routinely; stopping at the first
49# blank line would leave the rest of a fabricated list in the answer.
50_LLM_CITATION_BLOCK_RE = re.compile(
51 _CITE_HEADING + r"\n\s*" + _CITE_LIST_LINE + r"(?:\n+" + _CITE_LIST_LINE + r")*",
52 re.IGNORECASE,
53)
55# A citation-style heading at the very end of the text, nothing after it yet.
56# Mid-stream this is ambiguous (the next line decides list vs prose), so it is
57# held back rather than shown; at end of stream it is a dangling artifact of a
58# citation block the model never finished, and is dropped either way.
59_TRAILING_HEADING_RE = re.compile(_CITE_HEADING + r"\s*$", re.IGNORECASE)
62_SOURCE_SEP = " · "
63# A page controls its title, so a rendered label never opens a link, image, code
64# span, tag, entity, emphasis or strikethrough of its own, and never starts a block.
65_MARKDOWN_SPECIAL = "\\[]`<&*_~"
66_MARKDOWN_SPECIAL_RE = re.compile(f"[{re.escape(_MARKDOWN_SPECIAL)}]")
67_BLOCK_MARK_RE = re.compile(
68 f"^[{re.escape(''.join(c for c in string.punctuation if c not in _MARKDOWN_SPECIAL))}]"
69)
70_LIST_NUMBER_RE = re.compile(r"^(\d+)([.)])(?=\s|$)")
73def source_label(source: str, title: str = "") -> str:
74 """The plain one-line name of a source: ``host · title`` for a crawled page with a
75 stored *title*, else ``host · slug``; a local file's documents-dir-relative path."""
76 if not is_web_source(source):
77 return collapse_whitespace(source)
78 *folders, file_name = source.removeprefix(WEB_PREFIX).split("/")
79 stem = page_stem(file_name)
80 segments = folders if stem == INDEX_STEM else [*folders, stem]
81 if not segments:
82 return collapse_whitespace(source)
83 host = segments[0].removeprefix("www.")
84 slug = collapse_whitespace(title) or segments[-1]
85 return host if slug == segments[0] else f"{host}{_SOURCE_SEP}{slug}"
88def markdown_text(text: str) -> str:
89 """*text* on one line as markdown that shows exactly that text, with no markup, block
90 or control character."""
91 plain = strip_control_characters(collapse_whitespace(text))
92 escaped = _MARKDOWN_SPECIAL_RE.sub(r"\\\g<0>", plain)
93 escaped = _BLOCK_MARK_RE.sub(r"\\\g<0>", escaped)
94 return _LIST_NUMBER_RE.sub(r"\1\\\2", escaped)
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 _location_suffix(result: SearchChunk) -> str:
109 """The page or line span of a chunk, or empty when neither applies.
111 Zero means "no location": PDF chunks whose page metadata was missing are
112 stored with page 0, so a locator is only rendered when at least one end
113 of the span is set. Sole owner of the rule: a Sources line shows the
114 ``SourceRef.locator`` it computes.
115 """
116 if result.content_type == "pdf" and (result.page_start or result.page_end):
117 ps, pe = result.page_start, result.page_end
118 return f"page {ps}" if ps == pe else f"pages {ps}-{pe}"
119 if result.content_type == CODE_CONTENT_TYPE and (result.line_start or result.line_end):
120 ls, le = result.line_start, result.line_end
121 return f"line {ls}" if ls == le else f"lines {ls}-{le}"
122 return ""
125# Every locator ``_location_suffix`` writes. A client saves a ref's locator back, and
126# only these forms are accepted.
127_LOCATOR_RE = re.compile(r"(?:page|line) \d+|(?:pages|lines) \d+-\d+")
130def is_locator(text: str) -> bool:
131 """Whether *text* is a locator lilbee writes, or '' for none."""
132 return not text or _LOCATOR_RE.fullmatch(text) is not None
135def source_ref_link(ref: SourceRef) -> str:
136 """*ref* as the ``[label](file-url)`` link of a Sources line, the label as plain text;
137 the plain label when no path resolves."""
138 return _linked_label(ref.source, ref.label)
141def source_ref_line(ref: SourceRef) -> str:
142 """*ref* as a Sources line shows it: its link and its page or line locator."""
143 return located(source_ref_link(ref), ref)
146def markdown_source_line(ref: SourceRef) -> str:
147 """*ref* as a Markdown renderer shows its Sources line: its label and locator as
148 written, linked to the file of its own source."""
149 line = _linked_label(ref.source, markdown_text(ref.label))
150 return f"{line}, {markdown_text(ref.locator)}" if ref.locator else line
153def located(name: str, ref: SourceRef) -> str:
154 """*name*, the way a Sources line names *ref*, followed by its locator when it has one."""
155 return f"{name}, {ref.locator}" if ref.locator else name
158def _linked_label(source: str, label: str) -> str:
159 """*label* linked to the file of *source*, or plain when *source* has no file."""
160 if not source or is_memory_source(source):
161 return label
162 url = _source_file_url(source)
163 return f"[{label}]({url})" if url else label
166def unique_sources(results: list[SearchChunk]) -> list[SearchChunk]:
167 """The first chunk of each distinct source, in retrieval order. A source's
168 1-based position here is the citation number the model emits (see
169 ``build_context``) and the number shown in the Sources block, so the answer's
170 ``[n]`` markers and the source list always agree."""
171 seen: set[str] = set()
172 out: list[SearchChunk] = []
173 for r in results:
174 if r.source not in seen:
175 seen.add(r.source)
176 out.append(r)
177 return out
180@dataclass(frozen=True)
181class SourceRef:
182 """One source as a client names it: its ``label``, the ``[n]`` an answer cites it
183 by (None when no answer numbers it) and the plain-text ``locator`` of its page or
184 line span, such as ``page 3`` ('' when it has none)."""
186 source: str
187 label: str
188 citation_number: int | None
189 locator: str = ""
192def _chunk_ref(result: SearchChunk, title: str, number: int | None) -> SourceRef:
193 """The ref of *result*'s source, located by *result*. A memory row has no title or locator."""
194 if result.memory_id is not None:
195 return SourceRef(result.source, source_label(result.source), number)
196 return SourceRef(
197 result.source, source_label(result.source, title), number, _location_suffix(result)
198 )
201def source_refs(results: list[SearchChunk], titles: Mapping[str, str]) -> list[SourceRef]:
202 """Each unique source of *results*, with its label from its title in *titles*, its
203 locator from its first chunk and its citation number: its position in
204 ``unique_sources``, the ``[n]`` the model cites."""
205 return [
206 _chunk_ref(r, titles.get(r.source, ""), i) for i, r in enumerate(unique_sources(results), 1)
207 ]
210def _context_header(result: SearchChunk) -> str:
211 """One-line provenance for a context block: source name plus location.
213 Without it the answering model sees bare numbered text: it cannot
214 attribute a claim to a named document, notice two chunks share a source,
215 or confirm it is reading the document the user asked about.
216 """
217 location = _location_suffix(result)
218 if location:
219 return f"{result.source}, {location}"
220 return result.source
223def build_context(results: list[SearchChunk]) -> str:
224 """Number each passage by its source file, not its position, so citation
225 numbers are stable while streaming and map 1:1 to the Sources block. Passages
226 from the same file share a number.
228 Each block carries a provenance header: without it the answering model sees
229 bare numbered text and cannot attribute a claim to a named document, notice
230 two passages share a source, or confirm it is reading the document asked for.
231 """
232 order: dict[str, int] = {}
233 for r in results:
234 order.setdefault(r.source, len(order) + 1)
235 return "\n\n".join(f"[{order[r.source]}] ({_context_header(r)})\n{r.chunk}" for r in results)
238# Grepped by consumers to know an answer carries its own Sources list (the
239# no-results toast; the pill row, which must not stack a second list).
240SOURCES_BLOCK_MARKER = "\n\nSources:\n"
241# A Sources line's ``[label](file-url)`` link. One source per line, so the label
242# runs to the line's last ``](file://``: a label holding that text cannot end it.
243FILE_LINK_RE = re.compile(r"\[(?P<label>[^\n]+)\]\((?P<url>file://[^)\s]+)\)")
244# Parsed after a text; if it lands inside a fence, the text left that fence open.
245_FENCE_PROBE = "\n\n# probe"
248def format_sources_block(results: list[SearchChunk], titles: Mapping[str, str]) -> str:
249 """The authoritative numbered ``Sources:`` block appended to an answer. Each
250 unique source is numbered to match the ``[n]`` markers the model emitted, so
251 every inline citation resolves to a line here; *titles* holds the stored title
252 of each crawled page. '' when there are no sources."""
253 refs = source_refs(results, titles)
254 return _sources_list(refs, source_ref_line) if refs else ""
257@dataclass(frozen=True)
258class AnswerSources:
259 """Emitted once after a grounded answer's last token: the retrieved *results* the
260 answer was written from and the stored *titles* of their crawled pages."""
262 results: list[SearchChunk]
263 titles: Mapping[str, str]
265 def refs(self) -> list[SourceRef]:
266 """The answer's Sources list, numbered as the answer cites it."""
267 return source_refs(self.results, self.titles)
269 def block(self, answer: str) -> str:
270 """The ``Sources:`` block that follows *answer*, closing a code block it left open."""
271 refs = self.refs()
272 if not refs:
273 return ""
274 return close_open_fence(answer).removeprefix(answer) + _sources_list(refs, source_ref_line)
277def append_sources_block(
278 content: str,
279 refs: Sequence[SourceRef],
280 render: Callable[[SourceRef], str] = source_ref_line,
281) -> str:
282 """*content* ending in the numbered ``Sources:`` list of *refs*, each shown by
283 *render*, whatever the content says. A code block the content leaves open is
284 closed first, so the list is never inside it."""
285 return close_open_fence(content.rstrip()) + _sources_list(refs, render)
288def _sources_list(refs: Sequence[SourceRef], render: Callable[[SourceRef], str]) -> str:
289 """The ``Sources:`` list of *refs*, each shown by *render* under its citation number
290 (its position when it has none). The one renderer of the live and restored lists."""
291 numbers = [i if r.citation_number is None else r.citation_number for i, r in enumerate(refs, 1)]
292 if numbers == list(range(1, len(refs) + 1)):
293 # A markdown ordered list, so a Markdown renderer puts each source on its own
294 # line and the list number is the citation number the model cited inline.
295 lines = "\n".join(f"{n}. {render(r)}" for n, r in zip(numbers, refs, strict=True))
296 else:
297 # A markdown list shows its items numbered on from the first, so a cited
298 # subset [1], [3] would show 1. and 2.: escaped numbers on hard-broken
299 # lines show each number as written.
300 lines = "\\\n".join(f"{n}\\. {render(r)}" for n, r in zip(numbers, refs, strict=True))
301 return SOURCES_BLOCK_MARKER + "\n" + lines
304def split_sources_list(content: str) -> tuple[str, str]:
305 """*content* split before its last Sources list, unless that list is quoted in a code block.
307 A list after a code block the answer never closed (a cut-off answer) is lilbee's;
308 a list inside a code block that closes after it is pasted text.
309 """
310 text, marker, stored_sources = content.rpartition(SOURCES_BLOCK_MARKER)
311 if not marker:
312 return content, ""
313 fence = open_code_fence(text)
314 if fence is None or _never_closed(fence, content):
315 return text, marker + stored_sources
316 return content, ""
319def _never_closed(fence: Token, content: str) -> bool:
320 """Whether *fence*, opened in a prefix of *content*, is still open at its end."""
321 end = open_code_fence(content)
322 return (
323 end is not None
324 and end.map is not None
325 and fence.map is not None
326 and end.map[0] == fence.map[0]
327 )
330def terminal_sources_text(sources: str) -> str:
331 """The Sources list *sources* with every control character and escape sequence removed.
333 The one place a stored title or source label is made safe for a terminal. A page controls
334 its title, and a terminal reads an escape sequence in it as a command. Line breaks stay.
335 """
336 return "\n".join(strip_control_characters(line) for line in sources.split("\n"))
339def terminal_sources_block(content: str) -> str:
340 """*content* with the Sources list it ends in made safe for a terminal or a TUI widget."""
341 text, sources = split_sources_list(content)
342 return text + terminal_sources_text(sources)
345@cache
346def commonmark() -> MarkdownIt:
347 """The shared CommonMark parser."""
348 return MarkdownIt("commonmark")
351def open_code_fence(text: str) -> Token | None:
352 """The top-level code fence *text* leaves open, or None when every fence is closed."""
353 last = commonmark().parse(text + _FENCE_PROBE)[-1]
354 return last if last.type == "fence" else None
357def close_open_fence(text: str) -> str:
358 """*text* with a code fence it leaves open closed, so it cannot swallow what follows.
360 The parser decides: when a heading placed after *text* ends up inside a fence,
361 that fence is closed. A fence in a list item never does, because the heading
362 ends the item, so only a top-level fence is closed, and at column 0.
363 """
364 fence = open_code_fence(text)
365 return f"{text}\n{fence.markup}" if fence else text
368def _extract_cited_indices(text: str) -> set[int]:
369 """Extract citation references from LLM answer text: [1], [1, 2], [1-3]."""
370 indices: set[int] = set()
371 for m in _CITE_GROUP_RE.finditer(text):
372 group = m.group(1)
373 remainder = _CITE_RANGE_RE.sub("", group)
374 for start, end in _CITE_RANGE_RE.findall(group):
375 lo, hi = int(start), int(end)
376 if lo <= hi <= lo + _MAX_CITE_RANGE:
377 indices.update(range(lo, hi + 1))
378 indices.update(int(n) for n in re.findall(r"\d+", remainder))
379 return indices
382def _identifier_shaped(stem: str) -> bool:
383 """Whether a filename stem is distinctive enough to match in prose.
385 A stem carrying a digit or a separator ("survey_report", "ARC-00000482")
386 only appears in an answer when the model names the document; a bare word
387 stem ("notes") collides with ordinary prose and cannot be trusted.
388 """
389 return any(c.isdigit() or c in "_-" for c in stem)
392def cited_subset(answer: str, sources: list[SearchChunk]) -> list[SearchChunk]:
393 """The sources the answer actually referenced, in order (empty if none).
395 ``[n]`` markers are the primary signal, and ``n`` indexes the unique sources
396 (matching ``build_context``/``format_sources_block``). Name mentions count
397 too: context blocks show the model each source's name, and models often
398 attribute by name ("according to survey_report.pdf") instead of by marker,
399 which otherwise reads as an ungrounded answer to JSON consumers.
400 """
401 uniq = unique_sources(sources)
402 cited = _extract_cited_indices(answer)
403 picked = {i - 1 for i in cited if 1 <= i <= len(uniq)}
404 lowered = answer.lower()
405 for i, source in enumerate(uniq):
406 if i in picked:
407 continue
408 name = Path(source.source).name.lower()
409 stem = Path(source.source).stem
410 if _mentions(lowered, name) or (
411 _identifier_shaped(stem) and _mentions(lowered, stem.lower())
412 ):
413 picked.add(i)
414 return [uniq[i] for i in sorted(picked)]
417def _mentions(answer_lower: str, needle: str) -> bool:
418 """Whether *answer_lower* names *needle* as a whole token.
420 Plain containment marks a source cited whenever its name embeds in a
421 longer one ("log-1" inside "catalog-10", "notes.md" inside
422 "footnotes.md"), which inflates the grounding signal. Filename
423 characters are what must not abut the match; surrounding punctuation
424 and whitespace still count as a mention.
425 """
426 return re.search(rf"(?<![\w.-]){re.escape(needle)}(?![\w-])", answer_lower) is not None
429def _stream_safe_prefix(text: str) -> str:
430 """The prefix of *text* safe to show: a model-authored citation block
431 (heading plus list) is removed, and a bare trailing citation heading is
432 withheld until whatever follows disambiguates it.
434 Kept rstrip-free so the streaming filter can track emitted length exactly;
435 ``strip_llm_citations`` layers the final rstrip on top for one-shot answers.
436 """
437 cleaned = _LLM_CITATION_BLOCK_RE.sub("", text)
438 return _TRAILING_HEADING_RE.sub("", cleaned)
441def strip_llm_citations(text: str) -> str:
442 """Remove an LLM-generated trailing citation block (or dangling citation
443 heading) from answer text. Prose that merely mentions such a heading stays."""
444 return _stream_safe_prefix(text).rstrip()
447class StreamingCitationFilter:
448 """Suppress a model-generated ``Sources:``/``References:`` citation block as
449 it streams, so only lilbee's authoritative source list reaches the reader.
451 A one-shot answer can be stripped after the fact, but streamed tokens are
452 already on screen by the time the block is recognizable. This feeds the
453 answer in incrementally and only releases text once it is certain not to be
454 the start of a citation block: everything up to the last newline is safe to
455 show, the final (possibly partial) line is held back until more text
456 arrives, and a bare citation heading is held until the next line shows
457 whether a list (a citation block, dropped) or prose (a legitimate mention,
458 shown) follows. A heading left dangling when the stream ends is dropped.
459 """
461 def __init__(self) -> None:
462 self._buffer = ""
463 self._emitted = 0
465 def feed(self, text: str) -> str:
466 """Accept the next answer chunk; return the portion safe to show now."""
467 self._buffer += text
468 stripped = _stream_safe_prefix(self._buffer)
469 cut = stripped.rfind("\n")
470 committed = stripped if cut == -1 else stripped[:cut]
471 if len(committed) > self._emitted:
472 out = committed[self._emitted :]
473 self._emitted = len(committed)
474 return out
475 return ""
477 def flush(self) -> str:
478 """Release any remaining safe text once the stream has ended."""
479 final = _stream_safe_prefix(self._buffer)
480 if len(final) > self._emitted:
481 out = final[self._emitted :]
482 self._emitted = len(final)
483 return out
484 return ""
486 @property
487 def answer(self) -> str:
488 """The full answer shown so far, with any citation block removed."""
489 return strip_llm_citations(self._buffer)