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