Coverage for src/lilbee/cli/commands/ingest

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_prune_ignored_option = typer.Option(

48 False,

49 "--prune-ignored",

50 help="Also drop indexed documents a .lilbeeignore now excludes. Source files are kept.",

51)

52_ocr_timeout_option = typer.Option(

53 None,

54 "--ocr-timeout",

55 help="Per-page timeout in seconds for vision OCR (default: 300, 0 = no limit).",

56)

59def _apply_ocr_overrides(ocr: bool | None, ocr_timeout: float | None) -> None:

60 """Apply --ocr/--no-ocr and --ocr-timeout CLI overrides to config.

62 The CLI is a single-shot, single-process invocation, so mutating the global

63 cfg here is safe (it mirrors ``apply_overrides`` for the data dir). The

64 daemon-shared per-request OCR override uses a ContextVar instead; see

65 ``temporary_ocr_config``.

66 """

67 if ocr is not None:

68 cfg.enable_ocr = ocr

69 if ocr_timeout is not None:

70 cfg.ocr_timeout = ocr_timeout

73_paths_argument = typer.Argument(

74 ...,

75 help="Files, directories, or URLs to add to the knowledge base.",

76)

78_force_option = typer.Option(False, "--force", "-f", help="Overwrite existing files.")

79_max_cpus_option = typer.Option(

80 None,

81 "--max-cpus",

82 min=1,

83 help="Cap the workers used to discover and hash files. Unset = auto (all available cores).",

84)

85_processes_option = typer.Option(

86 None,

87 "--processes",

88 min=0,

89 help=(

90 "Ingest worker processes, one GPU each: N explicit, 0 = auto (one per card),"

91 " 1 = this process only."

92 ),

93)

94_crawl_option = typer.Option(

95 False,

96 "--crawl",

97 help="Recursively crawl URLs (whole site by default; see --depth and --max-pages).",

98)

99_depth_option = typer.Option(

100 None,

101 "--depth",

102 help="Cap link-follow depth for --crawl. Unset = unbounded; 0 = single URL only.",

103)

104_max_pages_option = typer.Option(

105 None,

106 "--max-pages",

107 help="Cap pages for --crawl. Unset = protective default; 0 = unlimited; N = hard cap.",

108)

109_include_subdomains_option = typer.Option(

110 False,

111 "--include-subdomains",

112 help=(

113 "Allow --crawl to follow links into sibling subdomains of the start "

114 "host (e.g. en.wikipedia.org plus af.wikipedia.org). Default scopes "

115 "the crawl to the exact start host only."

116 ),

117)

118

119

120def _partition_inputs(inputs: list[str]) -> tuple[list[Path], list[str]]:

121 """Split inputs into file paths and URLs."""

122 paths: list[Path] = []

123 urls: list[str] = []

124 for inp in inputs:

125 if is_url(inp):

126 urls.append(inp)

127 else:

128 paths.append(Path(inp))

129 return paths, urls

130

131

132def _crawl_urls_blocking(

133 urls: list[str],

134 *,

135 crawl: bool,

136 depth: int | None,

137 max_pages: int | None,

138 include_subdomains: bool = False,

139) -> list[Path]:

140 """Crawl URLs synchronously (for CLI), returning paths written.

141

142 Without --crawl, each URL is fetched as a single page (depth=0).

143 With --crawl, the default is whole-site unbounded (depth=None, pages=None).

144 Explicit --depth / --max-pages override both.

145

146 Ctrl-C is handled by running the crawl through _run_crawl_with_signal_cancel,

147 which installs a signal.signal handler that sets a threading.Event passed

148 into crawl_and_save. crawl_recursive polls the event between pages so the

149 signal flows through as a clean cancel instead of asyncio.run's default

150 KeyboardInterrupt-raising (which left browser contexts mid-teardown).

151 """

152 from rich.progress import Progress, SpinnerColumn, TaskID, TextColumn

153

154 from lilbee.crawler import crawl_and_save

155 from lilbee.runtime.progress import (

156 CrawlDoneEvent,

157 CrawlPageEvent,

158 EventType,

159 ProgressEvent,

160 )

161

162 if crawl:

163 effective_depth = depth

164 effective_pages = max_pages

165 else:

166 effective_depth = 0

167 effective_pages = None

168

169 cancel_event = threading.Event()

170

171 from rich.console import Console as RichConsole

172

173 err_console = RichConsole(stderr=True)

174 all_paths: list[Path] = []

175 with Progress(

176 SpinnerColumn(),

177 TextColumn("{task.description}"),

178 transient=True,

179 console=err_console,

180 disable=cfg.json_mode,

181 ) as progress:

182 for url in urls:

183 if cancel_event.is_set():

184 break

185 ptask = progress.add_task(f"Crawling {url}...", total=None)

186 crawled: dict[str, int] = {}

187

188 def _make_callback(

189 _t: TaskID = ptask, _crawled: dict[str, int] = crawled

190 ) -> DetailedProgressCallback:

191 def on_progress(event_type: EventType, data: ProgressEvent) -> None:

192 if event_type == EventType.CRAWL_PAGE:

193 if not isinstance(data, CrawlPageEvent):

194 raise TypeError(f"Expected CrawlPageEvent, got {type(data).__name__}")

195 total_str = str(data.total) if data.total > 0 else "?"

196 progress.update(

197 _t,

198 description=f"Crawled {data.current}/{total_str}: {data.url}",

199 )

200 elif event_type == EventType.CRAWL_DONE and isinstance(data, CrawlDoneEvent):

201 _crawled["n"] = data.pages_crawled

202

203 return on_progress

204

205 paths = _run_crawl_with_signal_cancel(

206 url,

207 depth=effective_depth,

208 max_pages=effective_pages,

209 on_progress=_make_callback(),

210 cancel_event=cancel_event,

211 crawl_and_save=crawl_and_save,

212 include_subdomains=include_subdomains,

213 )

214 all_paths.extend(paths)

215 progress.update(ptask, description=f"Done: {url} ({len(paths)} pages)")

216 # No explicit cap given and the crawl filled the protective default:

217 # tell the user how to go unlimited without editing settings.

218 default_cap = cfg.crawl_max_pages or cfg.crawl_safety_max_pages

219 if crawl and max_pages is None and crawled.get("n", 0) >= default_cap:

220 err_console.print(

221 f"Stopped at the default {default_cap}-page limit; "

222 f"pass --max-pages 0 to crawl unlimited (or --max-pages N for a higher cap)."

223 )

224 return all_paths

225

226

227def _run_crawl_with_signal_cancel(

228 url: str,

229 *,

230 depth: int | None,

231 max_pages: int | None,

232 on_progress: DetailedProgressCallback,

233 cancel_event: threading.Event,

234 crawl_and_save: Callable[..., Awaitable[list[Path]]],

235 include_subdomains: bool = False,

236) -> list[Path]:

237 """Run crawl_and_save on a dedicated event loop with a SIGINT->cancel hook.

238

239 asyncio.run() installs its own SIGINT handler that raises

240 KeyboardInterrupt, which tears the crawl down ungracefully. Registering a

241 plain signal.signal handler on the main thread AND running the crawl on a

242 loop we own (instead of asyncio.run) lets Ctrl-C set our threading.Event,

243 which crawl_recursive polls between pages so it can close the stream and

244 stop dispatch cleanly.

245 """

246 import signal

247

248 # signal.signal raises ValueError when called off the main thread (e.g.

249 # under pytest-xdist workers). Skip the SIGINT hook in that case; the

250 # cancel_event can still be driven externally.

251 _on_main_thread = threading.current_thread() is threading.main_thread()

252 previous_handler = signal.getsignal(signal.SIGINT) if _on_main_thread else None

253

254 def _on_sigint(_signum: int, _frame: object) -> None:

255 # Set the cancel event that crawl_recursive polls between pages, so

256 # a Ctrl-C flows through as a clean cancel instead of asyncio.run's

257 # default KeyboardInterrupt-raising dance.

258 cancel_event.set()

259

260 if _on_main_thread:

261 signal.signal(signal.SIGINT, _on_sigint)

262 # Manage the event loop explicitly. In the CLI this runs once per process,

263 # but under pytest-xdist the same worker thread runs many tests; leaving a

264 # closed loop set as the "current" loop for the thread poisons every later

265 # asyncio.get_event_loop() call and hangs macOS 3.12/3.13 unit-test CI.

266 # Always clear the thread-current loop in finally.

267 loop = asyncio.new_event_loop()

268 try:

269 asyncio.set_event_loop(loop)

270 coro = crawl_and_save(

271 url,

272 depth=depth,

273 max_pages=max_pages,

274 on_progress=on_progress,

275 cancel=cancel_event,

276 quiet=cfg.json_mode,

277 include_subdomains=include_subdomains,

278 )

279 result: list[Path] = loop.run_until_complete(coro)

280 return result

281 finally:

282 loop.close()

283 asyncio.set_event_loop(None)

284 if _on_main_thread:

285 signal.signal(signal.SIGINT, previous_handler)

286

287

288def _cancellable_progress(

289 cancel_event: threading.Event, chain: DetailedProgressCallback

290) -> DetailedProgressCallback:

291 """Wrap *chain* so a set *cancel_event* aborts the in-flight file cooperatively.

292

293 The ingest pipeline and the per-page vision OCR loop both call the progress

294 callback between units of work; raising :class:`TaskCancelledError` there is

295 the established cooperative-cancel signal, so a Ctrl+C stops a long OCR

296 between pages instead of after the whole document.

297 """

298 from lilbee.runtime.cancellation import TaskCancelledError

299

300 def _callback(event_type: object, data: object) -> None:

301 if cancel_event.is_set():

302 raise TaskCancelledError

303 chain(event_type, data) # type: ignore[arg-type]

304

305 return _callback

306

307

308def _run_sync_with_signal_cancel(

309 *,

310 force_rebuild: bool = False,

311 retry_skipped: bool = False,

312 prune_ignored: bool = False,

313 on_progress: DetailedProgressCallback | None = None,

314) -> object:

315 """Run ``sync`` on a dedicated loop with a SIGINT->cancel hook (no traceback on Ctrl+C).

316

317 Mirrors the crawl path: a plain signal handler sets a ``threading.Event``

318 that ``sync`` polls between files and the OCR loop polls between pages, so

319 Ctrl+C aborts cleanly rather than raising KeyboardInterrupt mid-ingest.

320 """

321 import signal

322

323 from lilbee.data.ingest import sync

324 from lilbee.runtime.progress import noop_callback

325

326 # Batch ingest is a headless one-shot: skip the eager warm so services init

327 # doesn't spawn every role. With lazy per-role spawn, the sync brings up only

328 # the embed server (plus vision/chat if those steps actually run), instead of

329 # holding an idle chat server's VRAM for the whole build.

330 cfg.worker_pool_eager_start = False

331

332 cancel_event = threading.Event()

333 callback = _cancellable_progress(cancel_event, on_progress or noop_callback)

334 # signal.signal raises ValueError when called off the main thread (e.g.

335 # under pytest-xdist workers). Skip the SIGINT hook in that case; the

336 # cancel_event can still be driven externally.

337 _on_main_thread = threading.current_thread() is threading.main_thread()

338 previous_handler = signal.getsignal(signal.SIGINT) if _on_main_thread else None

339

340 def _on_sigint(_signum: int, _frame: object) -> None:

341 cancel_event.set()

342

343 if _on_main_thread:

344 signal.signal(signal.SIGINT, _on_sigint)

345 loop = asyncio.new_event_loop()

346 try:

347 asyncio.set_event_loop(loop)

348 return loop.run_until_complete(

349 sync(

350 force_rebuild=force_rebuild,

351 quiet=cfg.json_mode,

352 on_progress=callback,

353 cancel=cancel_event,

354 retry_skipped=retry_skipped,

355 prune_ignored=prune_ignored,

356 )

357 )

358 finally:

359 loop.close()

360 asyncio.set_event_loop(None)

361 if _on_main_thread:

362 signal.signal(signal.SIGINT, previous_handler)

363

364

365def sync_cmd(

366 data_dir: Path | None = data_dir_option,

367 use_global: bool = global_option,

368 ocr: bool | None = _ocr_option,

369 ocr_timeout: float | None = _ocr_timeout_option,

370 retry_skipped: bool = _retry_skipped_option,

371 prune_ignored: bool = _prune_ignored_option,

372 max_cpus: int | None = _max_cpus_option,

373 processes: int | None = _processes_option,

374) -> None:

375 """Manually trigger document sync."""

376 apply_overrides(data_dir=data_dir, use_global=use_global)

377 _apply_ocr_overrides(ocr, ocr_timeout)

378 if max_cpus is not None:

379 cfg.ingest_workers = max_cpus

380 if processes is not None:

381 cfg.ingest_processes = processes

382

383 try:

384 result = _run_sync_with_signal_cancel(

385 retry_skipped=retry_skipped, prune_ignored=prune_ignored

386 )

387 except RuntimeError as exc:

388 if cfg.json_mode:

389 json_output({"error": str(exc)})

390 raise SystemExit(1) from None

391 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] {exc}")

392 raise SystemExit(1) from None

393 if cfg.json_mode:

394 json_output(sync_result_to_json(result))

395 return

396 console.print(result)

397

398

399def rebuild(

400 data_dir: Path | None = data_dir_option,

401 use_global: bool = global_option,

402 ocr: bool | None = _ocr_option,

403 ocr_timeout: float | None = _ocr_timeout_option,

404 max_cpus: int | None = _max_cpus_option,

405 processes: int | None = _processes_option,

406) -> None:

407 """Nuke the DB and re-ingest everything from documents/."""

408 apply_overrides(data_dir=data_dir, use_global=use_global)

409 _apply_ocr_overrides(ocr, ocr_timeout)

410 if max_cpus is not None:

411 cfg.ingest_workers = max_cpus

412 if processes is not None:

413 cfg.ingest_processes = processes

414 from lilbee.data.ingest import SyncResult

415

416 try:

417 result = _run_sync_with_signal_cancel(force_rebuild=True)

418 except RuntimeError as exc:

419 if cfg.json_mode:

420 json_output({"error": str(exc)})

421 raise SystemExit(1) from None

422 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] {exc}")

423 raise SystemExit(1) from None

424 if not isinstance(result, SyncResult):

425 raise TypeError(f"Expected SyncResult, got {type(result).__name__}")

426 if cfg.json_mode:

427 json_output({"command": "rebuild", "ingested": len(result.added)})

428 return

429 console.print(f"Rebuilt: {len(result.added)} documents ingested")

430

431

432def index(

433 data_dir: Path | None = data_dir_option,

434 use_global: bool = global_option,

435) -> None:

436 """Build the search indexes now (vector ANN + full-text).

437

438 Useful before publishing a large index so downloaders get fast search

439 without waiting for it to build on first query. Forces the vector index

440 even below the auto-build threshold.

441 """

442 apply_overrides(data_dir=data_dir, use_global=use_global)

443 store = get_services().store

444 store.ensure_fts_index()

445 store.ensure_scalar_indexes()

446 built = store.ensure_vector_index(force=True)

447 if cfg.json_mode:

448 json_output({"command": "index", "vector_index": built})

449 return

450 if built:

451 console.print("Search indexes built (vector ANN + full-text).")

452 else:

453 console.print("Full-text index built; vector index needs more chunks.")

454

455

456def _validate_file_paths(file_paths: list[Path]) -> None:

457 """Exit on the first missing path; respects ``cfg.json_mode``."""

458 for fp in file_paths:

459 if fp.exists():

460 continue

461 if cfg.json_mode:

462 json_output({"error": f"Path not found: {fp}"})

463 raise SystemExit(1)

464 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] Path not found: {fp}")

465 raise SystemExit(1)

466

467

468def _crawl_urls_step(

469 urls: list[str],

470 *,

471 crawl: bool,

472 depth: int | None,

473 max_pages: int | None,

474 include_subdomains: bool,

475) -> list[Path]:

476 """Crawl URLs (or fail fast when crawler extra is missing). Returns saved paths."""

477 if not urls:

478 return []

479 from lilbee.crawler import crawler_available

480

481 if not crawler_available():

482 console.print(

483 f"[{theme.ERROR}]Web crawling requires: pip install 'lilbee[crawler]'[/{theme.ERROR}]"

484 )

485 raise SystemExit(1)

486 crawled_paths = _crawl_urls_blocking(

487 urls,

488 crawl=crawl,

489 depth=depth,

490 max_pages=max_pages,

491 include_subdomains=include_subdomains,

492 )

493 if not cfg.json_mode:

494 console.print(

495 f"[{theme.MUTED}]Crawled {len(crawled_paths)} page(s)"

496 f" from {len(urls)} URL(s)[/{theme.MUTED}]"

497 )

498 return crawled_paths

499

500

501def _add_json_mode(file_paths: list[Path], crawled_paths: list[Path], *, force: bool) -> None:

502 """Run the JSON-mode finish: register roots, sync, emit one structured result."""

503 from lilbee.data.ingest import sync

504

505 reg_result = RegisterResult()

506 if file_paths:

507 reg_result = register_sources(file_paths, force=force)

508 # Headless one-shot ingest: only the embed server is needed, so suppress eager

509 # start (matching the interactive path) instead of warming every role's VRAM.

510 cfg.worker_pool_eager_start = False

511 result = asyncio.run(sync(quiet=True))

512 json_output(

513 {

514 "command": "add",

515 "copied": reg_result.registered,

516 "skipped": reg_result.skipped,

517 "tracked": reg_result.tracked,

518 "crawled": len(crawled_paths),

519 "sync": sync_result_to_json(result),

520 }

521 )

522

523

524def add(

525 paths: list[str] = _paths_argument,

526 data_dir: Path | None = data_dir_option,

527 use_global: bool = global_option,

528 force: bool = _force_option,

529 ocr: bool | None = _ocr_option,

530 ocr_timeout: float | None = _ocr_timeout_option,

531 crawl: bool = _crawl_option,

532 depth: int | None = _depth_option,

533 max_pages: int | None = _max_pages_option,

534 include_subdomains: bool = _include_subdomains_option,

535 max_cpus: int | None = _max_cpus_option,

536 processes: int | None = _processes_option,

537) -> None:

538 """Link files or crawl URLs into the knowledge base and ingest them."""

539 apply_overrides(data_dir=data_dir, use_global=use_global)

540 _apply_ocr_overrides(ocr, ocr_timeout)

541 if max_cpus is not None:

542 cfg.ingest_workers = max_cpus

543 if processes is not None:

544 cfg.ingest_processes = processes

545

546 file_paths, urls = _partition_inputs(paths)

547 _validate_file_paths(file_paths)

548

549 try:

550 crawled_paths = _crawl_urls_step(

551 urls,

552 crawl=crawl,

553 depth=depth,

554 max_pages=max_pages,

555 include_subdomains=include_subdomains,

556 )

557

558 if cfg.json_mode:

559 _add_json_mode(file_paths, crawled_paths, force=force)

560 return

561

562 if file_paths:

563 add_paths(file_paths, console, force=force, run_sync=_run_sync_with_signal_cancel)

564 elif urls:

565 # URLs already saved; just trigger sync (Ctrl+C-cancellable)

566 result = _run_sync_with_signal_cancel()

567 console.print(result)

568 except RuntimeError as exc:

569 if cfg.json_mode:

570 json_output({"error": str(exc)})

571 raise SystemExit(1) from None

572 console.print(f"[{theme.ERROR}]Error:[/{theme.ERROR}] {exc}")

573 raise SystemExit(1) from None

574

575

576_chunks_source_argument = typer.Argument(..., help="Source name to inspect chunks for.")

577

578

579def chunks(

580 source: str = _chunks_source_argument,

581 data_dir: Path | None = data_dir_option,

582 use_global: bool = global_option,

583) -> None:

584 """Show chunks a document was split into (useful for debugging retrieval)."""

585 apply_overrides(data_dir=data_dir, use_global=use_global)

586

587 store = get_services().store

588 known = {s["filename"] for s in store.get_sources()}

589 if source not in known:

590 if cfg.json_mode:

591 json_output({"error": f"Source not found: {source}"})

592 raise SystemExit(1)

593 console.print(f"[{theme.ERROR}]Source not found:[/{theme.ERROR}] {source}")

594 raise SystemExit(1)

595

596 raw_chunks = store.get_chunks_by_source(source)

597 cleaned = sorted(

598 [clean_result(c) for c in raw_chunks],

599 key=lambda c: c.get("chunk_index", 0),

600 )

601

602 if cfg.json_mode:

603 json_output({"command": "chunks", "source": source, "chunks": cleaned})

604 return

605

606 console.print(

607 f"[{theme.LABEL}]{len(cleaned)}[/{theme.LABEL}]"

608 f" chunks from [{theme.ACCENT}]{source}[/{theme.ACCENT}]\n"

609 )

610 for c in cleaned:

611 idx = c.get("chunk_index", "?")

612 preview = c.get("chunk", "")[:CHUNK_PREVIEW_LEN]

613 if len(c.get("chunk", "")) > CHUNK_PREVIEW_LEN:

614 preview += "..."

615 console.print(f" [{idx}] {preview}")

616

617

618_remove_names_argument = typer.Argument(

619 ..., help="Source name(s), folder(s), or glob pattern(s) to remove from the knowledge base."

620)

621

622_remove_yes_option = typer.Option(

623 False, "--yes", "-y", help="Skip the confirmation prompt when a name expands to many documents."

624)

625

626

627def remove(

628 names: list[str] = _remove_names_argument,

629 data_dir: Path | None = data_dir_option,

630 use_global: bool = global_option,

631 yes: bool = _remove_yes_option,

632) -> None:

633 """Remove documents from the knowledge base by source name, folder, or glob pattern.

634

635 A folder name removes every document indexed beneath it; a glob pattern

636 (containing ``*``, ``?``, or ``[]``) removes every source it matches. Source

637 files on disk are never deleted.

638 """

639 apply_overrides(data_dir=data_dir, use_global=use_global)

640

641 known = [s["filename"] for s in get_services().store.get_sources()]

642 targets = expand_remove_targets(names, known=known)

643 expanded = sorted(set(targets)) != sorted(set(names))

644 if expanded and not yes and not cfg.json_mode:

645 # Count only what actually exists; not-found names are kept in targets.

646 removable = sum(1 for t in targets if t in set(known))

647 typer.confirm(f"Remove {removable} document(s)? Source files on disk are kept.", abort=True)

648

649 result = remove_documents_durably(names, targets=targets)

650

651 if cfg.json_mode:

652 payload: dict = {"command": "remove", "removed": result.removed}

653 if result.not_found:

654 payload["not_found"] = result.not_found

655 json_output(payload)

656 if not result.removed and result.not_found:

657 raise SystemExit(1)

658 return

659

660 for name in result.removed:

661 console.print(f"Removed [{theme.ACCENT}]{name}[/{theme.ACCENT}]")

662 for name in result.not_found:

663 console.print(f"[{theme.ERROR}]Not found:[/{theme.ERROR}] {name}")

664 if not result.removed and result.not_found:

665 raise SystemExit(1)

Coverage for src/lilbee/cli/commands/ingest_sync.py: 100%

276 statements