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