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