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