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 

7 

8import typer 

9from rich.text import Text 

10 

11from lilbee.cli import theme 

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

13from lilbee.cli.commands._shared import fail 

14from lilbee.cli.helpers import json_output 

15from lilbee.core.config import cfg 

16from lilbee.runtime.cancellation import TaskCancelledError, sigint_cancel 

17from lilbee.runtime.console import styled 

18 

19_export_output_argument = typer.Argument( 

20 Path("pages.parquet"), 

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

22) 

23_import_dataset_argument = typer.Argument( 

24 ..., 

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

26) 

27_format_option = typer.Option( 

28 "", 

29 "--format", 

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

31) 

32_export_source_option = typer.Option( 

33 None, 

34 "--source", 

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

36) 

37 

38 

39def export_cmd( 

40 output: Path = _export_output_argument, 

41 fmt: str = _format_option, 

42 source: str | None = _export_source_option, 

43 data_dir: Path | None = data_dir_option, 

44 use_global: bool = global_option, 

45) -> None: 

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

47 apply_overrides(data_dir=data_dir, use_global=use_global) 

48 from lilbee.app.dataset import DatasetError, export_to_path 

49 

50 try: 

51 with sigint_cancel() as cancel: 

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

53 except TaskCancelledError: 

54 fail("Export cancelled; the partial file was removed.") 

55 except (DatasetError, OSError) as exc: # OSError: output path unwritable or library unreadable 

56 fail(str(exc)) 

57 

58 if cfg.json_mode: 

59 json_output(summary.model_dump()) 

60 return 

61 console.print( 

62 Text.assemble( 

63 "Wrote ", 

64 (str(summary.pages), theme.LABEL), 

65 " pages from ", 

66 (str(summary.sources), theme.LABEL), 

67 " source(s) to ", 

68 (str(output), theme.ACCENT), 

69 ), 

70 soft_wrap=True, 

71 ) 

72 

73 

74def import_cmd( 

75 dataset: Path = _import_dataset_argument, 

76 fmt: str = _format_option, 

77 data_dir: Path | None = data_dir_option, 

78 use_global: bool = global_option, 

79) -> None: 

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

81 apply_overrides(data_dir=data_dir, use_global=use_global) 

82 from lilbee.app.dataset import DatasetError, import_from_path 

83 from lilbee.app.ingest import LIBRARY_REFUSALS 

84 

85 try: 

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

87 except (DatasetError, *LIBRARY_REFUSALS) as exc: 

88 fail(str(exc)) 

89 

90 if cfg.json_mode: 

91 json_output(summary.model_dump()) 

92 return 

93 console.print( 

94 styled( 

95 "Imported ", 

96 (str(len(summary.sources)), theme.LABEL), 

97 " source(s) (", 

98 (str(summary.pages), theme.LABEL), 

99 " pages, ", 

100 (str(summary.chunks), theme.LABEL), 

101 " chunks)", 

102 ) 

103 )