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 ] 

120 

121 _SEARCH_DEBOUNCE_SECONDS = 0.08 

122 

123 def __init__( 

124 self, 

125 *, 

126 scope: PickerScope, 

127 options: list[ModelOption], 

128 ) -> None: 

129 super().__init__() 

130 self._scope: PickerScope = scope 

131 self._options = _PickerOptions(options=options) 

132 self._search_timer: Timer | None = None 

133 

134 def compose(self) -> ComposeResult: 

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

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

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

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

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

140 

141 def on_mount(self) -> None: 

142 self._refresh_list("") 

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

144 

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

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

147 ml.set_rows(self._options.to_sections(search)) 

148 

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

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

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

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

153 if self._search_timer is not None: 

154 self._search_timer.stop() 

155 search = event.value.strip() 

156 self._search_timer = self.set_timer( 

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

158 ) 

159 

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

161 def _on_search_submitted(self) -> None: 

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

163 if ml.option_count: 

164 ml.action_select() 

165 

166 @on(ModelList.Selected) 

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

168 event.stop() 

169 self.dismiss(event.row.ref) 

170 

171 def action_focus_search(self) -> None: 

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