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