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 

8 

9from lilbee.app.services import get_services 

10from lilbee.app.settings_map import SETTINGS_MAP 

11from lilbee.cli.tui import messages as msg 

12from lilbee.cli.tui.app import apply_active_model 

13from lilbee.providers.worker.transport import WorkerRole 

14 

15if TYPE_CHECKING: 

16 from textual.app import App 

17 from textual.widget import Widget 

18 

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

20 

21log = logging.getLogger(__name__) 

22 

23# Single source of truth for "after a model-key write, which worker pool role 

24# needs to respawn so the next call picks up the new ref?". Used by both the 

25# Settings picker dismiss path and the chat-screen model rail's button. 

26_MODEL_KEY_TO_WORKER_ROLE: dict[str, WorkerRole] = { 

27 "chat_model": WorkerRole.CHAT, 

28 "embedding_model": WorkerRole.EMBED, 

29 "reranker_model": WorkerRole.RERANK, 

30 "vision_model": WorkerRole.VISION, 

31} 

32 

33 

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

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

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

37 

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

39 if sc == scope: 

40 return key 

41 raise KeyError(scope) 

42 

43 

44def apply_model_pick( 

45 host: Widget, 

46 *, 

47 key: str, 

48 ref: str | None, 

49 on_done: Callable[[], None], 

50 reload_worker: bool = True, 

51) -> None: 

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

53 

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

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

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

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

58 Embedding-model swaps against a populated store route through a 

59 confirm modal first so the user is not surprised by the rebuild 

60 requirement. ``on_done`` runs after a successful write, never after 

61 a cancel and never after the catalog jump. 

62 

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

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

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

66 """ 

67 if ref is None: 

68 return 

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

70 

71 if ref == BROWSE_CATALOG_REF: 

72 _open_catalog_for_key(host, key) 

73 return 

74 defn = SETTINGS_MAP.get(key) 

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

76 return 

77 if key == "embedding_model" and ref and get_services().store.has_chunks(): 

78 _push_embed_swap_confirm(host, key, ref, on_done, reload_worker) 

79 return 

80 _persist(host.app, key, ref, on_done, reload_worker) 

81 

82 

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

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

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

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

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

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

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

90 model_field_to_picker_scope, 

91 picker_scope_to_task, 

92 ) 

93 

94 scope = model_field_to_picker_scope().get(key) 

95 if scope is None: 

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

97 return 

98 tab_id = TASK_TO_TAB_ID[picker_scope_to_task(scope)] 

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

100 

101 

102def _push_embed_swap_confirm( 

103 host: Widget, key: str, ref: str, on_done: Callable[[], None], reload_worker: bool 

104) -> None: 

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

106 

107 host.app.push_screen( 

108 ConfirmDialog(msg.EMBED_SWAP_CONFIRM_TITLE, msg.EMBED_SWAP_CONFIRM_MESSAGE), 

109 lambda confirmed: _on_embed_confirm(host.app, key, ref, confirmed, on_done, reload_worker), 

110 ) 

111 

112 

113def _on_embed_confirm( 

114 app: App, 

115 key: str, 

116 ref: str, 

117 confirmed: bool | None, 

118 on_done: Callable[[], None], 

119 reload_worker: bool, 

120) -> None: 

121 if not confirmed: 

122 app.notify(msg.EMBED_SWAP_CANCELLED) 

123 return 

124 _persist(app, key, ref, on_done, reload_worker) 

125 

126 

127def _persist( 

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

129) -> None: 

130 apply_active_model(app, key, ref) 

131 role = _MODEL_KEY_TO_WORKER_ROLE.get(key) 

132 if reload_worker and role is not None: 

133 get_services().reload_role(role) 

134 on_done()