Coverage for src/lilbee/cli/commands/ingest_sync.py: 100%
279 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-09-28 17:20 +0000
1"""Sync, rebuild, add, chunks, and remove commands."""
3from __future__ import annotations
5import asyncio
6import threading
7from pathlib import Path
8from typing import TYPE_CHECKING
10import typer
11from rich.text import Text
13if TYPE_CHECKING:
14 from collections.abc import Awaitable, Callable
16 from lilbee.runtime.progress import DetailedProgressCallback
18from lilbee.app.ingest import (
19 RegisterResult,
20 expand_remove_targets,
21 register_sources,
22 removable_names,
23 remove_documents_durably,
24)
25from lilbee.app.search import clean_result
26from lilbee.app.services import get_services
27from lilbee.cli import theme
28from lilbee.cli.app import (
29 apply_overrides,
30 console,
31 data_dir_option,
32 global_option,
33)
34from lilbee.cli.commands._shared import CHUNK_PREVIEW_LEN
35from lilbee.cli.helpers import (
36 add_paths,
37 json_output,
38 print_prefixed,
39 sync_result_to_json,
40)
41from lilbee.core.config import cfg
42from lilbee.crawler import is_url
44_ocr_option = typer.Option(
45 None,
46 "--ocr/--no-ocr",
47 help=(
48 "Turn OCR on/off for scanned PDFs; off applies to every backend, vision "
49 "included, and on behaves the same as leaving this option unset."
50 ),
51)
52_retry_skipped_option = typer.Option(
53 False,
54 "--retry-skipped",
55 help="Retry files that were skipped on a previous sync (clears the failed-file markers).",
56)
57_prune_ignored_option = typer.Option(
58 False,
59 "--prune-ignored",
60 help="Also drop indexed documents a .lilbeeignore now excludes. Source files are kept.",
61)
62_ocr_timeout_option = typer.Option(
63 None,
64 "--ocr-timeout",
65 help="Per-page timeout in seconds for vision OCR (default: 300, 0 = no limit).",
66)
69def _apply_ocr_overrides(ocr: bool | None, ocr_timeout: float | None) -> None:
70 """Apply --ocr/--no-ocr and --ocr-timeout CLI overrides to config.
72 The CLI is a single-shot, single-process invocation, so mutating the global
73 cfg here is safe (it mirrors ``apply_overrides`` for the data dir). The
74 daemon-shared per-request OCR override uses a ContextVar instead; see
75 ``temporary_ocr_config``.
76 """
77 if ocr is not None:
78 cfg.enable_ocr = ocr
79 if ocr_timeout is not None:
80 cfg.ocr_timeout = ocr_timeout
83_paths_argument = typer.Argument(
84 ...,
85 help="Files, directories, or URLs to add to the knowledge base.",
86)
88_force_option = typer.Option(False, "--force", "-f", help="Overwrite existing files.")
89_max_cpus_option = typer.Option(
90 None,
91 "--max-cpus",
92 min=1,
93 help="Cap the workers used to discover and hash files. Unset = auto (all available cores).",
94)
95_processes_option = typer.Option(
96 None,
97 "--processes",
98 min=0,
99 help=(
100 "Ingest worker processes, one GPU each: N explicit, 0 = auto (one per card),"
101 " 1 = this process only."
102 ),
103)
104_crawl_option = typer.Option(
105 False,
106 "--crawl",
107 help="Recursively crawl URLs (whole site by default; see --depth and --max-pages).",
108)
109_depth_option = typer.Option(
110 None,
111 "--depth",
112 help="Cap link-follow depth for --crawl. Unset = unbounded; 0 = single URL only.",
113)
114_max_pages_option = typer.Option(
115 None,
116 "--max-pages",
117 help="Cap pages for --crawl. Unset = protective default; 0 = unlimited; N = hard cap.",
118)
119_include_subdomains_option = typer.Option(
120 False,
121 "--include-subdomains",
122 help=(
123 "Allow --crawl to follow links into sibling subdomains of the start "
124 "host (e.g. en.wikipedia.org plus af.wikipedia.org). Default scopes "
125 "the crawl to the exact start host only."
126 ),
127)
130def _partition_inputs(inputs: list[str]) -> tuple[list[Path], list[str]]:
131 """Split inputs into file paths and URLs."""
132 paths: list[Path] = []
133 urls: list[str] = []
134 for inp in inputs:
135 if is_url(inp):
136 urls.append(inp)
137 else:
138 paths.append(Path(inp))
139 return paths, urls
142def _crawl_urls_blocking(
143 urls: list[str],
144 *,
145 crawl: bool,
146 depth: int | None,
147 max_pages: int | None,
148 include_subdomains: bool = False,
149) -> list[Path]:
150 """Crawl URLs synchronously (for CLI), returning paths written.
152 Without --crawl, each URL is fetched as a single page (depth=0).
153 With --crawl, the default is whole-site unbounded (depth=None, pages=None).
154 Explicit --depth / --max-pages override both.
156 Ctrl-C is handled by running the crawl through _run_crawl_with_signal_cancel,
157 which installs a signal.signal handler that sets a threading.Event passed
158 into crawl_and_save. crawl_recursive polls the event between pages so the
159 signal flows through as a clean cancel instead of asyncio.run's default
160 KeyboardInterrupt-raising (which left browser contexts mid-teardown).
161 """
162 from rich.progress import Progress, SpinnerColumn, TaskID
164 from lilbee.crawler import crawl_and_save
165 from lilbee.runtime.progress import (
166 CrawlDoneEvent,
167 CrawlPageEvent,
168 EventType,
169 ProgressEvent,
170 )
171 from lilbee.runtime.progress.columns import literal_text_column
173 if crawl:
174 effective_depth = depth
175 effective_pages = max_pages
176 else:
177 effective_depth = 0
178 effective_pages = None
180 cancel_event = threading.Event()
182 from lilbee.runtime.console import PlainConsole
184 err_console = PlainConsole(stderr=True)
185 all_paths: list[Path] = []
186 with Progress(
187 SpinnerColumn(),
188 literal_text_column("{task.description}"),
189 transient=True,
190 console=err_console,
191 disable=cfg.json_mode,
192 ) as progress:
193 for url in urls:
194 if cancel_event.is_set():
195 break
196 ptask = progress.add_task(f"Crawling {url}...", total=None)
197 crawled: dict[str, int] = {}
199 def _make_callback(
200 _t: TaskID = ptask, _crawled: dict[str, int] = crawled
201 ) -> DetailedProgressCallback:
202 def on_progress(event_type: EventType, data: ProgressEvent) -> None:
203 if event_type == EventType.CRAWL_PAGE:
204 if not isinstance(data, CrawlPageEvent):
205 raise TypeError(f"Expected CrawlPageEvent, got {type(data).__name__}")
206 total_str = str(data.total) if data.total > 0 else "?"
207 progress.update(
208 _t,
209 description=f"Crawled {data.current}/{total_str}: {data.url}",
210 )
211 elif event_type == EventType.CRAWL_DONE and isinstance(data, CrawlDoneEvent):
212 _crawled["n"] = data.pages_crawled
214 return on_progress
216 paths = _run_crawl_with_signal_cancel(
217 url,
218 depth=effective_depth,
219 max_pages=effective_pages,
220 on_progress=_make_callback(),
221 cancel_event=cancel_event,
222 crawl_and_save=crawl_and_save,
223 include_subdomains=include_subdomains,
224 )
225 all_paths.extend(paths)
226 progress.update(ptask, description=f"Done: {url} ({len(paths)} pages)")
227 # No explicit cap given and the crawl filled the protective default:
228 # tell the user how to go unlimited without editing settings.
229 default_cap = cfg.crawl_max_pages or cfg.crawl_safety_max_pages
230 if crawl and max_pages is None and crawled.get("n", 0) >= default_cap:
231 err_console.print(
232 f"Stopped at the default {default_cap}-page limit; "
233 f"pass --max-pages 0 to crawl unlimited (or --max-pages N for a higher cap).",
234 )
235 return all_paths
238def _run_crawl_with_signal_cancel(
239 url: str,
240 *,
241 depth: int | None,
242 max_pages: int | None,
243 on_progress: DetailedProgressCallback,
244 cancel_event: threading.Event,
245 crawl_and_save: Callable[..., Awaitable[list[Path]]],
246 include_subdomains: bool = False,
247) -> list[Path]:
248 """Run crawl_and_save on a dedicated event loop with a SIGINT->cancel hook.
250 asyncio.run() installs its own SIGINT handler that raises
251 KeyboardInterrupt, which tears the crawl down ungracefully. Registering a
252 plain signal.signal handler on the main thread AND running the crawl on a
253 loop we own (instead of asyncio.run) lets Ctrl-C set our threading.Event,
254 which crawl_recursive polls between pages so it can close the stream and
255 stop dispatch cleanly.
256 """
257 import signal
259 # signal.signal raises ValueError when called off the main thread (e.g.
260 # under pytest-xdist workers). Skip the SIGINT hook in that case; the
261 # cancel_event can still be driven externally.
262 _on_main_thread = threading.current_thread() is threading.main_thread()
263 previous_handler = signal.getsignal(signal.SIGINT) if _on_main_thread else None
265 def _on_sigint(_signum: int, _frame: object) -> None:
266 # Set the cancel event that crawl_recursive polls between pages, so
267 # a Ctrl-C flows through as a clean cancel instead of asyncio.run's
268 # default KeyboardInterrupt-raising dance.
269 cancel_event.set()
271 if _on_main_thread:
272 signal.signal(signal.SIGINT, _on_sigint)
273 # Manage the event loop explicitly. In the CLI this runs once per process,
274 # but under pytest-xdist the same worker thread runs many tests; leaving a
275 # closed loop set as the "current" loop for the thread poisons every later
276 # asyncio.get_event_loop() call and hangs macOS 3.12/3.13 unit-test CI.
277 # Always clear the thread-current loop in finally.
278 loop = asyncio.new_event_loop()
279 try:
280 asyncio.set_event_loop(loop)
281 coro = crawl_and_save(
282 url,
283 depth=depth,
284 max_pages=max_pages,
285 on_progress=on_progress,
286 cancel=cancel_event,
287 quiet=cfg.json_mode,
288 include_subdomains=include_subdomains,
289 )
290 result: list[Path] = loop.run_until_complete(coro)
291 return result
292 finally:
293 loop.close()
294 asyncio.set_event_loop(None)
295 if _on_main_thread:
296 signal.signal(signal.SIGINT, previous_handler)
299def _cancellable_progress(
300 cancel_event: threading.Event, chain: DetailedProgressCallback
301) -> DetailedProgressCallback:
302 """Wrap *chain* so a set *cancel_event* aborts the in-flight file cooperatively.
304 The ingest pipeline and the per-page vision OCR loop both call the progress
305 callback between units of work; raising :class:`TaskCancelledError` there is
306 the established cooperative-cancel signal, so a Ctrl+C stops a long OCR
307 between pages instead of after the whole document.
308 """
309 from lilbee.runtime.cancellation import TaskCancelledError
311 def _callback(event_type: object, data: object) -> None:
312 if cancel_event.is_set():
313 raise TaskCancelledError
314 chain(event_type, data) # type: ignore[arg-type]
316 return _callback
319def _run_sync_with_signal_cancel(
320 *,
321 force_rebuild: bool = False,
322 retry_skipped: bool = False,
323 prune_ignored: bool = False,
324 on_progress: DetailedProgressCallback | None = None,
325) -> object:
326 """Run ``sync`` on a dedicated loop with a SIGINT->cancel hook (no traceback on Ctrl+C).
328 Mirrors the crawl path: a plain signal handler sets a ``threading.Event``
329 that ``sync`` polls between files and the OCR loop polls between pages, so
330 Ctrl+C aborts cleanly rather than raising KeyboardInterrupt mid-ingest.
331 """
332 import signal
334 from lilbee.data.ingest import sync
335 from lilbee.runtime.progress import noop_callback
337 # Batch ingest is a headless one-shot: skip the eager warm so services init
338 # doesn't spawn every role. With lazy per-role spawn, the sync brings up only
339 # the embed server (plus vision/chat if those steps actually run), instead of
340 # holding an idle chat server's VRAM for the whole build.
341 cfg.worker_pool_eager_start = False
343 cancel_event = threading.Event()
344 callback = _cancellable_progress(cancel_event, on_progress or noop_callback)
345 # signal.signal raises ValueError when called off the main thread (e.g.
346 # under pytest-xdist workers). Skip the SIGINT hook in that case; the
347 # cancel_event can still be driven externally.
348 _on_main_thread = threading.current_thread() is threading.main_thread()
349 previous_handler = signal.getsignal(signal.SIGINT) if _on_main_thread else None
351 def _on_sigint(_signum: int, _frame: object) -> None:
352 cancel_event.set()
354 if _on_main_thread:
355 signal.signal(signal.SIGINT, _on_sigint)
356 loop = asyncio.new_event_loop()
357 try:
358 asyncio.set_event_loop(loop)
359 return loop.run_until_complete(
360 sync(
361 force_rebuild=force_rebuild,
362 quiet=cfg.json_mode,
363 on_progress=callback,
364 cancel=cancel_event,
365 retry_skipped=retry_skipped,
366 prune_ignored=prune_ignored,
367 )
368 )
369 finally:
370 loop.close()
371 asyncio.set_event_loop(None)
372 if _on_main_thread:
373 signal.signal(signal.SIGINT, previous_handler)
376def sync_cmd(
377 data_dir: Path | None = data_dir_option,
378 use_global: bool = global_option,
379 ocr: bool | None = _ocr_option,
380 ocr_timeout: float | None = _ocr_timeout_option,
381 retry_skipped: bool = _retry_skipped_option,
382 prune_ignored: bool = _prune_ignored_option,
383 max_cpus: int | None = _max_cpus_option,
384 processes: int | None = _processes_option,
385) -> None:
386 """Manually trigger document sync."""
387 apply_overrides(data_dir=data_dir, use_global=use_global)
388 _apply_ocr_overrides(ocr, ocr_timeout)
389 if max_cpus is not None:
390 cfg.ingest_workers = max_cpus
391 if processes is not None:
392 cfg.ingest_processes = processes
394 try:
395 result = _run_sync_with_signal_cancel(
396 retry_skipped=retry_skipped, prune_ignored=prune_ignored
397 )
398 except RuntimeError as exc:
399 if cfg.json_mode:
400 json_output({"error": str(exc)})
401 raise SystemExit(1) from None
402 print_prefixed(console, "Error: ", exc, style=theme.ERROR)
403 raise SystemExit(1) from None
404 if cfg.json_mode:
405 json_output(sync_result_to_json(result))
406 return
407 console.print(result)
410def rebuild(
411 data_dir: Path | None = data_dir_option,
412 use_global: bool = global_option,
413 ocr: bool | None = _ocr_option,
414 ocr_timeout: float | None = _ocr_timeout_option,
415 max_cpus: int | None = _max_cpus_option,
416 processes: int | None = _processes_option,
417) -> None:
418 """Nuke the DB and re-ingest everything from documents/."""
419 apply_overrides(data_dir=data_dir, use_global=use_global)
420 _apply_ocr_overrides(ocr, ocr_timeout)
421 if max_cpus is not None:
422 cfg.ingest_workers = max_cpus
423 if processes is not None:
424 cfg.ingest_processes = processes
425 from lilbee.data.ingest import SyncResult
427 try:
428 result = _run_sync_with_signal_cancel(force_rebuild=True)
429 except RuntimeError as exc:
430 if cfg.json_mode:
431 json_output({"error": str(exc)})
432 raise SystemExit(1) from None
433 print_prefixed(console, "Error: ", exc, style=theme.ERROR)
434 raise SystemExit(1) from None
435 if not isinstance(result, SyncResult):
436 raise TypeError(f"Expected SyncResult, got {type(result).__name__}")
437 if cfg.json_mode:
438 json_output({"command": "rebuild", "ingested": len(result.added)})
439 return
440 console.print(f"Rebuilt: {len(result.added)} documents ingested")
443def index(
444 data_dir: Path | None = data_dir_option,
445 use_global: bool = global_option,
446) -> None:
447 """Build the search indexes now (vector ANN + full-text).
449 Useful before publishing a large index so downloaders get fast search
450 without waiting for it to build on first query. Forces the vector index
451 even below the auto-build threshold.
452 """
453 apply_overrides(data_dir=data_dir, use_global=use_global)
454 store = get_services().store
455 store.ensure_fts_index()
456 store.ensure_scalar_indexes()
457 built = store.ensure_vector_index(force=True)
458 if cfg.json_mode:
459 json_output({"command": "index", "vector_index": built})
460 return
461 if built:
462 console.print("Search indexes built (vector ANN + full-text).")
463 else:
464 console.print("Full-text index built; vector index needs more chunks.")
467def _validate_file_paths(file_paths: list[Path]) -> None:
468 """Exit on the first missing path; respects ``cfg.json_mode``."""
469 for fp in file_paths:
470 if fp.exists():
471 continue
472 if cfg.json_mode:
473 json_output({"error": f"Path not found: {fp}"})
474 raise SystemExit(1)
475 print_prefixed(console, "Error: ", f"Path not found: {fp}", style=theme.ERROR)
476 raise SystemExit(1)
479def _crawl_urls_step(
480 urls: list[str],
481 *,
482 crawl: bool,
483 depth: int | None,
484 max_pages: int | None,
485 include_subdomains: bool,
486) -> list[Path]:
487 """Crawl URLs (or fail fast when crawler extra is missing). Returns saved paths."""
488 if not urls:
489 return []
490 from lilbee.crawler import crawler_available
492 if not crawler_available():
493 console.print(
494 "Web crawling requires: pip install 'lilbee[crawler]'",
495 style=theme.ERROR,
496 soft_wrap=True,
497 )
498 raise SystemExit(1)
499 crawled_paths = _crawl_urls_blocking(
500 urls,
501 crawl=crawl,
502 depth=depth,
503 max_pages=max_pages,
504 include_subdomains=include_subdomains,
505 )
506 if not cfg.json_mode:
507 console.print(
508 f"Crawled {len(crawled_paths)} page(s) from {len(urls)} URL(s)",
509 style=theme.MUTED,
510 )
511 return crawled_paths
514def _add_json_mode(file_paths: list[Path], crawled_paths: list[Path], *, force: bool) -> None:
515 """Run the JSON-mode finish: register roots, sync, emit one structured result."""
516 from lilbee.data.ingest import sync
518 reg_result = RegisterResult()
519 if file_paths:
520 reg_result = register_sources(file_paths, force=force)
521 # Headless one-shot ingest: only the embed server is needed, so suppress eager
522 # start (matching the interactive path) instead of warming every role's VRAM.
523 cfg.worker_pool_eager_start = False
524 # A sync is a whole-vault pass; run it only when something named reached the corpus.
525 result = asyncio.run(sync(quiet=True)) if reg_result.reached_corpus or crawled_paths else None
526 json_output(
527 {
528 "command": "add",
529 "copied": reg_result.registered,
530 "name_taken": reg_result.name_taken,
531 "overlapping": reg_result.overlapping,
532 "tracked": reg_result.tracked,
533 "refused": reg_result.refused,
534 "crawled": len(crawled_paths),
535 "sync": None if result is None else sync_result_to_json(result),
536 }
537 )
540def add(
541 paths: list[str] = _paths_argument,
542 data_dir: Path | None = data_dir_option,
543 use_global: bool = global_option,
544 force: bool = _force_option,
545 ocr: bool | None = _ocr_option,
546 ocr_timeout: float | None = _ocr_timeout_option,
547 crawl: bool = _crawl_option,
548 depth: int | None = _depth_option,
549 max_pages: int | None = _max_pages_option,
550 include_subdomains: bool = _include_subdomains_option,
551 max_cpus: int | None = _max_cpus_option,
552 processes: int | None = _processes_option,
553) -> None:
554 """Link files or crawl URLs into the knowledge base and ingest them."""
555 apply_overrides(data_dir=data_dir, use_global=use_global)
556 _apply_ocr_overrides(ocr, ocr_timeout)
557 if max_cpus is not None:
558 cfg.ingest_workers = max_cpus
559 if processes is not None:
560 cfg.ingest_processes = processes
562 file_paths, urls = _partition_inputs(paths)
563 _validate_file_paths(file_paths)
565 try:
566 crawled_paths = _crawl_urls_step(
567 urls,
568 crawl=crawl,
569 depth=depth,
570 max_pages=max_pages,
571 include_subdomains=include_subdomains,
572 )
574 if cfg.json_mode:
575 _add_json_mode(file_paths, crawled_paths, force=force)
576 return
578 if file_paths:
579 add_paths(file_paths, console, force=force, run_sync=_run_sync_with_signal_cancel)
580 elif urls:
581 # URLs already saved; just trigger sync (Ctrl+C-cancellable)
582 result = _run_sync_with_signal_cancel()
583 console.print(result)
584 except RuntimeError as exc:
585 if cfg.json_mode:
586 json_output({"error": str(exc)})
587 raise SystemExit(1) from None
588 print_prefixed(console, "Error: ", exc, style=theme.ERROR)
589 raise SystemExit(1) from None
592_chunks_source_argument = typer.Argument(..., help="Source name to inspect chunks for.")
595def chunks(
596 source: str = _chunks_source_argument,
597 data_dir: Path | None = data_dir_option,
598 use_global: bool = global_option,
599) -> None:
600 """Show chunks a document was split into (useful for debugging retrieval)."""
601 apply_overrides(data_dir=data_dir, use_global=use_global)
603 store = get_services().store
604 known = {s["filename"] for s in store.get_sources()}
605 if source not in known:
606 if cfg.json_mode:
607 json_output({"error": f"Source not found: {source}"})
608 raise SystemExit(1)
609 print_prefixed(console, "Source not found: ", source, style=theme.ERROR)
610 raise SystemExit(1)
612 raw_chunks = store.get_chunks_by_source(source)
613 cleaned = sorted(
614 [clean_result(c) for c in raw_chunks],
615 key=lambda c: c.get("chunk_index", 0),
616 )
618 if cfg.json_mode:
619 json_output({"command": "chunks", "source": source, "chunks": cleaned})
620 return
622 console.print(
623 Text.assemble(
624 (str(len(cleaned)), theme.LABEL), " chunks from ", (source, theme.ACCENT), "\n"
625 ),
626 soft_wrap=True,
627 )
628 for c in cleaned:
629 idx = c.get("chunk_index", "?")
630 preview = c.get("chunk", "")[:CHUNK_PREVIEW_LEN]
631 if len(c.get("chunk", "")) > CHUNK_PREVIEW_LEN:
632 preview += "..."
633 console.print(Text.assemble(f" [{idx}] ", preview), soft_wrap=True)
636_remove_names_argument = typer.Argument(
637 ..., help="Source name(s), folder(s), or glob pattern(s) to remove from the knowledge base."
638)
640_remove_yes_option = typer.Option(
641 False, "--yes", "-y", help="Skip the confirmation prompt when a name expands to many documents."
642)
645def remove(
646 names: list[str] = _remove_names_argument,
647 data_dir: Path | None = data_dir_option,
648 use_global: bool = global_option,
649 yes: bool = _remove_yes_option,
650) -> None:
651 """Remove documents from the knowledge base by source name, folder, or glob pattern.
653 A folder name removes every document indexed beneath it; a glob pattern
654 (containing ``*``, ``?``, or ``[]``) removes every source it matches. A file
655 held out because its ingestion failed can be named too; it stays out of later
656 syncs. Source files on disk are never deleted.
657 """
658 apply_overrides(data_dir=data_dir, use_global=use_global)
659 # Remove only touches the store, never the engine; skip the eager fleet warm.
660 cfg.worker_pool_eager_start = False
662 known = removable_names()
663 targets = expand_remove_targets(names, known=known)
664 expanded = sorted(set(targets)) != sorted(set(names))
665 if expanded and not yes and not cfg.json_mode:
666 # Count only what actually exists; not-found names are kept in targets.
667 removable = sum(1 for t in targets if t in set(known))
668 typer.confirm(f"Remove {removable} document(s)? Source files on disk are kept.", abort=True)
670 result = remove_documents_durably(names, targets=targets)
672 if cfg.json_mode:
673 payload: dict = {"command": "remove", "removed": result.removed}
674 if result.not_found:
675 payload["not_found"] = result.not_found
676 json_output(payload)
677 if not result.removed and result.not_found:
678 raise SystemExit(1)
679 return
681 for name in result.removed:
682 console.print(Text.assemble("Removed ", (name, theme.ACCENT)), soft_wrap=True)
683 for name in result.not_found:
684 print_prefixed(console, "Not found: ", name, style=theme.ERROR)
685 if not result.removed and result.not_found:
686 raise SystemExit(1)