Coverage for src/lilbee/cli/commands/search_chat.py: 100%
268 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-17 10:02 +0000
1"""Search, ask, chat, and topics commands."""
3from __future__ import annotations
5import re
6import sys
7from collections.abc import Callable
8from pathlib import Path
9from typing import Any, NoReturn
11import typer
12from rich.table import Table
14from lilbee.app.search import clean_result
15from lilbee.app.services import get_services
16from lilbee.cli import theme
17from lilbee.cli.app import (
18 apply_overrides,
19 chat_model_overridden,
20 console,
21 data_dir_option,
22 global_option,
23 model_option,
24 num_ctx_option,
25 repeat_penalty_option,
26 seed_option,
27 temperature_option,
28 top_k_sampling_option,
29 top_p_option,
30)
31from lilbee.cli.commands._shared import CHUNK_PREVIEW_LEN
32from lilbee.cli.helpers import (
33 announce_cold_start,
34 announce_ready,
35 announce_retrieval_query,
36 auto_sync,
37 json_output,
38)
39from lilbee.cli.log_routing import route_diagnostics_to_log_file
40from lilbee.core.config import cfg
41from lilbee.data.store import EmbeddingModelMismatchError, SearchScope, scope_to_chunk_type
42from lilbee.providers.base import ProviderError
43from lilbee.providers.roles import WorkerRole
45# How many top concepts to show inline before truncating with a ``+N more`` tail.
46_TOPIC_PREVIEW_LIMIT = 5
47# Upper bound on retrieved results, matching the REST search route's le=100 cap.
48_MAX_TOP_K = 100
50_EMBED_MISMATCH_ADOPT_HINT = (
51 "Run `lilbee use-embedder {model}` to search this index with its embedder."
52)
53_EMBED_MISMATCH_REBUILD_HINT = (
54 "This index needs a {dim}-dim embedder; run `lilbee rebuild` to re-embed it "
55 "under your current model."
56)
59def _exit_embedding_mismatch(exc: EmbeddingModelMismatchError) -> NoReturn:
60 """Print a surface-appropriate mismatch error and exit non-zero.
62 Headless: never switches embedder silently. Names the index's embedder and,
63 when adoptable (same dim), the one command that makes it searchable.
64 """
65 hint = (
66 _EMBED_MISMATCH_ADOPT_HINT.format(model=exc.persisted_model)
67 if exc.dims_match
68 else _EMBED_MISMATCH_REBUILD_HINT.format(dim=exc.persisted_dim)
69 )
70 if cfg.json_mode:
71 json_output({"error": str(exc), "hint": hint, "persisted_model": exc.persisted_model})
72 raise SystemExit(1)
73 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] {exc}")
74 console.print(hint)
75 raise SystemExit(1)
78_scope_option = typer.Option(
79 SearchScope.BOTH,
80 "--scope",
81 "-s",
82 help="Restrict the pool to raw chunks, wiki pages, or both (default).",
83 case_sensitive=False,
84)
87def _swap_stale_models_to_installed(chat_overridden: bool = False) -> None:
88 """Swap a stale chat/embedding ref to an installed model for this run only.
90 The persisted config is never rewritten from a one-shot command; durable
91 swaps belong to the TUI's interactive startup canonicalization. An explicit
92 --model override is honored as given, so an unusable value surfaces the
93 engine's not-installed error instead of a silent substitute. The notice
94 prints to stderr in every mode: stdout stays answer-only and JSON stays
95 parseable.
96 """
97 from rich.console import Console
99 from lilbee.app.settings import apply_ephemeral_model_swap
100 from lilbee.modelhub.model_manager import (
101 ValidationResult,
102 canonicalize_chat_model,
103 canonicalize_embedding_model,
104 )
106 checks = [(canonicalize_embedding_model(), "embedding_model", "Embedding")]
107 if not chat_overridden:
108 checks.insert(0, (canonicalize_chat_model(), "chat_model", "Chat"))
109 err = Console(stderr=True)
110 for canon, field, label in checks:
111 if canon.status == ValidationResult.OK or canon.effective == canon.original:
112 continue
113 apply_ephemeral_model_swap(field, canon.effective)
114 if canon.original:
115 notice = (
116 f"{label} model {canon.original!r} is unavailable ({canon.reason}); "
117 f"using installed {canon.effective!r} for this run."
118 )
119 else:
120 notice = f"No {label.lower()} model configured; using installed {canon.effective!r}."
121 err.print(notice, style=theme.WARNING)
124_MD_FILE_LINK_RE = re.compile(r"\[([^\]]+)\]\((file://[^)]+)\)")
127def _print_answer_stream(stream: Any, on_first_token: Callable[[], None]) -> None:
128 """Stream an answer to stdout verbatim, then render its Sources block.
130 Tokens print with markup off: model text is data, and Rich markup would eat
131 the ``[label]`` of every markdown source link (and let the model restyle the
132 terminal). On a terminal the Sources block's ``[label](file://...)`` links
133 become OSC 8 hyperlinks, clickable even when the path wraps; piped output
134 and legacy Windows consoles keep the raw markdown so the URL survives (Rich
135 emits no OSC 8 on the legacy path, which would drop it). The marker can span
136 tokens, so a marker-sized tail is held back until it can be classified.
137 """
138 from lilbee.retrieval.query.formatting import SOURCES_BLOCK_MARKER
139 from lilbee.retrieval.reasoning import RetrievalNotice
141 buf = ""
142 hold = len(SOURCES_BLOCK_MARKER)
143 in_sources = False
144 flushed = False
145 try:
146 for token in stream:
147 if isinstance(token, RetrievalNotice):
148 announce_retrieval_query(token.query)
149 continue
150 on_first_token()
151 if token.is_reasoning:
152 # Reasoning is not filtered by StreamingCitationFilter, so a
153 # thinking trace drafting a Sources: list must never trip the
154 # marker scan; it prints live and bypasses the buffer.
155 console.print(token.content, end="", markup=False)
156 continue
157 buf += token.content
158 if in_sources:
159 continue
160 if SOURCES_BLOCK_MARKER in buf:
161 head, _, buf = buf.partition(SOURCES_BLOCK_MARKER)
162 console.print(head, end="", markup=False)
163 in_sources = True
164 elif len(buf) > hold:
165 console.print(buf[:-hold], end="", markup=False)
166 buf = buf[-hold:]
167 flushed = True
168 finally:
169 if not flushed and buf and not in_sources:
170 # An exception escaped the stream mid-answer: the held tail is
171 # real answer text; print it before the error surfaces.
172 console.print(buf, markup=False)
173 if not in_sources:
174 console.print(buf, markup=False)
175 return
176 console.print(SOURCES_BLOCK_MARKER, end="", markup=False)
177 _print_sources_block(buf)
180def _print_sources_block(block: str) -> None:
181 """Render the Sources block, turning ``[label](file://...)`` into terminal links."""
182 from rich.markup import escape
184 if not console.is_terminal or console.legacy_windows:
185 console.print(block, markup=False, highlight=False)
186 return
187 parts: list[str] = []
188 last = 0
189 for m in _MD_FILE_LINK_RE.finditer(block):
190 parts.append(escape(block[last : m.start()]))
191 parts.append(f"[link={m.group(2)}]{escape(m.group(1))}[/link]")
192 last = m.end()
193 parts.append(escape(block[last:]))
194 console.print("".join(parts), highlight=False)
197def _reject_if_empty(value: str, label: str) -> None:
198 """Exit with a uniform error if *value* is empty/whitespace (matches REST)."""
199 if value and value.strip():
200 return
201 msg = f"{label} must not be empty"
202 if cfg.json_mode:
203 json_output({"error": msg})
204 raise SystemExit(1)
205 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] {msg}")
206 raise SystemExit(1)
209def _display_score(result: dict[str, Any]) -> float:
210 """Relevance score, else distance, else 0.0. Explicit None checks keep a
211 legitimate 0.0 from falling through a truthy ``or`` chain."""
212 score = result.get("relevance_score")
213 if score is None:
214 score = result.get("distance")
215 return 0.0 if score is None else score
218def search(
219 query: str = typer.Argument(..., help="Search query"),
220 top_k: int = typer.Option(None, "--top-k", "-k", min=1, help="Number of results"),
221 scope: SearchScope = _scope_option,
222 data_dir: Path | None = data_dir_option,
223 use_global: bool = global_option,
224) -> None:
225 """Search the knowledge base for relevant chunks."""
226 apply_overrides(data_dir=data_dir, use_global=use_global)
227 route_diagnostics_to_log_file()
229 _reject_if_empty(query, "query")
231 err = announce_cold_start(WorkerRole.EMBED, str(cfg.embedding_model))
232 try:
233 results = get_services().searcher.search(
234 query,
235 top_k=min(top_k or cfg.top_k, _MAX_TOP_K),
236 chunk_type=scope_to_chunk_type(scope),
237 )
238 announce_ready(err, WorkerRole.EMBED)
239 except EmbeddingModelMismatchError as exc:
240 _exit_embedding_mismatch(exc)
241 except Exception as exc:
242 if cfg.json_mode:
243 json_output({"error": str(exc)})
244 raise SystemExit(1) from None
245 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] {exc}")
246 raise SystemExit(1) from None
247 cleaned = [clean_result(r) for r in results]
249 if cfg.json_mode:
250 json_output({"command": "search", "query": query, "results": cleaned})
251 return
253 if not cleaned:
254 console.print("No results found.")
255 return
257 has_relevance = any("relevance_score" in r for r in cleaned)
258 table = Table(title="Search Results")
259 table.add_column("Source", style=theme.ACCENT)
260 table.add_column("Chunk", max_width=80)
261 score_label = "Score" if has_relevance else "Distance"
262 table.add_column(score_label, justify="right", style=theme.MUTED)
264 for r in cleaned:
265 chunk_text = r["chunk"]
266 preview = chunk_text[:CHUNK_PREVIEW_LEN]
267 if len(chunk_text) > CHUNK_PREVIEW_LEN:
268 preview += "..."
269 table.add_row(r["source"], preview, f"{_display_score(r):.4f}")
270 console.print(table)
273def ask(
274 question: str = typer.Argument(..., help="Question to ask"),
275 scope: SearchScope = _scope_option,
276 data_dir: Path | None = data_dir_option,
277 model: str | None = model_option,
278 use_global: bool = global_option,
279 temperature: float | None = temperature_option,
280 top_p: float | None = top_p_option,
281 top_k_sampling: int | None = top_k_sampling_option,
282 repeat_penalty: float | None = repeat_penalty_option,
283 num_ctx: int | None = num_ctx_option,
284 seed: int | None = seed_option,
285 no_sync: bool = typer.Option(
286 False, "--no-sync", help="Skip the pre-answer auto-sync (useful on large static corpora)."
287 ),
288) -> None:
289 """Ask a one-shot question."""
290 apply_overrides(
291 data_dir=data_dir,
292 model=model,
293 use_global=use_global,
294 temperature=temperature,
295 top_p=top_p,
296 top_k_sampling=top_k_sampling,
297 repeat_penalty=repeat_penalty,
298 num_ctx=num_ctx,
299 seed=seed,
300 )
301 route_diagnostics_to_log_file()
302 _reject_if_empty(question, "question")
304 try:
305 from lilbee.app.settings import apply_settings_update
306 from lilbee.modelhub.models import ensure_chat_model
308 pulled = ensure_chat_model()
309 if pulled is not None:
310 apply_settings_update({"chat_model": pulled})
311 _swap_stale_models_to_installed(chat_overridden=chat_model_overridden())
312 get_services().embedder.validate_model()
313 if cfg.auto_sync and not no_sync:
314 if cfg.json_mode:
315 from rich.console import Console as _QuietConsole
317 auto_sync(_QuietConsole(quiet=True))
318 else:
319 auto_sync(console)
321 chunk_type = scope_to_chunk_type(scope)
323 if cfg.json_mode:
324 result = get_services().searcher.ask_raw(question, chunk_type=chunk_type)
325 json_output(
326 {
327 "command": "ask",
328 "question": question,
329 "answer": result.answer,
330 "sources": [clean_result(s) for s in result.sources],
331 "cited_sources": [clean_result(s) for s in result.cited_sources],
332 "retrieval_query": result.retrieval_query,
333 "dropped_sources": [clean_result(s) for s in result.dropped_sources],
334 }
335 )
336 return
338 err = announce_cold_start(WorkerRole.CHAT, str(cfg.chat_model))
339 stream = get_services().searcher.ask_stream(question, chunk_type=chunk_type)
340 first = True
342 def _on_first_token() -> None:
343 nonlocal first
344 if first:
345 announce_ready(err, WorkerRole.CHAT)
346 first = False
348 _print_answer_stream(stream, on_first_token=_on_first_token)
349 except EmbeddingModelMismatchError as exc:
350 _exit_embedding_mismatch(exc)
351 except (RuntimeError, ProviderError) as exc:
352 if cfg.json_mode:
353 json_output({"error": str(exc)})
354 raise SystemExit(1) from None
355 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] {exc}")
356 raise SystemExit(1) from None
359def use_embedder(
360 ref: str = typer.Argument(
361 ..., help="Embedding model ref to adopt (copy it from a downloaded index's error)."
362 ),
363 data_dir: Path | None = data_dir_option,
364 use_global: bool = global_option,
365) -> None:
366 """Switch to embedder REF, downloading it if needed, without rebuilding the index."""
367 apply_overrides(data_dir=data_dir, use_global=use_global)
369 from lilbee.app.models import adopt_embedder
370 from lilbee.catalog.compat import UnsupportedArchError
372 try:
373 result = adopt_embedder(ref)
374 except (RuntimeError, ValueError, OSError, UnsupportedArchError) as exc:
375 if cfg.json_mode:
376 json_output({"error": str(exc)})
377 raise SystemExit(1) from None
378 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] {exc}")
379 raise SystemExit(1) from None
381 if cfg.json_mode:
382 json_output(
383 {"command": "use-embedder", "model": result.model, "status": result.status.value}
384 )
385 return
386 console.print(f"Now embedding with [{theme.ACCENT}]{result.model}[/{theme.ACCENT}].")
389def chat(
390 data_dir: Path | None = data_dir_option,
391 model: str | None = model_option,
392 use_global: bool = global_option,
393 temperature: float | None = temperature_option,
394 top_p: float | None = top_p_option,
395 top_k_sampling: int | None = top_k_sampling_option,
396 repeat_penalty: float | None = repeat_penalty_option,
397 num_ctx: int | None = num_ctx_option,
398 seed: int | None = seed_option,
399) -> None:
400 """Interactive chat loop. Press S in the TUI to sync pending documents."""
401 apply_overrides(
402 data_dir=data_dir,
403 model=model,
404 use_global=use_global,
405 temperature=temperature,
406 top_p=top_p,
407 top_k_sampling=top_k_sampling,
408 repeat_penalty=repeat_penalty,
409 num_ctx=num_ctx,
410 seed=seed,
411 )
412 # No diagnostics routing here: chat hands off to the TUI, which owns its
413 # logging (tui.log). Routing first would latch captureWarnings and leave a
414 # NOTSET cli.log handler running under the TUI.
416 if cfg.json_mode:
417 json_output({"error": "Chat requires a terminal, not --json"})
418 raise SystemExit(1)
419 if not sys.stdin.isatty() or not sys.stdout.isatty():
420 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] Chat requires a terminal.")
421 raise SystemExit(1)
422 from lilbee.cli.tui import run_tui
424 run_tui()
427def topics(
428 query: str = typer.Argument(None, help="Optional query to find related concepts."),
429 top_k: int = typer.Option(10, "--top-k", "-k", help="Number of results."),
430 data_dir: Path | None = data_dir_option,
431 use_global: bool = global_option,
432) -> None:
433 """Show top concept communities or concepts related to a query."""
434 apply_overrides(data_dir=data_dir, use_global=use_global)
436 from lilbee.retrieval.concepts import concepts_available
438 if not concepts_available():
439 msg = "Concept graph requires: pip install 'lilbee[graph]'"
440 if cfg.json_mode:
441 json_output({"error": msg})
442 raise SystemExit(1)
443 console.print(f"[{theme.ERROR}]{msg}[/{theme.ERROR}]")
444 raise SystemExit(1)
446 if not cfg.concept_graph:
447 if cfg.json_mode:
448 json_output({"error": "Concept graph is disabled (LILBEE_CONCEPT_GRAPH=false)"})
449 raise SystemExit(1)
450 console.print(
451 f"[{theme.ERROR}]Concept graph is disabled.[/{theme.ERROR}] "
452 "Enable with LILBEE_CONCEPT_GRAPH=true"
453 )
454 raise SystemExit(1)
456 if not get_services().concepts.get_graph():
457 if cfg.json_mode:
458 json_output({"error": "Concept graph not available"})
459 raise SystemExit(1)
460 console.print(f"[{theme.ERROR}]Concept graph not available.[/{theme.ERROR}]")
461 raise SystemExit(1)
463 if query:
464 _topics_for_query(query)
465 else:
466 _topics_overview(top_k)
469def _topics_for_query(query: str) -> None:
470 """Show concepts related to a query."""
471 cg = get_services().concepts
472 concepts = cg.extract_concepts(query)
473 related = cg.expand_query(query)
474 all_concepts = concepts + [r for r in related if r not in concepts]
476 if cfg.json_mode:
477 json_output({"command": "topics", "query": query, "concepts": all_concepts})
478 return
479 if not all_concepts:
480 console.print("No concepts found for this query.")
481 return
482 console.print(f"Concepts related to [{theme.ACCENT}]{query}[/{theme.ACCENT}]:")
483 for c in all_concepts:
484 console.print(f" {c}")
487def _topics_overview(top_k: int) -> None:
488 """Show top concept communities."""
489 from dataclasses import asdict
491 communities = get_services().concepts.top_communities(k=top_k)
492 if cfg.json_mode:
493 json_output({"command": "topics", "communities": [asdict(c) for c in communities]})
494 return
495 if not communities:
496 console.print("No concept communities found. Try syncing some documents first.")
497 return
498 table = Table(title="Concept Communities")
499 table.add_column("Cluster", justify="right", style=theme.MUTED)
500 table.add_column("Size", justify="right")
501 table.add_column("Top Concepts", style=theme.ACCENT)
502 for comm in communities:
503 preview = ", ".join(comm.concepts[:_TOPIC_PREVIEW_LIMIT])
504 if len(comm.concepts) > _TOPIC_PREVIEW_LIMIT:
505 preview += f" (+{len(comm.concepts) - _TOPIC_PREVIEW_LIMIT} more)"
506 table.add_row(str(comm.cluster_id), str(comm.size), preview)
507 console.print(table)