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 

10from rich.text import Text 

11 

12from lilbee.cli import theme 

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

14from lilbee.cli.helpers import json_output, print_prefixed, sigint_cancel 

15from lilbee.core.config import cfg 

16from lilbee.runtime.cancellation import TaskCancelledError 

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 _fail(message: str) -> NoReturn: 

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

41 if cfg.json_mode: 

42 json_output({"error": message}) 

43 else: 

44 print_prefixed(console, "Error: ", message, style=theme.ERROR) 

45 raise SystemExit(1) 

46 

47 

48def export_cmd( 

49 output: Path = _export_output_argument, 

50 fmt: str = _format_option, 

51 source: str | None = _export_source_option, 

52 data_dir: Path | None = data_dir_option, 

53 use_global: bool = global_option, 

54) -> None: 

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

56 apply_overrides(data_dir=data_dir, use_global=use_global) 

57 from lilbee.app.dataset import DatasetError, export_to_path 

58 

59 try: 

60 with sigint_cancel() as cancel: 

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

62 except TaskCancelledError: 

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

64 except DatasetError as exc: 

65 _fail(str(exc)) 

66 

67 if cfg.json_mode: 

68 json_output(summary.model_dump()) 

69 return 

70 console.print( 

71 Text.assemble( 

72 "Wrote ", 

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

74 " pages from ", 

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

76 " source(s) to ", 

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

78 ), 

79 soft_wrap=True, 

80 ) 

81 

82 

83def import_cmd( 

84 dataset: Path = _import_dataset_argument, 

85 fmt: str = _format_option, 

86 data_dir: Path | None = data_dir_option, 

87 use_global: bool = global_option, 

88) -> None: 

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

90 apply_overrides(data_dir=data_dir, use_global=use_global) 

91 from lilbee.app.dataset import DatasetError, import_from_path 

92 

93 try: 

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

95 except DatasetError as exc: 

96 _fail(str(exc)) 

97 

98 if cfg.json_mode: 

99 json_output(summary.model_dump()) 

100 return 

101 console.print( 

102 styled( 

103 "Imported ", 

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

105 " source(s) (", 

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

107 " pages, ", 

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

109 " chunks)", 

110 ) 

111 )