Coverage for src/lilbee/cli/launchers/warm_render.py: 100%

1"""Render chat-model warm progress from the server's SSE stream. 

2 

3A launcher consumes ``/api/warm/stream`` and drives a rich progress display so 

4the user sees a real read-phase byte bar, then an engine-load spinner, while a 

5large chat model loads, instead of a frozen line. Designed to degrade cleanly: 

6when the stream can't be opened (server still binding, a transient transport 

7error, or the endpoint absent) the caller falls back to a plain readiness poll. 

8""" 

9 

10from __future__ import annotations 

11 

12import json 

13from collections.abc import Iterator 

14 

15import httpx 

16from rich.progress import ( 

17 BarColumn, 

18 DownloadColumn, 

19 Progress, 

20 SpinnerColumn, 

21 TaskID, 

22 TextColumn, 

23 TimeElapsedColumn, 

24) 

25 

26from lilbee.catalog.formatting import display_label_for_ref 

27from lilbee.cli.app import console 

28from lilbee.providers.warm_progress import WarmPhase, WarmProgress 

29 

30_WARM_STREAM_PATH = "/api/warm/stream" 

31_SSE_DATA_PREFIX = "data:" 

32_STREAM_CONNECT_TIMEOUT_S = 5.0 

33_DEFAULT_MODEL_LABEL = "chat model" 

34 

35 

36def _model_label(ref: str | None) -> str: 

37 """The canonical UI label for the warming model, or a generic fallback.""" 

38 return display_label_for_ref(ref) if ref else _DEFAULT_MODEL_LABEL 

39 

40 

41def _iter_warm_events(base_url: str, timeout_s: float) -> Iterator[WarmProgress]: 

42 """Yield ``WarmProgress`` snapshots parsed from the SSE warm stream. 

43 

44 Raises ``httpx.HTTPError`` if the stream cannot be opened (e.g. a server 

45 without the endpoint), so the caller can fall back to a plain poll. 

46 """ 

47 timeout = httpx.Timeout(timeout_s, connect=_STREAM_CONNECT_TIMEOUT_S) 

48 with httpx.stream("GET", f"{base_url}{_WARM_STREAM_PATH}", timeout=timeout) as resp: 

49 resp.raise_for_status() 

50 for line in resp.iter_lines(): 

51 if not line.startswith(_SSE_DATA_PREFIX): 

52 continue 

53 payload = line[len(_SSE_DATA_PREFIX) :].strip() 

54 if not payload: 

55 continue 

56 try: 

57 data = json.loads(payload) 

58 except json.JSONDecodeError: 

59 continue 

60 # The terminal ``done`` event carries ``{}`` (no phase); only real 

61 # snapshots are yielded, so the renderer ignores it naturally. 

62 if isinstance(data, dict) and "phase" in data: 

63 yield WarmProgress.model_validate(data) 

64 

65 

66def _apply(progress: Progress, task_id: TaskID, snap: WarmProgress) -> None: 

67 """Reflect one warm snapshot onto the rich progress task.""" 

68 if snap.phase is WarmPhase.READING_WEIGHTS: 

69 progress.update( 

70 task_id, 

71 description=f"Reading {_model_label(snap.model_ref)} weights", 

72 total=snap.bytes_total or None, 

73 completed=snap.bytes_done, 

74 detail=snap.detail or "", 

75 ) 

76 elif snap.phase is WarmPhase.LOADING_ENGINE: 

77 # No byte signal during the VRAM load: drop to an indeterminate spinner. 

78 progress.update( 

79 task_id, 

80 description="Loading engine", 

81 total=None, 

82 detail=snap.detail or "", 

83 ) 

84 elif snap.phase is WarmPhase.READY: 

85 total = snap.bytes_total or None 

86 progress.update(task_id, description="Chat model ready", total=total, completed=total or 0) 

87 elif snap.phase is WarmPhase.ERROR: 

88 progress.update(task_id, description="Chat model failed to load", detail=snap.error or "") 

89 else: # STARTING 

90 progress.update(task_id, description="Preparing chat model", total=None, detail="") 

91 

92 

93def render_warm(base_url: str, timeout_s: float) -> bool | None: 

94 """Drive a progress display from the warm stream. 

95 

96 Returns ``True`` once the chat engine reports ready, ``False`` on an error 

97 phase or if the stream ends before ready (the caller proceeds either way), 

98 and ``None`` when the stream could not be used at all so the caller falls 

99 back to a plain readiness poll. 

100 """ 

101 columns = ( 

102 SpinnerColumn(), 

103 TextColumn("[progress.description]{task.description}"), 

104 BarColumn(), 

105 DownloadColumn(), 

106 TextColumn("{task.fields[detail]}"), 

107 TimeElapsedColumn(), 

108 ) 

109 saw_event = False 

110 try: 

111 with Progress(*columns, console=console, transient=True) as progress: 

112 task_id = progress.add_task("Preparing chat model", total=None, detail="") 

113 reached_ready = False 

114 for snap in _iter_warm_events(base_url, timeout_s): 

115 saw_event = True 

116 _apply(progress, task_id, snap) 

117 if snap.phase is WarmPhase.READY: 

118 reached_ready = True 

119 break 

120 if snap.phase is WarmPhase.ERROR: 

121 return False 

122 return reached_ready 

123 except httpx.HTTPError: 

124 # Only None means "stream never ran, poll instead". If it opened and 

125 # yielded events before dropping, report not-ready so the caller does not 

126 # double-spend the full warm budget on a second poll. 

127 return None if not saw_event else False