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_PAGE_FAILED = "crawl_page_failed" 

20 CRAWL_DONE = "crawl_done" 

21 SETUP_START = "setup_start" 

22 SETUP_PROGRESS = "setup_progress" 

23 SETUP_DONE = "setup_done" 

24 WIKI_PHASE = "wiki_phase" 

25 WIKI_PAGE = "wiki_page" 

26 

27 

28class SseEvent(StrEnum): 

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

30 

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

32 REASONING = "reasoning" 

33 SOURCES = "sources" 

34 RETRIEVAL_QUERY = "retrieval_query" 

35 ERROR = "error" 

36 DONE = "done" 

37 PROGRESS = "progress" 

38 HEARTBEAT = "heartbeat" 

39 ALREADY_INGESTING = "already_ingesting" 

40 WARMING = "warming" 

41 WARM = "warm" 

42 COMPACTING = "compacting" 

43 COMPACTION = "compaction" 

44 MEMORY_EXTRACTED = "memory_extracted" 

45 GPU_STATS = "gpu_stats" 

46 

47 

48class SseErrorCode(StrEnum): 

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

50 

51 MODEL_TOO_LARGE = "model_too_large" 

52 MODEL_NOT_INSTALLED = "model_not_installed" 

53 INDEX_EMBEDDER_MISMATCH = "index_embedder_mismatch" 

54 

55 

56class FileStartEvent(BaseModel): 

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

58 

59 file: str 

60 total_files: int 

61 current_file: int 

62 

63 

64class FileDoneEvent(BaseModel): 

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

66 

67 file: str 

68 status: str 

69 chunks: int 

70 

71 

72class BatchStatus(StrEnum): 

73 """Status values for BatchProgressEvent.status.""" 

74 

75 INGESTED = "ingested" 

76 SKIPPED = "skipped" 

77 FAILED = "failed" 

78 RASTERIZING = "rasterizing" 

79 

80 

81class BatchProgressEvent(BaseModel): 

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

83 

84 file: str 

85 status: BatchStatus 

86 current: int 

87 total: int 

88 

89 

90class ExtractEvent(BaseModel): 

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

92 

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

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

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

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

97 """ 

98 

99 file: str 

100 page: int 

101 total_pages: int 

102 

103 

104class EmbedEvent(BaseModel): 

105 """Emitted per batch during embedding.""" 

106 

107 file: str 

108 chunk: int 

109 total_chunks: int 

110 

111 

112class CrawlStartEvent(BaseModel): 

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

114 

115 url: str 

116 depth: int 

117 

118 

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

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

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

122CRAWL_TOTAL_UNKNOWN = -1 

123 

124 

125class CrawlPageEvent(BaseModel): 

126 """Emitted per page during crawling.""" 

127 

128 url: str 

129 current: int 

130 total: int 

131 

132 

133class CrawlPageFailedEvent(BaseModel): 

134 """Emitted when a crawled page yields nothing to save, with the reason why.""" 

135 

136 url: str 

137 reason: str 

138 

139 

140class CrawlDoneEvent(BaseModel): 

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

142 

143 pages_crawled: int 

144 files_written: int 

145 

146 

147class SyncDoneEvent(BaseModel): 

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

149 

150 added: int 

151 updated: int 

152 removed: int 

153 failed: int 

154 skipped: int = 0 

155 relocated: int = 0 

156 

157 

158class SetupStartEvent(BaseModel): 

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

160 

161 component: str 

162 size_estimate_bytes: int | None = None 

163 

164 

165class SetupProgressEvent(BaseModel): 

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

167 

168 component: str 

169 downloaded_bytes: int 

170 total_bytes: int | None = None 

171 detail: str = "" 

172 

173 

174class SetupDoneEvent(BaseModel): 

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

176 

177 component: str 

178 success: bool 

179 error: str | None = None 

180 

181 

182class WikiPhase(StrEnum): 

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

184 

185 EXTRACT = "extract" 

186 GENERATE = "generate" 

187 INDEX = "index" 

188 

189 

190class WikiPhaseEvent(BaseModel): 

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

192 

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

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

195 """ 

196 

197 phase: WikiPhase 

198 total: int = 0 

199 

200 

201class WikiPageEvent(BaseModel): 

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

203 

204 label: str 

205 pages: int 

206 current: int 

207 total: int