Coverage for src/lilbee/retrieval/query/searcher.py: 100%
649 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-14 11:46 +0000
1"""RAG search pipeline -- embed, search, expand, rerank, generate."""
3from __future__ import annotations
5import logging
6import re
7from collections.abc import Callable, Generator
8from datetime import datetime
9from enum import StrEnum
10from typing import TYPE_CHECKING, Any, NamedTuple, cast
12from pydantic import BaseModel, Field
13from typing_extensions import TypedDict
15from lilbee.core.config import Config
16from lilbee.core.config.enums import ChatMode
17from lilbee.core.llm_json import json_reply_format
18from lilbee.core.vectors import Vector
19from lilbee.data.store import (
20 ChunkType,
21 MemoryKind,
22 MemoryRow,
23 SearchChunk,
24 Store,
25 cosine_sim,
26 human_recall_predicate,
27)
28from lilbee.providers.base import (
29 LLMProvider,
30 ProviderError,
31 ProviderErrorKind,
32 estimate_budget_tokens,
33 prompt_token_budget,
34)
35from lilbee.retrieval.embedder import Embedder
36from lilbee.retrieval.language import noun_variants, query_language
37from lilbee.retrieval.query.compaction import (
38 COMPACT_PROMPT,
39 CompactionResult,
40 merge_notes,
41 plan_compaction,
42 summary_cap,
43 summary_word_budget,
44)
45from lilbee.retrieval.query.dedup import (
46 _greedy_cover,
47 _relevance_weight,
48 filter_results,
49 order_by_fusion,
50 prepare_results,
51)
52from lilbee.retrieval.query.expansion import (
53 CONDENSE_HISTORY_TURNS,
54 CONDENSE_MAX_TOKENS,
55 CONDENSE_PROMPT,
56 EXPANSION_MAX_TOKENS,
57 EXPANSION_PROMPT,
58 HYDE_MAX_TOKENS,
59)
60from lilbee.retrieval.query.formatting import (
61 CONTEXT_TEMPLATE,
62 StreamingCitationFilter,
63 build_context,
64 cited_subset,
65 format_sources_block,
66 strip_llm_citations,
67)
68from lilbee.retrieval.query.history_window import estimate_text_tokens
69from lilbee.retrieval.query.intent import (
70 INTENT_CLASSIFY_MAX_TOKENS,
71 INTENT_CLASSIFY_PROMPT,
72 AggregateKind,
73 AggregateQuery,
74 document_references,
75 matches_reference,
76 matches_stored_title,
77 matches_title,
78 parse_aggregate,
79 parse_llm_aggregate,
80 title_candidates,
81)
82from lilbee.retrieval.query.memory import format_memory_block
83from lilbee.retrieval.query.neighbors import expand_neighbors
84from lilbee.retrieval.query.structural import is_structural_chunk
85from lilbee.retrieval.query.tokenize import _idf_weights, _tokenize
86from lilbee.retrieval.reasoning import (
87 StreamToken,
88 cap_events_as_stream_tokens,
89 effective_reasoning_cap,
90 split_reasoning,
91 stream_chat_with_cap,
92 strip_reasoning,
93)
95if TYPE_CHECKING:
96 from lilbee.retrieval.concepts import ConceptGraph
97 from lilbee.retrieval.reranker import Reranker
99log = logging.getLogger(__name__)
101# BM25 probe needs at least this many hits to compare top vs. runner-up
102# scores for the expansion-skip heuristic.
103_MIN_BM25_PROBE_RESULTS = 2
105# Substring candidates fetched per document reference before token-exact
106# disambiguation picks the unique winner.
107_KNOWN_ITEM_CANDIDATES = 50
109# Content-based known-item resolution: BM25 hits probed for a reference, and
110# the fraction of them one source must own to count as that document. A
111# docket-style number lives in a document's text, not its filename, so
112# filename matching alone can never resolve it; concentration keeps the
113# fallback conservative, since a number cited across many filings spreads.
114_KNOWN_ITEM_PROBE_K = 6
115_KNOWN_ITEM_PROBE_MAJORITY = 0.75
118# Structured-query mode names (the ``mode:`` prefix shortcut). Single source for
119# both the prefix parser and the dispatch in ``_search_structured``. "term"/"vec"/
120# "hyde" pick a retrieval strategy; "wiki"/"raw" are ChunkType scope shortcuts.
121class QueryMode(StrEnum):
122 """Structured-query prefixes: ``term:``, ``vec:``, ``hyde:``, ``wiki:``, ``raw:``."""
124 TERM = "term"
125 VEC = "vec"
126 HYDE = "hyde"
127 WIKI = ChunkType.WIKI.value
128 RAW = ChunkType.RAW.value
131# Leading list markers models prepend to expansion output despite the prompt:
132# "1.", "2)", "-", "*", "•".
133_LIST_MARKER_RE = re.compile(r"^\s*(?:\d+[.)]\s*|[-*•]\s+)")
136def _strip_list_marker(line: str) -> str:
137 """Drop a leading list marker from an expansion variant line."""
138 return _LIST_MARKER_RE.sub("", line).strip()
141# Half-saturation constant for BM25 confidence: a raw score of 5 reads as 0.5,
142# 20 as 0.8, 45 as 0.9. A plain sigmoid saturated at ~0.99 for any raw score
143# above 5, which made the top-vs-runner-up gap condition unsatisfiable exactly
144# when BM25 was most certain, so the skip never fired.
145_BM25_HALF_SATURATION = 5.0
148def _bm25_confidence(score: float | None) -> float:
149 """Squash a raw, unbounded BM25 score into (0, 1) without saturating.
151 ``s / (s + k)`` keeps strong scores distinguishable (unlike a sigmoid,
152 which flattens everything past ~5 to within 0.01 of 1.0). Absent or
153 non-positive scores read as 0, so a missing FTS signal never trips the
154 expansion skip.
155 """
156 if score is None or score <= 0.0:
157 return 0.0
158 return score / (score + _BM25_HALF_SATURATION)
161def _noun_names_type(noun: str, type_name: str) -> bool:
162 """Whether the question's noun IS the type (modulo case/space/plural),
163 as opposed to reaching it through a synonym."""
164 named = noun_variants(type_name) | noun_variants(type_name.replace("_", " "))
165 return bool(noun_variants(noun) & named)
168# RAG mode answer when retrieval finds no usable sources: a grounded refusal
169# instead of free-wheeling on the model's parametric knowledge. Users who want
170# off-corpus answers can switch to chat mode.
171GROUNDED_REFUSAL = "I couldn't find anything in the indexed documents that answers that."
173# Ask/search answer when the library holds nothing yet. Distinct from the
174# grounded refusal, which implies a search ran and came up empty: here there is
175# nothing to search, so point the user at adding content. Shared across TUI,
176# CLI, HTTP, and MCP, so the phrasing stays surface-neutral (no slash commands).
177EMPTY_LIBRARY = (
178 "Your library is empty, so there's nothing to search yet. "
179 "Add documents to your library first, then ask again."
180)
182# Ask/search needs an embedder to ground an answer. When none is loaded, refuse
183# with an actionable message rather than hard-failing or silently answering
184# ungrounded; chat mode stays available for an off-corpus reply.
185SEARCH_NEEDS_EMBEDDER = (
186 "Search needs an embedding model to ground answers in your documents. "
187 "Add one, or switch to chat mode for an ungrounded reply."
188)
190# Association answers list at most this many groups before summarizing.
191_ASSOCIATION_LINES = 15
192# One retry after a provider context-overflow, refitting to this fraction of
193# the budget. The estimator is a heuristic; overflow must degrade, not fail.
194_OVERFLOW_RETRY_SCALE = 0.6
195# Approximate token cost of the Context/Question template wrapper.
196_CONTEXT_TEMPLATE_TOKENS = 16
197# Approximate per-source overhead: the "[i] " marker, the provenance header
198# (source path plus page/line span), and the blank-line separator.
199_PER_SOURCE_TOKENS = 24
202class ChatMessage(TypedDict):
203 """A single chat message with role and content."""
205 role: str
206 content: str
209class AskResult(BaseModel):
210 """Structured result from ask_raw: answer text, retrieved sources, and the cited subset.
212 ``sources`` is the full retrieved/reranked set; ``cited_sources`` is the subset the
213 answer actually referenced via [n] markers (empty when it cited nothing), so a JSON
214 consumer can tell whether the answer was grounded without re-parsing the text.
215 """
217 answer: str
218 sources: list[SearchChunk]
219 cited_sources: list[SearchChunk] = Field(default_factory=list)
222class StructuredQuery(NamedTuple):
223 """A mode-prefixed query split from its mode (see ``QueryMode``)."""
225 mode: QueryMode | None
226 query: str
229class RagContext(NamedTuple):
230 """Grounded context for one turn: the chunks and the prompt built on them.
232 ``base_results`` is the pre-widen selected set. An overflow retry refits from
233 it, not from ``results`` (whose neighbor text is baked in and can no longer
234 be shed), so a tighter fit drops expansion before it drops an original chunk.
235 """
237 results: list[SearchChunk]
238 messages: list[ChatMessage]
239 base_results: list[SearchChunk] | None = None
242class Searcher:
243 """RAG search pipeline -- embed, search, expand, rerank, generate.
244 All search and answer operations go through this class.
245 Constructed with injected dependencies via the Services container.
246 """
248 def __init__(
249 self,
250 config: Config,
251 provider: LLMProvider,
252 store: Store,
253 embedder: Embedder,
254 reranker: Reranker,
255 concepts: ConceptGraph,
256 ) -> None:
257 self._config = config
258 self._provider = provider
259 self._store = store
260 self._embedder = embedder
261 self._reranker = reranker
262 self._concepts = concepts
264 def _apply_temporal_filter(
265 self, results: list[SearchChunk], question: str
266 ) -> list[SearchChunk]:
267 if not self._config.temporal_filtering:
268 return results
269 from lilbee.runtime.temporal import detect_temporal, resolve_date_range
271 keyword = detect_temporal(question)
272 if keyword is None:
273 return results
274 date_range = resolve_date_range(keyword)
275 source_dates = self._store.source_ingested_at_map()
276 filtered: list[SearchChunk] = []
277 for r in results:
278 ingested_at = source_dates.get(r.source, "")
279 if not ingested_at:
280 filtered.append(r)
281 continue
282 try:
283 doc_date = datetime.fromisoformat(ingested_at)
284 if date_range.start <= doc_date <= date_range.end:
285 filtered.append(r)
286 except (ValueError, TypeError):
287 filtered.append(r)
288 return filtered if filtered else results
290 def _apply_guardrails(
291 self,
292 variants: list[tuple[str, Vector]],
293 question_vec: Vector,
294 ) -> list[tuple[str, Vector]]:
295 """Drop expansion variants whose embedding drifts too far from the question."""
296 if not self._config.expansion_guardrails:
297 return variants
298 threshold = self._config.expansion_similarity_threshold
299 return [(text, vec) for text, vec in variants if cosine_sim(question_vec, vec) >= threshold]
301 def _concept_query_expansion(self, question: str) -> list[str]:
302 if not self._config.concept_graph:
303 return []
304 try:
305 if not self._concepts.get_graph():
306 return []
307 return self._concepts.expand_query(question)
308 except Exception:
309 log.debug("Concept query expansion failed", exc_info=True)
310 return []
312 def _llm_expand(self, question: str, count: int) -> list[str]:
313 """Call the LLM to produce ``count`` alternative phrasings.
315 Reasoning is stripped before the line split: a reasoning chat model
316 otherwise contributes its deliberation as "variants" that get embedded
317 and searched. List numbering is stripped per line since models add it
318 despite the prompt, and "1." pollutes the BM25 arm of every variant
319 search.
320 """
321 prompt = EXPANSION_PROMPT.format(count=count, question=question)
322 messages = [{"role": "user", "content": prompt}]
323 response = self._provider.chat(
324 messages, stream=False, options={"num_predict": EXPANSION_MAX_TOKENS}
325 )
326 text = strip_reasoning(response.text).strip()
327 variants = [_strip_list_marker(line.strip()) for line in text.split("\n") if line.strip()]
328 kept = [v for v in variants if v][:count]
329 log.info("Query expansion produced %d variants", len(kept))
330 return kept
332 def _expand_query(self, question: str, question_vec: Vector) -> list[tuple[str, Vector]]:
333 """Return ``(variant, variant_vec)`` pairs for downstream search.
335 LLM variants run through ``_apply_guardrails``; concept-graph
336 variants bypass it since they come from deterministic traversal.
337 Embeddings batch per source: one provider round-trip per source.
338 """
339 count = self._config.query_expansion_count
340 if count <= 0 and not self._config.concept_graph:
341 return []
342 # Short queries skip LLM expansion: BM25/vector signal is already strong
343 # and the LLM round-trip dominates latency on small local models.
344 # Concept-graph expansion still runs.
345 short_threshold = self._config.expansion_short_query_tokens
346 skip_llm = short_threshold > 0 and len(_tokenize(question)) <= short_threshold
347 try:
348 llm_variants: list[tuple[str, Vector]] = []
349 if count > 0 and not skip_llm:
350 llm_texts = list(self._llm_expand(question, count))
351 if llm_texts:
352 llm_vectors = self._embedder.embed_query_batch(llm_texts)
353 llm_variants = list(zip(llm_texts, llm_vectors, strict=True))
354 llm_variants = self._apply_guardrails(llm_variants, question_vec)
356 concept_texts = list(self._concept_query_expansion(question))
357 if concept_texts:
358 concept_vectors = self._embedder.embed_query_batch(concept_texts)
359 llm_variants.extend(zip(concept_texts, concept_vectors, strict=True))
361 return llm_variants
362 except Exception as exc:
363 log.warning("Query expansion disabled for this call: %s", exc)
364 log.debug("Query expansion exception", exc_info=True)
365 return []
367 def _should_skip_expansion(self, question: str, chunk_type: ChunkType | None = None) -> bool:
368 if self._config.expansion_skip_threshold <= 0:
369 return False
370 # Probe the same pool the scoped search returns, else a confident hit in
371 # the wrong sub-pool could skip expansion the scoped result actually needs.
372 results = self._store.bm25_probe(
373 question, top_k=_MIN_BM25_PROBE_RESULTS, chunk_type=chunk_type
374 )
375 if not results:
376 return False
377 top_raw = results[0].bm25_score or 0.0
378 if _bm25_confidence(top_raw) < self._config.expansion_skip_threshold:
379 return False
380 if len(results) < _MIN_BM25_PROBE_RESULTS:
381 return True
382 # Relative gap in raw score space: any squash compresses the spread
383 # between two strong scores toward zero, so a squashed-space gap test
384 # can never fire exactly when the lexical arm is most certain.
385 second_raw = results[1].bm25_score or 0.0
386 relative_gap = (top_raw - second_raw) / top_raw if top_raw > 0 else 0.0
387 skip = relative_gap >= self._config.expansion_skip_gap
388 if skip:
389 log.info(
390 "Query expansion skipped: BM25 confident (raw %.1f, gap %.0f%%)",
391 top_raw,
392 relative_gap * 100,
393 )
394 return skip
396 def _apply_concept_boost(self, results: list[SearchChunk], question: str) -> list[SearchChunk]:
397 if not self._config.concept_graph or not results:
398 return results
399 try:
400 if not self._concepts.get_graph():
401 return results
402 query_concepts = self._concepts.extract_concepts(question)
403 if not query_concepts:
404 return results
405 boosted = self._concepts.boost_results(results, query_concepts)
406 # boost_results returns copies with the canonical score raised, in
407 # input order; re-sort so the boost actually re-ranks for callers
408 # that consume search() order directly (CLI search, MCP search).
409 return order_by_fusion(boosted)
410 except Exception:
411 log.debug("Concept boost failed", exc_info=True)
412 return results
414 def _hyde_search(
415 self, question: str, top_k: int, chunk_type: ChunkType | None = None
416 ) -> list[SearchChunk]:
417 """Hypothetical Document Embedding search.
418 Gao et al. 2022, "Precise Zero-Shot Dense Retrieval without
419 Relevance Labels" -- generates a hypothetical answer passage,
420 embeds it, and uses the embedding to search for real documents.
422 The passage is deliberately embedded with ``embed_query`` (the query
423 instruction), not the document prefix: it stands in for the user's
424 query against the doc-prefixed index, staying in the same vector
425 space as every other query this searcher issues. Changing that is a
426 retrieval-quality experiment for the embedding bench, not a refactor.
427 """
428 try:
429 response = self._provider.chat(
430 [{"role": "user", "content": self._config.hyde_prompt.format(question=question)}],
431 stream=False,
432 options={"num_predict": HYDE_MAX_TOKENS},
433 )
434 # Reasoning models front-load deliberation; embedding it instead
435 # of the passage would search for the model's thought process.
436 text = strip_reasoning(response.text).strip()
437 if not text:
438 return []
439 hyde_vec = self._embedder.embed_query(text)
440 return self._store.search(hyde_vec, top_k=top_k, query_text=None, chunk_type=chunk_type)
441 except Exception:
442 log.debug("HyDE search failed", exc_info=True)
443 return []
445 def _refuse_wiki_scope(self, chunk_type: ChunkType | None) -> bool:
446 """Whether a wiki-scoped search must serve nothing because wiki is off.
448 Turning the setting off deletes no rows, so the pages a library was
449 wikified with are still there. Serving them would contradict the
450 setting and widening to the whole pool would answer a question the
451 caller did not ask, so the scope resolves to no results.
452 """
453 if chunk_type != ChunkType.WIKI or self._config.wiki:
454 return False
455 log.warning("wiki scope requested but the wiki is disabled; returning no results")
456 return True
458 def _retrieval_scope(self, chunk_type: ChunkType | None) -> ChunkType | None:
459 """The chunk filter to apply, keeping wiki rows out while wiki is off.
461 An unscoped search narrows to ``RAW``, which covers table chunks, so
462 the only rows it drops are generated wiki pages.
463 """
464 if chunk_type is None and not self._config.wiki:
465 return ChunkType.RAW
466 return chunk_type
468 def _parse_structured_query(self, question: str) -> StructuredQuery:
469 stripped = question.strip()
470 for mode in QueryMode:
471 prefix = f"{mode.value}:"
472 if stripped.lower().startswith(prefix):
473 return StructuredQuery(mode, stripped[len(prefix) :].strip())
474 return StructuredQuery(None, question)
476 def _search_structured(
477 self,
478 mode: QueryMode,
479 query: str,
480 top_k: int,
481 chunk_type: ChunkType | None = None,
482 ) -> list[SearchChunk]:
483 # QueryMode.WIKI / QueryMode.RAW are a chunk-type scope shortcut, and an
484 # explicit ``chunk_type`` arg beats the prefix. Resolving the scope up
485 # front routes every mode through the same wiki-disabled guard, so
486 # ``wiki:`` cannot bypass it.
487 requested = chunk_type
488 if requested is None and mode in (QueryMode.WIKI, QueryMode.RAW):
489 requested = ChunkType(mode.value)
490 if self._refuse_wiki_scope(requested):
491 return []
492 scope = self._retrieval_scope(requested)
493 if mode is QueryMode.TERM:
494 return self._store.bm25_probe(query, top_k=top_k, chunk_type=scope)
495 if mode is QueryMode.VEC:
496 query_vec = self._embedder.embed_query(query)
497 return self._store.search(query_vec, top_k=top_k, query_text=None, chunk_type=scope)
498 if mode is QueryMode.HYDE:
499 return self._hyde_search(query, top_k, chunk_type=scope)
500 query_vec = self._embedder.embed_query(query)
501 return self._store.search(query_vec, top_k=top_k, query_text=query, chunk_type=scope)
503 def select_context(
504 self, results: list[SearchChunk], question: str, max_sources: int | None = None
505 ) -> list[SearchChunk]:
506 """Pick ``max_sources`` chunks.
508 Results carrying ``rerank_score`` keep the cross-encoder order (top
509 ``max_sources``); otherwise greedy IDF-weighted set cover.
510 """
511 if max_sources is None:
512 max_sources = self._config.max_context_sources
513 if len(results) <= max_sources:
514 return results
515 if any(r.rerank_score is not None for r in results):
516 return results[:max_sources]
518 question_terms = set(_tokenize(question))
519 if not question_terms:
520 return results[:max_sources]
522 chunk_tokens = [set(_tokenize(r.chunk)) for r in results]
523 term_weights = _idf_weights(question_terms, chunk_tokens)
524 if not any(term_weights.values()):
525 return results[:max_sources]
527 weights = [_relevance_weight(r) for r in results]
528 selected = _greedy_cover(chunk_tokens, question_terms, term_weights, max_sources, weights)
529 selected.sort()
530 return [results[i] for i in selected]
532 def _merge_variant_results(
533 self,
534 question: str,
535 query_vec: Vector,
536 results: list[SearchChunk],
537 seen: set[tuple[str, int]],
538 top_k: int,
539 chunk_type: ChunkType | None,
540 ) -> None:
541 """Append unseen variant-search hits to ``results`` (in place)."""
542 for variant, variant_vec in self._expand_query(question, query_vec):
543 variant_results = self._store.search(
544 variant_vec,
545 top_k=top_k,
546 query_text=variant,
547 chunk_type=chunk_type,
548 )
549 for r in variant_results:
550 key = (r.source, r.chunk_index)
551 if key not in seen:
552 results.append(r)
553 seen.add(key)
555 def _merge_hyde_results(
556 self,
557 question: str,
558 results: list[SearchChunk],
559 seen: set[tuple[str, int]],
560 top_k: int,
561 chunk_type: ChunkType | None = None,
562 ) -> None:
563 """Append unseen HyDE hits to ``results`` (in place), down-weighted by
564 ``hyde_weight`` in canonical score space (a weight of 1.0 trusts HyDE
565 hits as much as direct hits; lower discounts them proportionally)."""
566 for r in self._hyde_search(question, top_k, chunk_type=chunk_type):
567 key = (r.source, r.chunk_index)
568 if key in seen:
569 continue
570 if r.score is not None:
571 r = r.model_copy(update={"score": r.score * self._config.hyde_weight})
572 results.append(r)
573 seen.add(key)
575 def search(
576 self,
577 question: str,
578 top_k: int = 0,
579 *,
580 chunk_type: ChunkType | None = None,
581 ) -> list[SearchChunk]:
582 """Embed question and search with expansion, HyDE, and concept boost.
583 Returns up to top_k*2 candidates for downstream filtering.
585 When *chunk_type* is set (``"raw"`` or ``"wiki"``), only chunks of
586 that type are returned. An explicit ``chunk_type`` always wins
587 over the ``wiki:``/``raw:`` prefix shortcut in *question* so the
588 user-facing scope choice has the final say.
590 While wiki generation is disabled, pages generated before it was
591 turned off stay out of every result: an unscoped search narrows to
592 document chunks, and a ``"wiki"`` scope returns nothing rather than
593 widening to the pool the caller did not ask for.
595 A ``mode:`` prefix (``term:``/``vec:``/``hyde:``/``wiki:``/``raw:``)
596 forces a single explicit retrieval strategy and so skips expansion and
597 concept boost, but the temporal date-range filter still applies -- it is
598 a filter, not a re-ranking, and a "recent" query must be honored in any
599 mode.
600 """
601 if top_k == 0:
602 top_k = self._config.top_k
603 mode, clean_query = self._parse_structured_query(question)
604 if mode is not None:
605 structured = self._search_structured(mode, clean_query, top_k, chunk_type=chunk_type)
606 return self._apply_temporal_filter(structured, clean_query)
607 if self._refuse_wiki_scope(chunk_type):
608 return []
609 chunk_type = self._retrieval_scope(chunk_type)
610 if self._config.intent_routing:
611 # A query naming one document wants that document on every
612 # retrieval surface, not just ask: without this, bare search
613 # (HTTP /api/search, MCP) returns similarity neighbors of the
614 # question's wording. The document's head, in document order,
615 # fills the standard return budget.
616 known_item = self._known_item_results(question, chunk_type)
617 if known_item:
618 return known_item[: top_k * 2]
619 query_vec = self._embedder.embed_query(question)
620 # Retrieve the reranker's candidate depth when one is loaded, else top_k.
621 retrieve_k = (
622 max(top_k, self._config.rerank_candidates) if self._config.reranker_model else top_k
623 )
624 results = self._store.search(
625 query_vec,
626 top_k=retrieve_k,
627 query_text=question,
628 chunk_type=chunk_type,
629 )
630 # Query expansion (variant + HyDE searches) is skipped for short/term
631 # queries, but concept boost is a separate graph re-rank that should still
632 # apply -- the early return used to drop it on the skip path.
633 if not self._should_skip_expansion(question, chunk_type):
634 seen = {(r.source, r.chunk_index) for r in results}
635 self._merge_variant_results(question, query_vec, results, seen, top_k, chunk_type)
636 if self._config.hyde:
637 self._merge_hyde_results(question, results, seen, top_k, chunk_type)
638 # Merged variant/HyDE hits arrive appended, not ranked; every consumer
639 # of this method (bare search surfaces included) gets one global order
640 # over the canonical score rather than insertion order.
641 results = order_by_fusion(results)
642 # Apply the date-range filter here so the bare search() path (e.g. /api/search)
643 # honors a "recent"/"today" query, matching the chat/ask path.
644 results = self._apply_temporal_filter(results, question)
645 # One relevance cutoff for every surface. The rule lives here (with
646 # the lexical-support exemption) rather than per surface: the CLI,
647 # HTTP, and MCP copies of a bare distance cutoff dropped both-arm
648 # rows the fusion layer deliberately keeps past max_distance.
649 results = filter_results(results, self._config.max_distance)
650 # Drop tables-of-contents and cover pages that only the vector arm
651 # surfaced: they dilute context precision without answering a question.
652 # Filtered from the top_k*2 candidate buffer so enough real passages
653 # remain for the downstream trim. Runs before the concept boost so a
654 # boost cannot promote a structural chunk into the rank-0 exemption.
655 if self._config.filter_structural_chunks:
656 # A lexical (BM25 or title) hit or the top-ranked row is content the
657 # answer may need, whatever its shape, so it is never dropped; only
658 # structural chunks the lexical arms did not support are removed.
659 results = [
660 r
661 for i, r in enumerate(results)
662 if r.bm25_score is not None or i == 0 or not is_structural_chunk(r.chunk)
663 ]
664 results = self._apply_concept_boost(results, question)
665 results = order_by_fusion(results)
666 # Rerank when a cross-encoder is loaded so every search surface (HTTP,
667 # MCP, CLI, ask) gets reranked order, not just the ask/chat path. A
668 # mode: prefix returned earlier and stays unreranked by design.
669 if self._config.reranker_model:
670 results = self._reranker.rerank(question, results)
671 return results[: top_k * 2]
673 def _condense_question(self, question: str, history: list[ChatMessage]) -> str:
674 """Rewrite a follow-up into a standalone retrieval query.
676 Retrieval sees only the query text; without this, "what about his
677 brother?" is embedded and BM25-matched with its pronouns. The
678 rewritten form drives retrieval only; the user's original wording
679 still reaches the answering prompt. Falls back to the original
680 question on any failure or empty rewrite.
681 """
682 recent = history[-CONDENSE_HISTORY_TURNS:]
683 transcript = "\n".join(f"{m['role']}: {m['content']}" for m in recent)
684 prompt = CONDENSE_PROMPT.format(history=transcript, question=question)
685 try:
686 response = self._provider.chat(
687 [{"role": "user", "content": prompt}],
688 stream=False,
689 options={"num_predict": CONDENSE_MAX_TOKENS},
690 )
691 rewritten = strip_reasoning(response.text).strip().splitlines()
692 first_line = rewritten[0].strip() if rewritten else ""
693 if first_line:
694 log.debug("Condensed follow-up %r -> %r", question, first_line)
695 return first_line
696 except Exception:
697 log.debug("History condensation failed; using the raw question", exc_info=True)
698 return question
700 def summarize_history(
701 self,
702 messages: list[ChatMessage],
703 previous_summary: str = "",
704 on_batch: Callable[[int, int], None] | None = None,
705 ) -> CompactionResult:
706 """Condense turns being dropped from the prompt into carry-forward notes.
708 Chat calls this when a conversation outgrows its token budget: without it
709 the oldest turns are dropped outright and the model silently loses a
710 conversation the user can still scroll. *previous_summary* is folded in
711 so summaries compound instead of each one forgetting the last.
713 Each batch is summarized ONCE, independently, and the notes are merged.
714 Feeding each batch the running summary instead would re-summarize the
715 summary once per batch: at a 2k window a long backlog is ~16 batches, so
716 the earliest turns would be a summary of a summary sixteen deep, which a
717 small model degrades into drift long before the budget runs out. Depth
718 stays at one, plus one merge-compression when the notes outgrow the cap.
720 Returns the notes and how many turns they cover; ``stranded`` counts turns
721 dropped without notes, which the caller must surface rather than hide.
722 ``on_batch`` hears ``(batch, total)`` before each model call, for progress UI.
723 """
724 ctx_target = self._config.chat_n_ctx_target
725 plan = plan_compaction(messages, ctx_target=ctx_target)
726 notes: list[str] = []
727 condensed = 0
728 stranded = plan.stranded
729 for index, batch in enumerate(plan.batches):
730 if on_batch is not None:
731 on_batch(index + 1, len(plan.batches))
732 note = self._summarize_batch(batch)
733 if note:
734 notes.append(note)
735 condensed += len(batch)
736 else:
737 # Count what landed, not what was planned: these turns are gone
738 # with nothing standing in for them.
739 stranded += len(batch)
740 merged = merge_notes(previous_summary, notes)
741 cap = summary_cap(ctx_target)
742 if estimate_text_tokens(merged) > cap:
743 merged = self._summarize_batch([{"role": "user", "content": merged}]) or merged
744 return CompactionResult(
745 summary=merged or previous_summary, condensed=condensed, stranded=stranded
746 )
748 def _summarize_batch(self, batch: list[ChatMessage]) -> str:
749 """Fold one batch of dropped turns into notes.
751 Each batch is summarized on its own, with no carried-forward notes in
752 the prompt: summarize_history merges the per-batch notes instead, which
753 keeps summary depth at one rather than re-summarizing the summary once
754 per batch.
756 An overflowing batch splits in half and each half folds on its own:
757 batch sizing is estimate-based, and the cost of an estimate miss here
758 is stranded turns, not a slow call. Depth is log2 of the batch.
760 Returns "" on any other failure, so the caller counts the batch as
761 stranded rather than reporting turns it has no notes for.
762 """
763 transcript = "\n".join(f"{m['role']}: {m['content']}" for m in batch)
764 prompt = COMPACT_PROMPT.format(
765 words=summary_word_budget(self._config.chat_n_ctx_target),
766 transcript=transcript,
767 )
768 try:
769 response = self._provider.chat(
770 [{"role": "user", "content": prompt}],
771 stream=False,
772 options={
773 "num_predict": summary_cap(self._config.chat_n_ctx_target),
774 # A thinking model spends the whole budget in a <think>
775 # block that llama.cpp force-closes and strip_reasoning
776 # deletes whole, leaving "" and stranding the batch.
777 "think": False,
778 # Deterministic: the same conversation folds the same way.
779 "temperature": 0,
780 },
781 )
782 summary = strip_reasoning(response.text).strip()
783 if summary:
784 return summary
785 # Only the native llama-server path honors think=False; elsewhere a
786 # reasoning model can leave nothing after the strip. Its reasoning
787 # is itself a summary of these turns, so recover it rather than
788 # strand them. Non-reasoning models never reach here.
789 reasoning = split_reasoning(response.text).reasoning.strip()
790 if reasoning:
791 return reasoning
792 log.warning("History compaction returned nothing for this batch")
793 except ProviderError as exc:
794 # A single message too big for the window cannot split; it falls
795 # through to the warning below.
796 if exc.kind is ProviderErrorKind.CONTEXT_OVERFLOW and len(batch) > 1:
797 mid = len(batch) // 2
798 first = self._summarize_batch(batch[:mid])
799 second = self._summarize_batch(batch[mid:])
800 merged = "\n".join(part for part in (first, second) if part.strip())
801 if merged.strip():
802 return merged
803 log.warning("History compaction failed for this batch", exc_info=True)
804 except Exception:
805 # warning, not debug: the user is told turns were dropped, so the
806 # reason must be in the log by default.
807 log.warning("History compaction failed for this batch", exc_info=True)
808 return ""
810 def _known_item_results(
811 self, question: str, chunk_type: ChunkType | None = None
812 ) -> list[SearchChunk]:
813 """Resolve a document named in *question* to its own chunks.
815 A question that names a document wants that document, not a ranking:
816 similarity search retrieves neighbors of the question's wording,
817 which for "summarize survey_214.pdf" is mostly noise. Resolution is
818 conservative: only a reference matching exactly one source routes;
819 anything ambiguous falls back to topical retrieval. Chunks come back
820 in document order with full canonical confidence, since their
821 relevance is established by the name match, not by similarity.
823 ``chunk_type`` scopes the content probe that resolves a reference
824 living in a document's text, so a scoped search cannot resolve
825 through rows it excludes. A wiki scope never routes here at all:
826 it asks for generated pages, not for a named document's own text.
827 """
828 if chunk_type == ChunkType.WIKI:
829 return []
830 for ref in document_references(question):
831 filename = self._resolve_reference_filename(ref, chunk_type)
832 chunks = self._document_chunks(filename)
833 if chunks:
834 log.info("Known-item route: %r resolved to %s", ref, filename)
835 return chunks
836 # No explicit reference: a known-item question shape may name the
837 # document by its human title ("summarize Frankenstein" against
838 # Frankenstein.txt), which has no filename, quote, or number cue.
839 for title in title_candidates(question):
840 filename = self._resolve_title_filename(title)
841 chunks = self._document_chunks(filename)
842 if chunks:
843 log.info("Known-item title route: %r resolved to %s", title, filename)
844 return chunks
845 return []
847 def _document_chunks(self, filename: str | None) -> list[SearchChunk]:
848 """A resolved document's chunks in document order at full confidence,
849 or empty for no resolution. Relevance is established by the name
850 match, not similarity, hence the canonical 1.0."""
851 if filename is None:
852 return []
853 chunks = self._store.get_chunks_by_source(filename)
854 chunks.sort(key=lambda c: c.chunk_index)
855 return [c.model_copy(update={"score": 1.0}) for c in chunks]
857 def _resolve_title_filename(self, title: str) -> str | None:
858 """The one source whose stem or stored title *title* names, or ``None``.
860 The article-stripped title pre-filters candidates by substring (over
861 filename and stored title), then the token-exact comparison decides;
862 only a unique winner routes, so shared titles fall back to topical
863 retrieval.
864 """
865 stripped = query_language().leading_article_pattern.sub("", title.strip())
866 candidates = self._store.get_sources(search=stripped, limit=_KNOWN_ITEM_CANDIDATES)
867 matches = [
868 s
869 for s in candidates
870 if matches_title(title, s["filename"]) or matches_stored_title(title, s.get("title"))
871 ]
872 if len(matches) == 1:
873 return str(matches[0]["filename"])
874 return None
876 def _resolve_reference_filename(
877 self, ref: str, chunk_type: ChunkType | None = None
878 ) -> str | None:
879 """The one source *ref* names, or ``None`` when nothing resolves uniquely.
881 Filename resolution first: substring search over-matches (a bare
882 "482" hits every zero-padded id containing it), so token-exact
883 matching disambiguates and only a unique winner routes. A unique
884 substring hit still routes for word refs (quoted titles never
885 token-match hyphenated filenames) but not numeric ones: "12" inside
886 "notes-2012" is the false match token comparison exists to reject.
888 When no filename knows the reference, it may be a docket-style number
889 living in the document's own text; a BM25 probe resolves it when the
890 hits concentrate in a single source.
891 """
892 candidates = self._store.get_sources(search=ref, limit=_KNOWN_ITEM_CANDIDATES)
893 matches = [s for s in candidates if matches_reference(ref, s["filename"])]
894 if len(matches) == 1:
895 return str(matches[0]["filename"])
896 if not matches and len(candidates) == 1 and not ref.strip().isdigit():
897 return str(candidates[0]["filename"])
898 if matches:
899 return None # several sources genuinely carry the reference
900 return self._resolve_reference_by_content(ref, chunk_type)
902 def _resolve_reference_by_content(
903 self, ref: str, chunk_type: ChunkType | None = None
904 ) -> str | None:
905 """Resolve *ref* to the single source whose text owns it, if any."""
906 hits = self._store.bm25_probe(ref, top_k=_KNOWN_ITEM_PROBE_K, chunk_type=chunk_type)
907 if len(hits) < _KNOWN_ITEM_PROBE_K:
908 return None
909 counts: dict[str, int] = {}
910 for hit in hits:
911 counts[hit.source] = counts.get(hit.source, 0) + 1
912 top_source, owned = max(counts.items(), key=lambda kv: kv[1])
913 if owned / len(hits) >= _KNOWN_ITEM_PROBE_MAJORITY:
914 return top_source
915 return None
917 def build_rag_context(
918 self,
919 question: str,
920 top_k: int = 0,
921 history: list[ChatMessage] | None = None,
922 *,
923 chunk_type: ChunkType | None = None,
924 ) -> RagContext | None:
925 """Build RAG context from search results.
927 ``chunk_type`` restricts the pool to ``"raw"`` (which covers table
928 chunks too) or ``"wiki"`` rows; ``None`` (default) searches the
929 mixed pool, or document chunks alone while the wiki is disabled.
930 """
931 retrieval_query = question
932 if history and self._config.history_rewrite:
933 retrieval_query = self._condense_question(question, history)
934 # Resolve a wiki:/raw: scope prefix the way search() does, so the scope
935 # it names reaches the known-item route and the wiki-disabled guard. Left
936 # in the query, the prefix sits in front of a document name and
937 # document_references resolves that name from the opposite pool; the
938 # search() call below re-parses the prefix for strategy routing.
939 mode, clean_query = self._parse_structured_query(retrieval_query)
940 requested = chunk_type
941 if requested is None and mode in (QueryMode.WIKI, QueryMode.RAW):
942 requested = ChunkType(mode.value)
943 scope = self._retrieval_scope(requested)
944 known_item = (
945 self._known_item_results(clean_query, scope) if self._config.intent_routing else []
946 )
947 if known_item:
948 # The named document IS the context; ranking and reranking would
949 # only reorder or drop parts of it. The budget fit below still
950 # trims to the context window, keeping the document's head.
951 results = known_item
952 else:
953 # search() reranks internally now, so the ask path no longer reranks
954 # again. It still requests the reranker's candidate depth so context
955 # assembly (prepare_results, select_context) works from a deep
956 # reranked pool rather than only top_k*2.
957 retrieve_k = top_k or self._config.top_k
958 if self._config.reranker_model:
959 retrieve_k = max(retrieve_k, self._config.rerank_candidates)
960 results = self.search(retrieval_query, top_k=retrieve_k, chunk_type=chunk_type)
961 results = filter_results(
962 results, self._config.max_distance, self._config.min_relevance_score
963 )
964 if not results:
965 # No relevant documents, but the user's stored memories may
966 # still ground the turn ("what's my name?"): facts recalled
967 # for this question answer via the memory-injected direct
968 # prompt instead of a refusal. Facts only -- always-injected
969 # preferences say nothing about answerability.
970 if self._memory_facts(question):
971 return RagContext([], self.direct_messages(question, history))
972 return None
973 results = prepare_results(results, self._config.diversity_max_per_source)
974 # Temporal filtering already ran inside search(); no need to repeat it here.
975 results = self.select_context(results, retrieval_query)
976 return self._finalize_context(results, question, history)
978 def _finalize_context(
979 self,
980 results: list[SearchChunk],
981 question: str,
982 history: list[ChatMessage] | None,
983 scale: float = 1.0,
984 ) -> RagContext:
985 """Fit *results* to the context budget and assemble the prompt.
987 Split from build_rag_context so an overflow retry can refit the same
988 retrieved set tighter without re-running retrieval or condensation.
989 """
990 system = self._system_with_memory(self._config.rag_system_prompt, question)
991 base_results = list(results)
992 budget = self._context_budget(system, question, history, scale)
993 results, used = self._fit_to_budget(results, budget)
994 results = self._widen_with_neighbors(results, max(0, budget - used))
995 context = build_context(results)
996 prompt = CONTEXT_TEMPLATE.format(context=context, question=question)
997 messages: list[ChatMessage] = [{"role": "system", "content": system}]
998 if history:
999 messages.extend(history)
1000 messages.append({"role": "user", "content": prompt})
1001 return RagContext(results, messages, base_results)
1003 def _context_budget(
1004 self,
1005 system: str,
1006 question: str,
1007 history: list[ChatMessage] | None,
1008 scale: float = 1.0,
1009 ) -> int:
1010 """Token budget left for source passages after the fixed prompt parts.
1012 The ceiling is the engine's ACTUAL per-slot window when known: the
1013 configured value is a target the dynamic picker aims for, and the
1014 server can come up smaller (the fleet divides context across slots).
1015 Budgeting against the target let a routed whole document overflow
1016 the real window and hard-fail the request with an HTTP 400.
1017 """
1018 configured = self._config.num_ctx or self._config.chat_n_ctx_target
1019 served = self._provider.served_chat_ctx()
1020 ctx = min(configured, served) if served else configured
1021 # Fit inside what the provider will actually accept: prompt_token_budget
1022 # already removes the generation reserve and the engine's margin, so the
1023 # sources get what is left after the rest of the prompt.
1024 non_source = (
1025 estimate_budget_tokens(system)
1026 + estimate_budget_tokens(question)
1027 + sum(estimate_budget_tokens(m["content"]) for m in history or [])
1028 + _CONTEXT_TEMPLATE_TOKENS
1029 )
1030 return int((prompt_token_budget(ctx) - non_source) * scale)
1032 def _fit_to_budget(
1033 self, results: list[SearchChunk], budget: int
1034 ) -> tuple[list[SearchChunk], int]:
1035 """Fit *results* into *budget*: the kept sources and the tokens they cost.
1037 ``max_context_sources`` caps by count; this caps by tokens so a
1038 retrieval-heavy query degrades gracefully instead of erroring with
1039 CONTEXT_OVERFLOW. The top-ranked source is always kept.
1041 Returning the spent total lets the caller derive the leftover for
1042 neighbor expansion instead of re-deriving the same per-chunk cost, so
1043 the two stages cannot drift apart on the accounting.
1044 """
1045 kept: list[SearchChunk] = []
1046 used = 0
1047 for r in results:
1048 cost = estimate_budget_tokens(r.chunk) + _PER_SOURCE_TOKENS
1049 if kept and used + cost > budget:
1050 break
1051 kept.append(r)
1052 used += cost
1053 if len(kept) < len(results):
1054 log.info(
1055 "Kept %d of %d sources to fit the model context window.",
1056 len(kept),
1057 len(results),
1058 )
1059 return kept, used
1061 def _widen_with_neighbors(self, results: list[SearchChunk], leftover: int) -> list[SearchChunk]:
1062 """Widen each fitted passage with adjacent same-source chunks.
1064 Spends only *leftover*, the budget the fit did not use, so a tight
1065 window sheds expansion first and never drops an original chunk for a
1066 neighbor. Widening keeps each passage's citation number and identity;
1067 its text and page/line span do change, so the sources block shows the
1068 widened range.
1069 """
1070 radius = self._config.neighbor_expansion
1071 if radius <= 0 or leftover <= 0:
1072 return results
1073 # With the structural filter on, expansion must not re-import the TOC
1074 # and cover text the filter dropped from the results.
1075 exclude = is_structural_chunk if self._config.filter_structural_chunks else None
1076 return expand_neighbors(
1077 results, self._store, radius, leftover, estimate_budget_tokens, exclude=exclude
1078 )
1080 def _system_with_memory(self, base_prompt: str, question: str) -> str:
1081 """Append the local-owner memory block to *base_prompt* when memory is enabled."""
1082 block = self._memory_block(question)
1083 return f"{base_prompt}\n\n{block}" if block else base_prompt
1085 def _memory_block(self, question: str) -> str:
1086 """Recall the local human's preferences and relevant facts as a system block.
1088 Preferences are always included; facts are recalled by similarity. Empty
1089 when memory is disabled or nothing matches. MCP agents never reach this path
1090 (their tools recall explicitly under their own owner).
1091 """
1092 if not self._config.memory_enabled:
1093 return ""
1094 # The human's answers see their own memories plus any an agent shared.
1095 preferences = self._store.get_memories(
1096 owner_predicate=human_recall_predicate(),
1097 kind=MemoryKind.PREFERENCE,
1098 )
1099 facts = self._memory_facts(question)
1100 return format_memory_block(preferences, facts, self._config.memory_token_budget)
1102 def _memory_facts(self, question: str) -> list[MemoryRow]:
1103 """Similarity-recalled facts for *question*, or empty.
1105 A non-empty result means memory can ground this turn on its own:
1106 facts are distance-gated against the question, unlike preferences,
1107 which are always injected and say nothing about answerability.
1108 """
1109 if not self._config.memory_enabled:
1110 return []
1111 if self._config.memory_top_k <= 0 or not self._embedder.embedding_available():
1112 return []
1113 vector = self._embedder.embed_query(question)
1114 return self._store.search_memories(
1115 vector,
1116 owner_predicate=human_recall_predicate(),
1117 top_k=self._config.memory_top_k,
1118 max_distance=self._config.memory_max_distance,
1119 )
1121 def _answer_aggregate(self, aggregate: AggregateQuery) -> str:
1122 """Answer a count-shaped question with an exact full-corpus scan.
1124 Top-k retrieval sees a handful of chunks out of the whole corpus, so
1125 it structurally cannot count; the faithful-but-useless outcome is a
1126 model hedging that "the context does not provide precise counts".
1127 Counting is a scan, and a scan needs no language model: the numbers
1128 below are exact, not generated.
1129 """
1130 if aggregate.kind is AggregateKind.TOTAL_SOURCES:
1131 sources = self._store.count_sources()
1132 chunks = self._store.count_chunks()
1133 return f"The index holds {sources} documents split into {chunks} searchable passages."
1134 if aggregate.kind is AggregateKind.TERM_MENTIONS:
1135 chunk_hits, source_hits = self._store.count_term_mentions(aggregate.term)
1136 return (
1137 f"Exact scan of the whole index: {source_hits} documents mention "
1138 f"{aggregate.term!r}, across {chunk_hits} passages. This counts literal "
1139 f"mentions of the phrase, not paraphrases."
1140 )
1141 if aggregate.kind in (AggregateKind.DISTINCT_TYPE, AggregateKind.TYPE_ASSOCIATION):
1142 typed = self._answer_typed_aggregate(aggregate)
1143 if typed is not None:
1144 return typed
1145 return self._decline_aggregate()
1147 def _decline_aggregate(self) -> str:
1148 """The honest no-capability answer, naming what IS countable."""
1149 from lilbee.retrieval.entities import load_schema
1151 schema = load_schema(self._store)
1152 if schema is not None and schema.types:
1153 countable = ", ".join(sorted(t.name.replace("_", " ") for t in schema.types))
1154 return (
1155 "That count isn't answerable from the extracted records. Countable "
1156 f"entity types in this index: {countable}. I can also count documents "
1157 "or passages that mention a specific term."
1158 )
1159 return (
1160 "Answering that count needs structured records (dates, identifiers, or "
1161 "entities) that aren't extracted from this corpus yet. I can count "
1162 "documents or passages that mention a specific term, or you can ask for "
1163 "the passages themselves and count from those."
1164 )
1166 def _answer_typed_aggregate(self, aggregate: AggregateQuery) -> str | None:
1167 """Exact answers over extracted entities, or None when the question's
1168 nouns don't resolve against the extraction schema."""
1169 from lilbee.retrieval.entities import load_schema
1171 schema = load_schema(self._store)
1172 counted = schema.type_for_noun(aggregate.noun) if schema else None
1173 if schema is None or counted is None:
1174 return None
1175 if aggregate.kind is AggregateKind.DISTINCT_TYPE:
1176 return self._answer_distinct_count(counted.name, asked_for=aggregate.noun)
1177 grouped = schema.type_for_noun(aggregate.group_noun)
1178 if grouped is None:
1179 return None
1180 return self._answer_association_count(counted.name, grouped.name)
1182 def _answer_distinct_count(self, type_name: str, asked_for: str = "") -> str:
1183 pretty = type_name.replace("_", " ")
1184 mentions, distinct = self._store.entity_value_counts(type_name)
1185 if mentions == 0:
1186 return (
1187 f"No {pretty} entities are extracted yet; "
1188 "run a sync with entity extraction enabled first."
1189 )
1190 answer = (
1191 f"Exact scan of the extracted records: {distinct} distinct "
1192 f"{pretty} values, across {mentions} mentions."
1193 )
1194 if asked_for and not _noun_names_type(asked_for, type_name):
1195 # A synonym resolved the question's noun to a proxy type; counting
1196 # one is not counting the other (one aircraft flies many flights),
1197 # so the answer must say which quantity it actually measured.
1198 answer += (
1199 f" Note: this counts {pretty} values, the closest extracted type to "
1200 f"{asked_for.strip()!r}, which may not be the same quantity."
1201 )
1202 return answer
1204 def _answer_association_count(self, counted: str, grouped: str) -> str:
1205 counted_pretty = counted.replace("_", " ")
1206 grouped_pretty = grouped.replace("_", " ")
1207 counts = self._store.entity_association_counts(counted, grouped_by=grouped)
1208 if not counts:
1209 return (
1210 f"No co-occurring {counted_pretty} and {grouped_pretty} entities are "
1211 "extracted yet; run a sync with entity extraction enabled first."
1212 )
1213 shown = list(counts.items())[:_ASSOCIATION_LINES]
1214 lines = "\n".join(f" {value}: {n}" for value, n in shown)
1215 more = len(counts) - len(shown)
1216 suffix = f"\n ... and {more} more" if more > 0 else ""
1217 return (
1218 f"Exact counts from the extracted records ({counted_pretty} "
1219 f"per {grouped_pretty}, by shared passage):\n{lines}{suffix}"
1220 )
1222 def route_direct_answer(self, question: str) -> str | None:
1223 """The exact-scan answer for a count-shaped question, else ``None``.
1225 Every retrieval entry point must consult this before building RAG
1226 context: ask_raw/ask_stream do (covering CLI and TUI), and the HTTP
1227 handlers call it themselves because they assemble their own prompts
1228 from build_rag_context. An entry point that skips it hedges at the
1229 count questions every other surface answers exactly.
1230 """
1231 if not self._config.intent_routing:
1232 return None
1233 aggregate = parse_aggregate(question)
1234 if self._config.intent_llm and (
1235 aggregate is None or aggregate.kind is AggregateKind.UNSUPPORTED
1236 ):
1237 # The deterministic patterns found no answerable count shape; let
1238 # the chat model classify phrasings (and languages) they miss. A
1239 # ``None`` here keeps whatever the patterns concluded, so an LLM
1240 # failure can never lose a deterministic decline.
1241 aggregate = self._llm_classify_aggregate(question) or aggregate
1242 if aggregate is None:
1243 return None
1244 log.info("Aggregate route: %s for %r", aggregate.kind.value, question)
1245 return self._answer_aggregate(aggregate)
1247 def _llm_classify_aggregate(self, question: str) -> AggregateQuery | None:
1248 """One short classification call, mapped conservatively to a route.
1250 Any provider failure or malformed reply means no route -- the same
1251 harmless degrade to topical retrieval as a deterministic miss.
1252 """
1253 prompt = INTENT_CLASSIFY_PROMPT.format(question=question)
1254 try:
1255 response = self._provider.chat(
1256 [{"role": "user", "content": prompt}],
1257 stream=False,
1258 options={
1259 "num_predict": INTENT_CLASSIFY_MAX_TOKENS,
1260 "response_format": json_reply_format(),
1261 },
1262 )
1263 except Exception:
1264 log.debug("LLM intent classification failed; using pattern result", exc_info=True)
1265 return None
1266 parsed = parse_llm_aggregate(strip_reasoning(response.text))
1267 if parsed is not None:
1268 log.info("LLM intent route: %s for %r", parsed.kind.value, question)
1269 return parsed
1271 def skip_retrieval(self) -> bool:
1272 """Whether this turn should bypass RAG: chat-only mode or no embedder."""
1273 return (
1274 self._config.chat_mode == ChatMode.CHAT.value
1275 or not self._embedder.embedding_available()
1276 )
1278 def search_unavailable(self) -> bool:
1279 """Search mode is active but retrieval can't run because no embedder is loaded.
1281 Ask refuses cleanly in this state (best UX: tell the user search needs an
1282 embedder) rather than silently answering ungrounded. Chat mode is exempt --
1283 it intentionally answers off-corpus, so it falls back instead of refusing.
1284 """
1285 return (
1286 self._config.chat_mode != ChatMode.CHAT.value
1287 and not self._embedder.embedding_available()
1288 )
1290 def library_empty(self) -> bool:
1291 """Whether the store holds no indexed content yet (nothing to search)."""
1292 return not self._store.has_chunks()
1294 def direct_messages(
1295 self, question: str, history: list[ChatMessage] | None = None
1296 ) -> list[ChatMessage]:
1297 """Build messages for direct LLM chat (no RAG context)."""
1298 messages: list[ChatMessage] = [
1299 {
1300 "role": "system",
1301 "content": self._system_with_memory(self._config.general_system_prompt, question),
1302 }
1303 ]
1304 if history:
1305 messages.extend(history)
1306 messages.append({"role": "user", "content": question})
1307 return messages
1309 def _messages_for_provider(self, messages: list[ChatMessage]) -> list[dict[str, str]]:
1310 """Convert ChatMessage list to provider-expected format."""
1311 return [{"role": m["role"], "content": m["content"]} for m in messages]
1313 def _direct_chat(
1314 self,
1315 question: str,
1316 history: list[ChatMessage] | None,
1317 options: dict[str, Any] | None,
1318 ) -> str:
1319 """Run a no-RAG chat turn and return the cleaned response."""
1320 messages = self.direct_messages(question, history)
1321 provider_messages = self._messages_for_provider(messages)
1322 opts = options if options is not None else self._config.generation_options()
1323 result = self._provider.chat(provider_messages, options=opts or None)
1324 raw = result.text
1325 return raw if self._config.show_reasoning else strip_reasoning(raw)
1327 def pre_retrieval_answer(self, question: str) -> str | None:
1328 """The canned answer a grounded turn gives before retrieval runs:
1329 the empty-library guidance, or a count question's exact scan.
1330 ``None`` means retrieval should proceed. One ladder shared by every
1331 surface (ask, stream, HTTP) so they cannot drift.
1333 An empty library with recalled memory facts falls through: memory is
1334 the user's own ground truth, so build_rag_context answers from it
1335 instead of this method telling the user to add documents.
1336 """
1337 if self.library_empty() and not self._memory_facts(question):
1338 return EMPTY_LIBRARY
1339 return self.route_direct_answer(question)
1341 def ask_raw(
1342 self,
1343 question: str,
1344 top_k: int = 0,
1345 history: list[ChatMessage] | None = None,
1346 options: dict[str, Any] | None = None,
1347 *,
1348 chunk_type: ChunkType | None = None,
1349 ) -> AskResult:
1350 """Ask a question. Refuses cleanly without an embedder (search can't
1351 ground); falls back to direct chat only when chat_mode is 'chat'."""
1352 if self.search_unavailable():
1353 return AskResult(answer=SEARCH_NEEDS_EMBEDDER, sources=[])
1354 if self.skip_retrieval():
1355 return AskResult(answer=self._direct_chat(question, history, options), sources=[])
1356 pre_answer = self.pre_retrieval_answer(question)
1357 if pre_answer is not None:
1358 return AskResult(answer=pre_answer, sources=[])
1359 rag = self.build_rag_context(question, top_k=top_k, history=history, chunk_type=chunk_type)
1360 if rag is None:
1361 return AskResult(answer=GROUNDED_REFUSAL, sources=[])
1362 results, messages = rag.results, rag.messages
1363 opts = options if options is not None else self._config.generation_options()
1364 try:
1365 result = self._provider.chat(
1366 self._messages_for_provider(messages), options=opts or None
1367 )
1368 except ProviderError as exc:
1369 if exc.kind is not ProviderErrorKind.CONTEXT_OVERFLOW or not results:
1370 raise
1371 # The budget estimator is a heuristic; when the engine still reports
1372 # overflow, refit tighter and retry once. Refit from the pre-widen
1373 # set so the tighter budget sheds neighbor expansion before it drops
1374 # an original chunk, not the reverse.
1375 log.warning("Context overflow despite budgeting; retrying with a tighter fit")
1376 retry = self._finalize_context(
1377 rag.base_results if rag.base_results is not None else results,
1378 question,
1379 history,
1380 scale=_OVERFLOW_RETRY_SCALE,
1381 )
1382 results, messages = retry.results, retry.messages
1383 result = self._provider.chat(
1384 self._messages_for_provider(messages), options=opts or None
1385 )
1386 raw = result.text
1387 clean = raw if self._config.show_reasoning else strip_reasoning(raw)
1388 # Citations are read off the prose only: a model that echoes its own
1389 # Sources list would otherwise mark every retrieved file cited.
1390 return AskResult(
1391 answer=clean,
1392 sources=results,
1393 cited_sources=cited_subset(strip_llm_citations(clean), results),
1394 )
1396 def ask(
1397 self,
1398 question: str,
1399 top_k: int = 0,
1400 history: list[ChatMessage] | None = None,
1401 options: dict[str, Any] | None = None,
1402 *,
1403 chunk_type: ChunkType | None = None,
1404 ) -> str:
1405 """Ask a question and get a formatted answer with citations."""
1406 result = self.ask_raw(
1407 question, top_k=top_k, history=history, options=options, chunk_type=chunk_type
1408 )
1409 if not result.sources:
1410 return result.answer
1411 answer = strip_llm_citations(result.answer)
1412 return f"{answer}{format_sources_block(result.sources)}"
1414 def _stream_direct(
1415 self,
1416 question: str,
1417 history: list[ChatMessage] | None,
1418 options: dict[str, Any] | None,
1419 ) -> Generator[StreamToken, None, None]:
1420 """Streaming branch with the general system prompt (no RAG context)."""
1421 messages = self.direct_messages(question, history)
1422 provider_messages = self._messages_for_provider(messages)
1423 opts = options if options is not None else self._config.generation_options()
1424 events = stream_chat_with_cap(
1425 self._provider,
1426 cast("list[dict[str, Any]]", provider_messages),
1427 options=opts,
1428 model=self._config.chat_model,
1429 show_reasoning=self._config.show_reasoning,
1430 cap_chars=effective_reasoning_cap(),
1431 )
1432 try:
1433 yield from cap_events_as_stream_tokens(events)
1434 except (ConnectionError, OSError) as exc:
1435 yield StreamToken(content=f"\n\n[Connection lost: {exc}]", is_reasoning=False)
1437 def ask_stream(
1438 self,
1439 question: str,
1440 top_k: int = 0,
1441 history: list[ChatMessage] | None = None,
1442 options: dict[str, Any] | None = None,
1443 *,
1444 chunk_type: ChunkType | None = None,
1445 ) -> Generator[StreamToken, None, None]:
1446 """Stream answer tokens with citations appended at the end."""
1447 if self.search_unavailable():
1448 yield StreamToken(content=SEARCH_NEEDS_EMBEDDER, is_reasoning=False)
1449 return
1450 if self.skip_retrieval():
1451 yield from self._stream_direct(question, history, options)
1452 return
1453 pre_answer = self.pre_retrieval_answer(question)
1454 if pre_answer is not None:
1455 yield StreamToken(content=pre_answer, is_reasoning=False)
1456 return
1458 rag = self.build_rag_context(question, top_k=top_k, history=history, chunk_type=chunk_type)
1459 if rag is None:
1460 yield StreamToken(content=GROUNDED_REFUSAL, is_reasoning=False)
1461 return
1462 results, messages = rag.results, rag.messages
1463 # No overflow retry here: a stream cannot be rebuilt once tokens have
1464 # been yielded, so the conservative budget the context fit already
1465 # applied is the streaming path's protection.
1466 provider_messages = self._messages_for_provider(messages)
1467 opts = options if options is not None else self._config.generation_options()
1468 events = stream_chat_with_cap(
1469 self._provider,
1470 cast("list[dict[str, Any]]", provider_messages),
1471 options=opts,
1472 model=self._config.chat_model,
1473 show_reasoning=self._config.show_reasoning,
1474 cap_chars=effective_reasoning_cap(),
1475 )
1476 yield from self._filtered_answer_tokens(events)
1477 # A model that emits its own trailing Sources block has had it dropped by
1478 # the filter above; this authoritative list is numbered to match the
1479 # ``[n]`` markers the model used, so every citation resolves to a line.
1480 block = format_sources_block(results)
1481 if block:
1482 yield StreamToken(content=block, is_reasoning=False)
1484 def _filtered_answer_tokens(
1485 self, events: Generator[Any, None, None]
1486 ) -> Generator[StreamToken, None, None]:
1487 """Pump model events through the streaming citation filter.
1489 Reasoning tokens pass through untouched; answer tokens are withheld
1490 while they could still be the start of a model-authored Sources
1491 block, and any held-back tail is released when the stream ends.
1492 """
1493 cite_filter = StreamingCitationFilter()
1494 try:
1495 for token in cap_events_as_stream_tokens(events):
1496 if token.is_reasoning:
1497 yield token
1498 continue
1499 shown = cite_filter.feed(token.content)
1500 if shown:
1501 yield StreamToken(content=shown, is_reasoning=False)
1502 except (ConnectionError, OSError) as exc:
1503 yield StreamToken(content=f"\n\n[Connection lost: {exc}]", is_reasoning=False)
1504 tail = cite_filter.flush()
1505 if tail:
1506 yield StreamToken(content=tail, is_reasoning=False)