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.data.export import ( 

14 DatasetFormat, 

15 build_page_dataset, 

16 decode_format, 

17 deserialize_dataset, 

18 import_dataset, 

19 load_page_dataset, 

20 resolve_format, 

21 serialize_dataset, 

22 write_dataset, 

23) 

24from lilbee.data.store import EmbeddingModelMismatchError, PageTextRecord 

25from lilbee.runtime.progress import DetailedProgressCallback, noop_callback 

26 

27if TYPE_CHECKING: 

28 import pyarrow as pa 

29 

30 

31class DatasetError(Exception): 

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

33 

34 

35class ExportSummary(BaseModel): 

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

37 

38 command: str = "export" 

39 format: str 

40 output: str 

41 pages: int 

42 sources: int 

43 

44 

45class ImportSummary(BaseModel): 

46 """Result of an import.""" 

47 

48 command: str = "import" 

49 sources: list[str] 

50 pages: int 

51 chunks: int 

52 

53 

54@dataclass 

55class ExportPayload: 

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

57 

58 data: bytes 

59 fmt: DatasetFormat 

60 pages: int 

61 sources: int 

62 

63 

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

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

66 if not value: 

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

68 try: 

69 return decode_format(value) 

70 except ValueError as exc: 

71 raise DatasetError(str(exc)) from None 

72 

73 

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

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

76 store = get_services().store 

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

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

79 table = build_page_dataset(store, source) 

80 if table.num_rows == 0: 

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

82 return table 

83 

84 

85def export_to_path( 

86 output: Path, 

87 fmt_value: str, 

88 source: str | None, 

89 *, 

90 cancel: threading.Event | None = None, 

91) -> ExportSummary: 

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

93 

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

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

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

97 """ 

98 try: 

99 fmt = resolve_format(fmt_value, output) 

100 except ValueError as exc: 

101 raise DatasetError(str(exc)) from None 

102 table = _build_validated(source) 

103 write_dataset(table, output, fmt, cancel) 

104 return ExportSummary( 

105 format=str(fmt), 

106 output=str(output), 

107 pages=table.num_rows, 

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

109 ) 

110 

111 

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

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

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

115 table = _build_validated(source) 

116 return ExportPayload( 

117 data=serialize_dataset(table, fmt), 

118 fmt=fmt, 

119 pages=table.num_rows, 

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

121 ) 

122 

123 

124async def _run_import( 

125 rows: list[PageTextRecord], on_progress: DetailedProgressCallback 

126) -> ImportSummary: 

127 """Re-embed *rows* into the store, mapping the mismatch error for surfaces.""" 

128 if not rows: 

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

130 store = get_services().store 

131 try: 

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

133 except EmbeddingModelMismatchError as exc: 

134 raise DatasetError(str(exc)) from None 

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

136 

137 

138async def import_from_path( 

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

140) -> ImportSummary: 

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

142 try: 

143 fmt = resolve_format(fmt_value, path) 

144 rows = load_page_dataset(path, fmt) 

145 except ValueError as exc: 

146 raise DatasetError(str(exc)) from None 

147 return await _run_import(rows, on_progress) 

148 

149 

150async def import_from_bytes( 

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

152) -> ImportSummary: 

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

154 fmt = require_format(fmt_value) 

155 try: 

156 rows = deserialize_dataset(data, fmt) 

157 except ValueError as exc: 

158 raise DatasetError(str(exc)) from None 

159 return await _run_import(rows, on_progress)