Coverage for src/lilbee/app/dataset.py: 100%

1"""Surface-neutral export/import use cases over the per-page text dataset.""" 

2 

3from __future__ import annotations 

4 

5import threading 

6from dataclasses import dataclass 

7from pathlib import Path 

8from typing import TYPE_CHECKING 

9 

10from pydantic import BaseModel 

11 

12from lilbee.app.services import get_services 

13from lilbee.core.config import active_config 

14from lilbee.data.export import ( 

15 DatasetFormat, 

16 build_page_dataset, 

17 decode_format, 

18 deserialize_dataset, 

19 import_dataset, 

20 load_page_dataset, 

21 resolve_format, 

22 serialize_dataset, 

23 write_dataset, 

24) 

25from lilbee.data.store import EmbeddingModelMismatchError, PageTextRecord 

26from lilbee.runtime.lock import sync_running 

27from lilbee.runtime.progress import DetailedProgressCallback, noop_callback 

28 

29if TYPE_CHECKING: 

30 import pyarrow as pa 

31 

32 

33class DatasetError(Exception): 

34 """User-facing export/import failure surfaces render as-is.""" 

35 

36 

37class ExportSummary(BaseModel): 

38 """Result of a path-based export.""" 

39 

40 command: str = "export" 

41 format: str 

42 output: str 

43 pages: int 

44 sources: int 

45 

46 

47class ImportSummary(BaseModel): 

48 """Result of an import.""" 

49 

50 command: str = "import" 

51 sources: list[str] 

52 pages: int 

53 chunks: int 

54 

55 

56@dataclass 

57class ExportPayload: 

58 """In-memory export for byte transport (HTTP download).""" 

59 

60 data: bytes 

61 fmt: DatasetFormat 

62 pages: int 

63 sources: int 

64 

65 

66def require_format(value: str) -> DatasetFormat: 

67 """Decode an explicit *value* into a format; there is no path to infer from.""" 

68 if not value: 

69 raise DatasetError("format is required (parquet or jsonl)") 

70 try: 

71 return decode_format(value) 

72 except ValueError as exc: 

73 raise DatasetError(str(exc)) from None 

74 

75 

76def _build_validated(source: str | None) -> pa.Table: 

77 """Build the dataset table for *source* (or all), validating the request.""" 

78 store = get_services().store 

79 if source is not None and source not in {s["filename"] for s in store.get_sources()}: 

80 raise DatasetError(f"Source not found: {source}") 

81 table = build_page_dataset(store, source) 

82 if table.num_rows == 0: 

83 raise DatasetError("Nothing to export: the store has no indexed pages.") 

84 return table 

85 

86 

87def export_to_path( 

88 output: Path, 

89 fmt_value: str, 

90 source: str | None, 

91 *, 

92 cancel: threading.Event | None = None, 

93) -> ExportSummary: 

94 """Write the per-page dataset to *output*; format from *fmt_value* or suffix. 

95 

96 Setting *cancel* stops between row groups and removes the partial file. The 

97 table build ahead of it is a single columnar scan with no boundary to poll, 

98 so the stop lands on the write rather than the read. 

99 """ 

100 try: 

101 fmt = resolve_format(fmt_value, output) 

102 except ValueError as exc: 

103 raise DatasetError(str(exc)) from None 

104 table = _build_validated(source) 

105 write_dataset(table, output, fmt, cancel) 

106 return ExportSummary( 

107 format=str(fmt), 

108 output=str(output), 

109 pages=table.num_rows, 

110 sources=len(table.column("source").unique()), 

111 ) 

112 

113 

114def export_to_bytes(fmt_value: str, source: str | None) -> ExportPayload: 

115 """Encode the per-page dataset to bytes; empty *fmt_value* defaults to parquet.""" 

116 fmt = require_format(fmt_value) if fmt_value else DatasetFormat.PARQUET 

117 table = _build_validated(source) 

118 return ExportPayload( 

119 data=serialize_dataset(table, fmt), 

120 fmt=fmt, 

121 pages=table.num_rows, 

122 sources=len(table.column("source").unique()), 

123 ) 

124 

125 

126async def _run_import( 

127 rows: list[PageTextRecord], on_progress: DetailedProgressCallback 

128) -> ImportSummary: 

129 """Re-embed *rows* into the store under the sync mark, so a reset refuses meanwhile.""" 

130 if not rows: 

131 raise DatasetError("Dataset has no pages to import.") 

132 store = get_services().store 

133 try: 

134 async with sync_running(active_config().data_root): 

135 result = await import_dataset(store, rows, on_progress=on_progress) 

136 except EmbeddingModelMismatchError as exc: 

137 raise DatasetError(str(exc)) from None 

138 return ImportSummary(sources=result.sources, pages=result.pages, chunks=result.chunks) 

139 

140 

141async def import_from_path( 

142 path: Path, fmt_value: str, on_progress: DetailedProgressCallback = noop_callback 

143) -> ImportSummary: 

144 """Load and import a dataset file; format from *fmt_value* or suffix.""" 

145 try: 

146 fmt = resolve_format(fmt_value, path) 

147 rows = load_page_dataset(path, fmt) 

148 except ValueError as exc: 

149 raise DatasetError(str(exc)) from None 

150 return await _run_import(rows, on_progress) 

151 

152 

153async def import_from_bytes( 

154 data: bytes, fmt_value: str, on_progress: DetailedProgressCallback = noop_callback 

155) -> ImportSummary: 

156 """Decode and import dataset *data*; *fmt_value* is required (no filename).""" 

157 fmt = require_format(fmt_value) 

158 try: 

159 rows = deserialize_dataset(data, fmt) 

160 except ValueError as exc: 

161 raise DatasetError(str(exc)) from None 

162 return await _run_import(rows, on_progress)