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

1"""Export and import the per-page text dataset.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6from pathlib import Path 

7from typing import NoReturn 

8 

9import typer 

10 

11from lilbee.cli import theme 

12from lilbee.cli.app import apply_overrides, console, data_dir_option, global_option 

13from lilbee.cli.helpers import json_output, sigint_cancel 

14from lilbee.core.config import cfg 

15from lilbee.runtime.cancellation import TaskCancelledError 

16 

17_export_output_argument = typer.Argument( 

18 Path("pages.parquet"), 

19 help="Output file (suffix sets the format unless --format is given).", 

20) 

21_import_dataset_argument = typer.Argument( 

22 ..., 

23 help="Dataset file to import (parquet or jsonl).", 

24) 

25_format_option = typer.Option( 

26 "", 

27 "--format", 

28 help="Dataset format: parquet or jsonl. Inferred from the file suffix when omitted.", 

29) 

30_export_source_option = typer.Option( 

31 None, 

32 "--source", 

33 help="Export only this source (default: every source).", 

34) 

35 

36 

37def _fail(message: str) -> NoReturn: 

38 """Emit *message* as an error in the active output mode and exit non-zero.""" 

39 if cfg.json_mode: 

40 json_output({"error": message}) 

41 else: 

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

43 raise SystemExit(1) 

44 

45 

46def export_cmd( 

47 output: Path = _export_output_argument, 

48 fmt: str = _format_option, 

49 source: str | None = _export_source_option, 

50 data_dir: Path | None = data_dir_option, 

51 use_global: bool = global_option, 

52) -> None: 

53 """Write a per-page {source, page, text} dataset (drops vectors).""" 

54 apply_overrides(data_dir=data_dir, use_global=use_global) 

55 from lilbee.app.dataset import DatasetError, export_to_path 

56 

57 try: 

58 with sigint_cancel() as cancel: 

59 summary = export_to_path(output, fmt, source, cancel=cancel) 

60 except TaskCancelledError: 

61 _fail("Export cancelled; the partial file was removed.") 

62 except DatasetError as exc: 

63 _fail(str(exc)) 

64 

65 if cfg.json_mode: 

66 json_output(summary.model_dump()) 

67 return 

68 console.print( 

69 f"Wrote [{theme.LABEL}]{summary.pages}[/{theme.LABEL}] pages from " 

70 f"[{theme.LABEL}]{summary.sources}[/{theme.LABEL}] source(s) to " 

71 f"[{theme.ACCENT}]{output}[/{theme.ACCENT}]" 

72 ) 

73 

74 

75def import_cmd( 

76 dataset: Path = _import_dataset_argument, 

77 fmt: str = _format_option, 

78 data_dir: Path | None = data_dir_option, 

79 use_global: bool = global_option, 

80) -> None: 

81 """Import a per-page text dataset, re-embedding it with the current model.""" 

82 apply_overrides(data_dir=data_dir, use_global=use_global) 

83 from lilbee.app.dataset import DatasetError, import_from_path 

84 

85 try: 

86 summary = asyncio.run(import_from_path(dataset, fmt)) 

87 except DatasetError as exc: 

88 _fail(str(exc)) 

89 

90 if cfg.json_mode: 

91 json_output(summary.model_dump()) 

92 return 

93 console.print( 

94 f"Imported [{theme.LABEL}]{len(summary.sources)}[/{theme.LABEL}] source(s) " 

95 f"([{theme.LABEL}]{summary.pages}[/{theme.LABEL}] pages, " 

96 f"[{theme.LABEL}]{summary.chunks}[/{theme.LABEL}] chunks)" 

97 )