Coverage for src/lilbee/cli/tui/widgets/model_pick.py: 100%

1"""Shared picker-dismiss logic for the model rail and settings screen.""" 

2 

3from __future__ import annotations 

4 

5import logging 

6from collections.abc import Callable 

7from typing import TYPE_CHECKING, cast 

8 

9from lilbee.app.services import get_services 

10from lilbee.app.settings import ( 

11 SettingsSnapshot, 

12 embed_pick_needs_rebuild, 

13 index_has_chunks, 

14 restore_settings, 

15 snapshot_settings, 

16) 

17from lilbee.app.settings_map import SETTINGS_MAP 

18from lilbee.cli.tui import messages as msg 

19from lilbee.cli.tui.app import apply_active_model 

20from lilbee.cli.tui.thread_safe import call_from_thread 

21from lilbee.core.config import cfg 

22from lilbee.providers.roles import MODEL_FIELD_TO_ROLE 

23 

24if TYPE_CHECKING: 

25 from textual.app import App 

26 from textual.widget import Widget 

27 

28 from lilbee.cli.tui.app import LilbeeApp 

29 from lilbee.cli.tui.screens.model_picker import PickerScope 

30 

31log = logging.getLogger(__name__) 

32 

33# Name for the thread worker that reloads a non-chat role off the event loop 

34# (the chat scope has its own worker in chat.py). 

35_PERSIST_WORKER_NAME = "model_swap_persist" 

36# What LanceDB raises from a read or a write, and what config.toml raises from a write. 

37STORE_ERRORS = (ValueError, OSError, RuntimeError) 

38 

39 

40def config_key_for_scope(scope: PickerScope) -> str: 

41 """Inverse of ``model_field_to_picker_scope``: scope -> config attribute name.""" 

42 from lilbee.cli.tui.screens.settings_widgets import model_field_to_picker_scope 

43 

44 for key, sc in model_field_to_picker_scope().items(): 

45 if sc == scope: 

46 return key 

47 raise KeyError(scope) 

48 

49 

50def apply_model_pick( 

51 host: Widget, 

52 *, 

53 key: str, 

54 ref: str | None, 

55 on_done: Callable[[], None], 

56 reload_worker: bool = True, 

57) -> None: 

58 """Persist a picker selection and reload the affected worker. 

59 

60 ``ref is None`` means the user cancelled (Esc); leave the field alone. 

61 ``ref == ""`` for a nullable field means the user picked the explicit 

62 "disabled" row; clear the field. ``ref == BROWSE_CATALOG_REF`` is the 

63 on-ramp action: open the Catalog focused on the role's task tab. 

64 An embedding-model swap against a populated store asks one question 

65 before anything changes: the swap needs a rebuild of the index. Yes 

66 applies the swap and starts the rebuild; No changes nothing. A swap with 

67 no question offers the rebuild after it when the write reports a stale 

68 index. ``on_done`` runs after a successful write and reload, never after a 

69 cancel, a refused write, a failed reload or the catalog jump. 

70 

71 Pass ``reload_worker=False`` when the caller resets the worker another 

72 way (the chat screen cancels its stream and resets services on a chat 

73 swap, so reloading the chat role here too would tear that work down twice). 

74 """ 

75 if ref is None: 

76 return 

77 from lilbee.cli.tui.screens.model_picker import BROWSE_CATALOG_REF 

78 

79 if ref == BROWSE_CATALOG_REF: 

80 _open_catalog_for_key(host, key) 

81 return 

82 defn = SETTINGS_MAP.get(key) 

83 if not ref and (defn is None or not defn.nullable): 

84 return 

85 if key == "embedding_model" and ref != cfg.embedding_model: 

86 try: 

87 ask = index_has_chunks() and embed_pick_needs_rebuild(ref) 

88 except STORE_ERRORS as exc: 

89 host.app.notify(msg.MODEL_ASSIGN_REJECTED.format(error=exc), severity="error") 

90 return 

91 if ask: 

92 _ask_embed_swap(host.app, key, ref, on_done, reload_worker) 

93 return 

94 _swap(host.app, key, ref, on_done, reload_worker, asked=False) 

95 

96 

97def _ask_embed_swap( 

98 app: App, key: str, ref: str, on_done: Callable[[], None], reload_worker: bool 

99) -> None: 

100 """Ask before an embedder swap that needs a rebuild; only Yes changes anything.""" 

101 from lilbee.cli.tui.widgets.confirm_dialog import ConfirmDialog 

102 

103 def _answered(rebuild: bool | None) -> None: 

104 if rebuild: 

105 _swap(app, key, ref, on_done, reload_worker, asked=True) 

106 

107 app.push_screen( 

108 ConfirmDialog(msg.CMD_REBUILD_CONFIRM_TITLE, msg.EMBED_SWAP_REBUILD_MESSAGE), 

109 _answered, 

110 ) 

111 

112 

113def _open_catalog_for_key(host: Widget, key: str) -> None: 

114 """Push CatalogScreen focused on the task tab matching the role's key.""" 

115 # circular: model_pick -> catalog (catalog imports settings_widgets which 

116 # imports model_picker, which transitively pulls in model_pick). 

117 from lilbee.cli.tui.screens.catalog import CatalogScreen 

118 from lilbee.cli.tui.screens.catalog_utils import TASK_TO_TAB_ID 

119 from lilbee.cli.tui.screens.settings_widgets import ( 

120 model_field_to_picker_scope, 

121 picker_scope_to_task, 

122 ) 

123 

124 scope = model_field_to_picker_scope().get(key) 

125 if scope is None: 

126 log.debug("Cannot open catalog for unknown model key %r", key) 

127 return 

128 tab_id = TASK_TO_TAB_ID[picker_scope_to_task(scope)] 

129 host.app.push_screen(CatalogScreen(focus_task=tab_id)) 

130 

131 

132def _swap( 

133 app: App, 

134 key: str, 

135 ref: str, 

136 on_done: Callable[[], None], 

137 reload_worker: bool, 

138 *, 

139 asked: bool, 

140) -> None: 

141 """Write the ref, reload its worker, restore the old ref on failure, then settle the rebuild. 

142 

143 A refused write ends here. A worker reload is a multi-second fleet restart, 

144 so it runs in a thread worker behind an indicator toast, and the rest runs 

145 back on the main thread. The write's verdict decides the rebuild: Yes 

146 (*asked*) starts it, and a swap with no question offers it. 

147 """ 

148 before = snapshot_settings([key]) 

149 result = apply_active_model(app, key, ref) 

150 if result is None: 

151 return 

152 lilbee_app = cast("LilbeeApp", app) 

153 

154 def _settle() -> None: 

155 on_done() 

156 if not result.reindex_required: 

157 return 

158 if asked: 

159 lilbee_app.start_rebuild() 

160 else: 

161 lilbee_app.offer_rebuild(msg.SETTINGS_REINDEX_MESSAGE) 

162 

163 role = MODEL_FIELD_TO_ROLE.get(key) 

164 if not (reload_worker and role is not None): 

165 _settle() 

166 return 

167 

168 target_role = role # narrowed to non-None; bind for the worker closure 

169 app.notify(msg.MODEL_SWAP_APPLYING) 

170 

171 def _runner() -> None: 

172 try: 

173 # wait=True runs the reload in this worker thread (already off the event 

174 # loop) so a failure is caught here and the done toast fires only once 

175 # the fleet has actually restarted, not before. 

176 get_services().reload_role(target_role, wait=True) 

177 except Exception as exc: # any reload failure becomes a toast, never a crash 

178 call_from_thread(app, _roll_back, lilbee_app, key, before, exc) 

179 return 

180 call_from_thread(app, _finish) 

181 

182 def _finish() -> None: 

183 app.notify(msg.MODEL_SWAP_DONE.format(name=ref)) 

184 _settle() 

185 

186 app.run_worker(_runner, thread=True, exit_on_error=False, name=_PERSIST_WORKER_NAME) 

187 

188 

189def _roll_back(app: LilbeeApp, key: str, before: SettingsSnapshot, error: Exception) -> None: 

190 """Put the old model and its config entry back; the toast names the model now in use.""" 

191 try: 

192 restore_settings(before) 

193 except STORE_ERRORS: 

194 log.warning("Could not restore %s after a failed reload", key, exc_info=True) 

195 in_use = getattr(cfg, key) 

196 toast = ( 

197 msg.MODEL_SWAP_FAILED.format(error=error) 

198 if in_use == before.values[key] 

199 else msg.MODEL_SWAP_STUCK.format(error=error, name=in_use) 

200 ) 

201 app.notify(toast, severity="error") 

202 app.publish_settings([key])