Coverage for src/lilbee/cli/tui/screens/model_picker.py: 100%

1"""Modal picker for chat / embedding model selection. 

2 

3Replaces the flat ``Select`` widgets in ModelBar. Opens with focus on a 

4search input that filters a virtualized ``ModelList`` body. Returns the 

5selected ref via ``dismiss``; the caller routes it through 

6``apply_active_model`` so the persistence path is unchanged. 

7""" 

8 

9from __future__ import annotations 

10 

11from dataclasses import dataclass 

12from typing import ClassVar, Literal 

13 

14from textual import on 

15from textual.app import ComposeResult 

16from textual.binding import Binding, BindingType 

17from textual.containers import Vertical 

18from textual.screen import ModalScreen 

19from textual.timer import Timer 

20from textual.widgets import Input, Label, Static 

21 

22from lilbee.catalog.types import ModelCompat 

23from lilbee.cli.tui import messages as msg 

24from lilbee.cli.tui.screens.catalog_utils import ( 

25 CatalogRow, 

26 LocalCatalogRow, 

27) 

28from lilbee.cli.tui.widgets.model_bar import ModelOption 

29from lilbee.cli.tui.widgets.model_list import ModelList, ModelListSection 

30 

31PickerScope = Literal["chat", "embed", "vision", "rerank"] 

32 

33# Sentinel ref returned by the picker when the user picks the "Browse catalog 

34# to download" row. apply_model_pick intercepts it and navigates to the 

35# Catalog focused on the role's task tab; it is never persisted as a real ref. 

36BROWSE_CATALOG_REF = "__browse_catalog__" 

37 

38 

39def _picker_title(scope: PickerScope) -> str: 

40 """Return the modal heading for the requested scope.""" 

41 if scope == "embed": 

42 return msg.MODEL_PICKER_TITLE_EMBED 

43 if scope == "vision": 

44 return msg.MODEL_PICKER_TITLE_VISION 

45 if scope == "rerank": 

46 return msg.MODEL_PICKER_TITLE_RERANK 

47 return msg.MODEL_PICKER_TITLE_CHAT 

48 

49 

50def _browse_catalog_row() -> CatalogRow: 

51 """The 'Browse catalog' action row appended to every picker.""" 

52 return LocalCatalogRow( 

53 name=msg.MODEL_PICKER_BROWSE_CATALOG, 

54 task="", 

55 params="--", 

56 size="--", 

57 quant="--", 

58 downloads="--", 

59 featured=False, 

60 installed=False, 

61 sort_downloads=0, 

62 sort_size=0.0, 

63 ref=BROWSE_CATALOG_REF, 

64 backend="", 

65 compat=ModelCompat.SUPPORTED, 

66 ) 

67 

68 

69@dataclass 

70class _PickerOptions: 

71 """Bridge from the dropdown's ``ModelOption`` shape to ``CatalogRow``.""" 

72 

73 options: list[ModelOption] 

74 

75 def to_sections(self, search: str) -> list[ModelListSection]: 

76 rows = [_option_to_row(o) for o in self.options if _matches(o, search)] 

77 # The browse-catalog action row is always present (filter-agnostic) so 

78 # the on-ramp is reachable even when the search box is empty AND when 

79 # it has narrowed the list to zero installed matches. 

80 rows.append(_browse_catalog_row()) 

81 return [ModelListSection(heading=None, rows=rows)] 

82 

83 

84def _option_to_row(option: ModelOption) -> CatalogRow: 

85 # Picker options are already-available models; mark SUPPORTED so the list 

86 # doesn't tag every row with the "?" unknown-compat indicator. 

87 return LocalCatalogRow( 

88 name=option.label, 

89 task="", 

90 params="--", 

91 size="--", 

92 quant="--", 

93 downloads="--", 

94 featured=False, 

95 installed=False, 

96 sort_downloads=0, 

97 sort_size=0.0, 

98 ref=option.ref, 

99 backend="", 

100 compat=ModelCompat.SUPPORTED, 

101 ) 

102 

103 

104def _matches(option: ModelOption, search: str) -> bool: 

105 if not search: 

106 return True 

107 needle = search.lower() 

108 return needle in option.label.lower() or needle in option.ref.lower() 

109 

110 

111class ModelPickerModal(ModalScreen[str | None]): 

112 """Searchable model list. Returns the selected ref or None on cancel.""" 

113 

114 CSS_PATH: ClassVar[str] = "model_picker.tcss" 

115 

116 BINDINGS: ClassVar[list[BindingType]] = [ 

117 Binding("escape", "dismiss(None)", "Cancel", show=True), 

118 Binding("slash", "focus_search", "Search", show=True), 

119 # Arrow keys reach the list even while the search input has focus, so 

120 # open -> arrows -> Enter works without knowing about Tab. When the 

121 # list itself is focused its own cursor bindings win. 

122 Binding("down", "list_cursor_down", "Next", show=False), 

123 Binding("up", "list_cursor_up", "Prev", show=False), 

124 ] 

125 

126 _SEARCH_DEBOUNCE_SECONDS = 0.08 

127 

128 def __init__( 

129 self, 

130 *, 

131 scope: PickerScope, 

132 options: list[ModelOption], 

133 ) -> None: 

134 super().__init__() 

135 self._scope: PickerScope = scope 

136 self._options = _PickerOptions(options=options) 

137 self._search_timer: Timer | None = None 

138 

139 def compose(self) -> ComposeResult: 

140 with Vertical(id="picker-root"): 

141 yield Label(_picker_title(self._scope), id="picker-title") 

142 yield Input(placeholder=msg.MODEL_PICKER_SEARCH_PLACEHOLDER, id="picker-search") 

143 yield ModelList(id="picker-list") 

144 yield Static(msg.MODEL_PICKER_HINT, id="picker-hint") 

145 

146 def on_mount(self) -> None: 

147 self._refresh_list("") 

148 self.query_one("#picker-search", Input).focus() 

149 

150 def _refresh_list(self, search: str) -> None: 

151 ml = self.query_one("#picker-list", ModelList) 

152 sections = self._options.to_sections(search) 

153 ml.set_rows(sections) 

154 # A default highlight makes Enter select the first match immediately. 

155 # The ever-present browse-catalog action row is not a match, so Enter 

156 # with no real matches stays a no-op instead of jumping to the catalog. 

157 has_match = any( 

158 row.ref != BROWSE_CATALOG_REF for section in sections for row in section.rows 

159 ) 

160 ml.highlighted = 0 if has_match else None 

161 

162 @on(Input.Changed, "#picker-search") 

163 def _on_search_changed(self, event: Input.Changed) -> None: 

164 # Debounce: rapid typing collapses to one rebuild after the user pauses. 

165 # Without this, each keystroke remounts every option (~50 ms each at 500 rows). 

166 if self._search_timer is not None: 

167 self._search_timer.stop() 

168 search = event.value.strip() 

169 self._search_timer = self.set_timer( 

170 self._SEARCH_DEBOUNCE_SECONDS, lambda: self._refresh_list(search) 

171 ) 

172 

173 @on(Input.Submitted, "#picker-search") 

174 def _on_search_submitted(self) -> None: 

175 ml = self.query_one("#picker-list", ModelList) 

176 if ml.option_count: 

177 ml.action_select() 

178 

179 @on(ModelList.Selected) 

180 def _on_model_list_selected(self, event: ModelList.Selected) -> None: 

181 event.stop() 

182 self.dismiss(event.row.ref) 

183 

184 def action_focus_search(self) -> None: 

185 self.query_one("#picker-search", Input).focus() 

186 

187 def action_list_cursor_down(self) -> None: 

188 self.query_one("#picker-list", ModelList).action_cursor_down() 

189 

190 def action_list_cursor_up(self) -> None: 

191 self.query_one("#picker-list", ModelList).action_cursor_up()