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.offload import owns_ingest_pool 

26from lilbee.data.store import EmbeddingModelMismatchError, PageTextRecord 

27from lilbee.runtime.lock import sync_running 

28from lilbee.runtime.progress import DetailedProgressCallback, noop_callback 

29 

30if TYPE_CHECKING: 

31 import pyarrow as pa 

32 

33 

34class DatasetError(Exception): 

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

36 

37 

38class ExportSummary(BaseModel): 

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

40 

41 command: str = "export" 

42 format: str 

43 output: str 

44 pages: int 

45 sources: int 

46 

47 

48class ImportSummary(BaseModel): 

49 """Result of an import.""" 

50 

51 command: str = "import" 

52 sources: list[str] 

53 pages: int 

54 chunks: int 

55 

56 

57@dataclass 

58class ExportPayload: 

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

60 

61 data: bytes 

62 fmt: DatasetFormat 

63 pages: int 

64 sources: int 

65 

66 

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

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

69 if not value: 

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

71 try: 

72 return decode_format(value) 

73 except ValueError as exc: 

74 raise DatasetError(str(exc)) from None 

75 

76 

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

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

79 store = get_services().store 

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

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

82 table = build_page_dataset(store, source) 

83 if table.num_rows == 0: 

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

85 return table 

86 

87 

88def export_to_path( 

89 output: Path, 

90 fmt_value: str, 

91 source: str | None, 

92 *, 

93 cancel: threading.Event | None = None, 

94) -> ExportSummary: 

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

96 

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

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

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

100 """ 

101 try: 

102 fmt = resolve_format(fmt_value, output) 

103 except ValueError as exc: 

104 raise DatasetError(str(exc)) from None 

105 table = _build_validated(source) 

106 write_dataset(table, output, fmt, cancel) 

107 return ExportSummary( 

108 format=str(fmt), 

109 output=str(output), 

110 pages=table.num_rows, 

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

112 ) 

113 

114 

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

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

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

118 table = _build_validated(source) 

119 return ExportPayload( 

120 data=serialize_dataset(table, fmt), 

121 fmt=fmt, 

122 pages=table.num_rows, 

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

124 ) 

125 

126 

127@owns_ingest_pool 

128async def _run_import( 

129 rows: list[PageTextRecord], on_progress: DetailedProgressCallback 

130) -> ImportSummary: 

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

132 

133 The import has its own ingest pool, shut down when it ends. 

134 """ 

135 if not rows: 

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

137 store = get_services().store 

138 try: 

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

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

141 except EmbeddingModelMismatchError as exc: 

142 raise DatasetError(str(exc)) from None 

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

144 

145 

146async def import_from_path( 

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

148) -> ImportSummary: 

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

150 try: 

151 fmt = resolve_format(fmt_value, path) 

152 rows = load_page_dataset(path, fmt) 

153 except ValueError as exc: 

154 raise DatasetError(str(exc)) from None 

155 return await _run_import(rows, on_progress) 

156 

157 

158async def import_from_bytes( 

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

160) -> ImportSummary: 

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

162 fmt = require_format(fmt_value) 

163 try: 

164 rows = deserialize_dataset(data, fmt) 

165 except ValueError as exc: 

166 raise DatasetError(str(exc)) from None 

167 return await _run_import(rows, on_progress)