Coverage for src/lilbee/cli/commands/search_chat.py: 100%
268 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-08 09:20 +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 }
334 )
335 return
337 err = announce_cold_start(WorkerRole.CHAT, str(cfg.chat_model))
338 stream = get_services().searcher.ask_stream(question, chunk_type=chunk_type)
339 first = True
341 def _on_first_token() -> None:
342 nonlocal first
343 if first:
344 announce_ready(err, WorkerRole.CHAT)
345 first = False
347 _print_answer_stream(stream, on_first_token=_on_first_token)
348 except EmbeddingModelMismatchError as exc:
349 _exit_embedding_mismatch(exc)
350 except (RuntimeError, ProviderError) as exc:
351 if cfg.json_mode:
352 json_output({"error": str(exc)})
353 raise SystemExit(1) from None
354 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] {exc}")
355 raise SystemExit(1) from None
358def use_embedder(
359 ref: str = typer.Argument(
360 ..., help="Embedding model ref to adopt (copy it from a downloaded index's error)."
361 ),
362 data_dir: Path | None = data_dir_option,
363 use_global: bool = global_option,
364) -> None:
365 """Switch to embedder REF, downloading it if needed, without rebuilding the index."""
366 apply_overrides(data_dir=data_dir, use_global=use_global)
368 from lilbee.app.models import adopt_embedder
369 from lilbee.catalog.compat import UnsupportedArchError
371 try:
372 result = adopt_embedder(ref)
373 except (RuntimeError, ValueError, OSError, UnsupportedArchError) as exc:
374 if cfg.json_mode:
375 json_output({"error": str(exc)})
376 raise SystemExit(1) from None
377 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] {exc}")
378 raise SystemExit(1) from None
380 if cfg.json_mode:
381 json_output(
382 {"command": "use-embedder", "model": result.model, "status": result.status.value}
383 )
384 return
385 console.print(f"Now embedding with [{theme.ACCENT}]{result.model}[/{theme.ACCENT}].")
388def chat(
389 data_dir: Path | None = data_dir_option,
390 model: str | None = model_option,
391 use_global: bool = global_option,
392 temperature: float | None = temperature_option,
393 top_p: float | None = top_p_option,
394 top_k_sampling: int | None = top_k_sampling_option,
395 repeat_penalty: float | None = repeat_penalty_option,
396 num_ctx: int | None = num_ctx_option,
397 seed: int | None = seed_option,
398) -> None:
399 """Interactive chat loop. Press S in the TUI to sync pending documents."""
400 apply_overrides(
401 data_dir=data_dir,
402 model=model,
403 use_global=use_global,
404 temperature=temperature,
405 top_p=top_p,
406 top_k_sampling=top_k_sampling,
407 repeat_penalty=repeat_penalty,
408 num_ctx=num_ctx,
409 seed=seed,
410 )
411 # No diagnostics routing here: chat hands off to the TUI, which owns its
412 # logging (tui.log). Routing first would latch captureWarnings and leave a
413 # NOTSET cli.log handler running under the TUI.
415 if cfg.json_mode:
416 json_output({"error": "Chat requires a terminal, not --json"})
417 raise SystemExit(1)
418 if not sys.stdin.isatty() or not sys.stdout.isatty():
419 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] Chat requires a terminal.")
420 raise SystemExit(1)
421 from lilbee.cli.tui import run_tui
423 run_tui()
426def topics(
427 query: str = typer.Argument(None, help="Optional query to find related concepts."),
428 top_k: int = typer.Option(10, "--top-k", "-k", help="Number of results."),
429 data_dir: Path | None = data_dir_option,
430 use_global: bool = global_option,
431) -> None:
432 """Show top concept communities or concepts related to a query."""
433 apply_overrides(data_dir=data_dir, use_global=use_global)
435 from lilbee.retrieval.concepts import concepts_available
437 if not concepts_available():
438 msg = "Concept graph requires: pip install 'lilbee[graph]'"
439 if cfg.json_mode:
440 json_output({"error": msg})
441 raise SystemExit(1)
442 console.print(f"[{theme.ERROR}]{msg}[/{theme.ERROR}]")
443 raise SystemExit(1)
445 if not cfg.concept_graph:
446 if cfg.json_mode:
447 json_output({"error": "Concept graph is disabled (LILBEE_CONCEPT_GRAPH=false)"})
448 raise SystemExit(1)
449 console.print(
450 f"[{theme.ERROR}]Concept graph is disabled.[/{theme.ERROR}] "
451 "Enable with LILBEE_CONCEPT_GRAPH=true"
452 )
453 raise SystemExit(1)
455 if not get_services().concepts.get_graph():
456 if cfg.json_mode:
457 json_output({"error": "Concept graph not available"})
458 raise SystemExit(1)
459 console.print(f"[{theme.ERROR}]Concept graph not available.[/{theme.ERROR}]")
460 raise SystemExit(1)
462 if query:
463 _topics_for_query(query)
464 else:
465 _topics_overview(top_k)
468def _topics_for_query(query: str) -> None:
469 """Show concepts related to a query."""
470 cg = get_services().concepts
471 concepts = cg.extract_concepts(query)
472 related = cg.expand_query(query)
473 all_concepts = concepts + [r for r in related if r not in concepts]
475 if cfg.json_mode:
476 json_output({"command": "topics", "query": query, "concepts": all_concepts})
477 return
478 if not all_concepts:
479 console.print("No concepts found for this query.")
480 return
481 console.print(f"Concepts related to [{theme.ACCENT}]{query}[/{theme.ACCENT}]:")
482 for c in all_concepts:
483 console.print(f" {c}")
486def _topics_overview(top_k: int) -> None:
487 """Show top concept communities."""
488 from dataclasses import asdict
490 communities = get_services().concepts.top_communities(k=top_k)
491 if cfg.json_mode:
492 json_output({"command": "topics", "communities": [asdict(c) for c in communities]})
493 return
494 if not communities:
495 console.print("No concept communities found. Try syncing some documents first.")
496 return
497 table = Table(title="Concept Communities")
498 table.add_column("Cluster", justify="right", style=theme.MUTED)
499 table.add_column("Size", justify="right")
500 table.add_column("Top Concepts", style=theme.ACCENT)
501 for comm in communities:
502 preview = ", ".join(comm.concepts[:_TOPIC_PREVIEW_LIMIT])
503 if len(comm.concepts) > _TOPIC_PREVIEW_LIMIT:
504 preview += f" (+{len(comm.concepts) - _TOPIC_PREVIEW_LIMIT} more)"
505 table.add_row(str(comm.cluster_id), str(comm.size), preview)
506 console.print(table)