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

1"""Granular progress callback protocol for streaming pipeline events.""" 

2 

3from collections.abc import Callable 

4from contextvars import ContextVar 

5from enum import StrEnum 

6from typing import Any 

7 

8from pydantic import BaseModel 

9 

10 

11class EventType(StrEnum): 

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

13 

14 FILE_START = "file_start" 

15 FILE_DONE = "file_done" 

16 BATCH_PROGRESS = "batch_progress" 

17 DONE = "done" 

18 EMBED = "embed" 

19 EXTRACT = "extract" 

20 CRAWL_START = "crawl_start" 

21 CRAWL_PAGE = "crawl_page" 

22 CRAWL_DONE = "crawl_done" 

23 SETUP_START = "setup_start" 

24 SETUP_PROGRESS = "setup_progress" 

25 SETUP_DONE = "setup_done" 

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 ERROR = "error" 

35 DONE = "done" 

36 PROGRESS = "progress" 

37 HEARTBEAT = "heartbeat" 

38 ALREADY_INGESTING = "already_ingesting" 

39 

40 

41class FileStartEvent(BaseModel): 

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

43 

44 file: str 

45 total_files: int 

46 current_file: int 

47 

48 

49class FileDoneEvent(BaseModel): 

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

51 

52 file: str 

53 status: str 

54 chunks: int 

55 

56 

57class BatchProgressEvent(BaseModel): 

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

59 

60 file: str 

61 status: str 

62 current: int 

63 total: int 

64 

65 

66class ExtractEvent(BaseModel): 

67 """Emitted per page during vision OCR extraction.""" 

68 

69 file: str 

70 page: int 

71 total_pages: int 

72 

73 

74class EmbedEvent(BaseModel): 

75 """Emitted per batch during embedding.""" 

76 

77 file: str 

78 chunk: int 

79 total_chunks: int 

80 

81 

82class CrawlStartEvent(BaseModel): 

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

84 

85 url: str 

86 depth: int 

87 

88 

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

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

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

92CRAWL_TOTAL_UNKNOWN = -1 

93 

94 

95class CrawlPageEvent(BaseModel): 

96 """Emitted per page during crawling.""" 

97 

98 url: str 

99 current: int 

100 total: int 

101 

102 

103class CrawlDoneEvent(BaseModel): 

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

105 

106 pages_crawled: int 

107 files_written: int 

108 

109 

110class SyncDoneEvent(BaseModel): 

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

112 

113 added: int 

114 updated: int 

115 removed: int 

116 failed: int 

117 

118 

119class SetupStartEvent(BaseModel): 

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

121 

122 component: str 

123 size_estimate_bytes: int | None = None 

124 

125 

126class SetupProgressEvent(BaseModel): 

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

128 

129 component: str 

130 downloaded_bytes: int 

131 total_bytes: int | None = None 

132 detail: str = "" 

133 

134 

135class SetupDoneEvent(BaseModel): 

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

137 

138 component: str 

139 success: bool 

140 error: str | None = None 

141 

142 

143ProgressEvent = ( 

144 FileStartEvent 

145 | FileDoneEvent 

146 | BatchProgressEvent 

147 | ExtractEvent 

148 | EmbedEvent 

149 | SyncDoneEvent 

150 | CrawlStartEvent 

151 | CrawlPageEvent 

152 | CrawlDoneEvent 

153 | SetupStartEvent 

154 | SetupProgressEvent 

155 | SetupDoneEvent 

156) 

157 

158DetailedProgressCallback = Callable[[EventType, ProgressEvent], None] 

159 

160# When set, vision updates the batch task's description instead of creating its own bar. 

161# Value is (Progress, batch_task_id). 

162shared_progress: ContextVar[tuple[Any, Any] | None] = ContextVar("shared_progress", default=None) 

163 

164 

165def noop_callback(event_type: EventType, data: ProgressEvent) -> None: 

166 """Default no-op callback — discards all events."""