Coverage for src/lilbee/runtime/progress/types.py: 100%

1"""Event type enums and Pydantic models for the progress protocol.""" 

2 

3from enum import StrEnum 

4 

5from pydantic import BaseModel 

6 

7 

8class EventType(StrEnum): 

9 """Progress event types emitted during sync/ingest.""" 

10 

11 FILE_START = "file_start" 

12 FILE_DONE = "file_done" 

13 BATCH_PROGRESS = "batch_progress" 

14 DONE = "done" 

15 EMBED = "embed" 

16 EXTRACT = "extract" 

17 CRAWL_START = "crawl_start" 

18 CRAWL_PAGE = "crawl_page" 

19 CRAWL_DONE = "crawl_done" 

20 SETUP_START = "setup_start" 

21 SETUP_PROGRESS = "setup_progress" 

22 SETUP_DONE = "setup_done" 

23 WIKI_PHASE = "wiki_phase" 

24 WIKI_PAGE = "wiki_page" 

25 

26 

27class SseEvent(StrEnum): 

28 """SSE event names used in the HTTP streaming protocol.""" 

29 

30 TOKEN = "token" # noqa: S105 -- SSE event name, not a credential 

31 REASONING = "reasoning" 

32 SOURCES = "sources" 

33 ERROR = "error" 

34 DONE = "done" 

35 PROGRESS = "progress" 

36 HEARTBEAT = "heartbeat" 

37 ALREADY_INGESTING = "already_ingesting" 

38 WARMING = "warming" 

39 WARM = "warm" 

40 COMPACTING = "compacting" 

41 COMPACTION = "compaction" 

42 MEMORY_EXTRACTED = "memory_extracted" 

43 GPU_STATS = "gpu_stats" 

44 

45 

46class SseErrorCode(StrEnum): 

47 """Stable ``code`` values on SSE error events for clients to branch on.""" 

48 

49 MODEL_TOO_LARGE = "model_too_large" 

50 MODEL_NOT_INSTALLED = "model_not_installed" 

51 INDEX_EMBEDDER_MISMATCH = "index_embedder_mismatch" 

52 

53 

54class FileStartEvent(BaseModel): 

55 """Emitted when a file begins ingestion.""" 

56 

57 file: str 

58 total_files: int 

59 current_file: int 

60 

61 

62class FileDoneEvent(BaseModel): 

63 """Emitted when a file finishes ingestion (success or error).""" 

64 

65 file: str 

66 status: str 

67 chunks: int 

68 

69 

70class BatchStatus(StrEnum): 

71 """Status values for BatchProgressEvent.status.""" 

72 

73 INGESTED = "ingested" 

74 SKIPPED = "skipped" 

75 FAILED = "failed" 

76 RASTERIZING = "rasterizing" 

77 

78 

79class BatchProgressEvent(BaseModel): 

80 """Emitted after each file completes during batch ingestion.""" 

81 

82 file: str 

83 status: BatchStatus 

84 current: int 

85 total: int 

86 

87 

88class ExtractEvent(BaseModel): 

89 """Emitted with page-level extraction progress. 

90 

91 OCR fires one event per page as xberg processes it, as a running count 

92 with ``total_pages == 0`` (the total is unknown mid-extraction). Extraction 

93 then fires once per file with ``page == total_pages`` so subscribers see 

94 "extracted N pages" before the embed phase ticks. 

95 """ 

96 

97 file: str 

98 page: int 

99 total_pages: int 

100 

101 

102class EmbedEvent(BaseModel): 

103 """Emitted per batch during embedding.""" 

104 

105 file: str 

106 chunk: int 

107 total_chunks: int 

108 

109 

110class CrawlStartEvent(BaseModel): 

111 """Emitted when a crawl operation begins.""" 

112 

113 url: str 

114 depth: int 

115 

116 

117# Sentinel used in CrawlPageEvent.total when the crawl's final page count is 

118# not yet known (BFS streaming, page N emitted before N+1 is discovered). 

119# Consumers (plugin, TUI, CLI) treat total <= 0 as indeterminate progress. 

120CRAWL_TOTAL_UNKNOWN = -1 

121 

122 

123class CrawlPageEvent(BaseModel): 

124 """Emitted per page during crawling.""" 

125 

126 url: str 

127 current: int 

128 total: int 

129 

130 

131class CrawlDoneEvent(BaseModel): 

132 """Emitted when a crawl operation completes.""" 

133 

134 pages_crawled: int 

135 files_written: int 

136 

137 

138class SyncDoneEvent(BaseModel): 

139 """Emitted when the sync operation completes.""" 

140 

141 added: int 

142 updated: int 

143 removed: int 

144 failed: int 

145 skipped: int = 0 

146 relocated: int = 0 

147 

148 

149class SetupStartEvent(BaseModel): 

150 """Emitted when a setup/bootstrap operation begins.""" 

151 

152 component: str 

153 size_estimate_bytes: int | None = None 

154 

155 

156class SetupProgressEvent(BaseModel): 

157 """Emitted periodically during a setup/bootstrap operation.""" 

158 

159 component: str 

160 downloaded_bytes: int 

161 total_bytes: int | None = None 

162 detail: str = "" 

163 

164 

165class SetupDoneEvent(BaseModel): 

166 """Emitted when a setup/bootstrap operation completes.""" 

167 

168 component: str 

169 success: bool 

170 error: str | None = None 

171 

172 

173class WikiPhase(StrEnum): 

174 """Stage of a wiki build or synthesis run.""" 

175 

176 EXTRACT = "extract" 

177 GENERATE = "generate" 

178 INDEX = "index" 

179 

180 

181class WikiPhaseEvent(BaseModel): 

182 """Emitted when a wiki run enters a new phase. 

183 

184 ``total`` is the number of units the phase will process (sources for a 

185 build, clusters for synthesis), 0 where the phase has no unit count. 

186 """ 

187 

188 phase: WikiPhase 

189 total: int = 0 

190 

191 

192class WikiPageEvent(BaseModel): 

193 """Emitted after each source (build) or cluster (synthesis) is written.""" 

194 

195 label: str 

196 pages: int 

197 current: int 

198 total: int