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.cli.tui import messages as msg 

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

24 CatalogRow, 

25 LocalCatalogRow, 

26) 

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

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

29 

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

31 

32 

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

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

35 if scope == "embed": 

36 return msg.MODEL_PICKER_TITLE_EMBED 

37 if scope == "vision": 

38 return msg.MODEL_PICKER_TITLE_VISION 

39 if scope == "rerank": 

40 return msg.MODEL_PICKER_TITLE_RERANK 

41 return msg.MODEL_PICKER_TITLE_CHAT 

42 

43 

44@dataclass 

45class _PickerOptions: 

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

47 

48 options: list[ModelOption] 

49 

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

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

52 return [ModelListSection(heading=None, rows=rows)] if rows else [] 

53 

54 

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

56 return LocalCatalogRow( 

57 name=option.label, 

58 task="", 

59 params="--", 

60 size="--", 

61 quant="--", 

62 downloads="--", 

63 featured=False, 

64 installed=False, 

65 sort_downloads=0, 

66 sort_size=0.0, 

67 ref=option.ref, 

68 backend="", 

69 ) 

70 

71 

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

73 if not search: 

74 return True 

75 needle = search.lower() 

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

77 

78 

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

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

81 

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

83 

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

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

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

87 ] 

88 

89 _SEARCH_DEBOUNCE_SECONDS = 0.08 

90 

91 def __init__( 

92 self, 

93 *, 

94 scope: PickerScope, 

95 options: list[ModelOption], 

96 ) -> None: 

97 super().__init__() 

98 self._scope: PickerScope = scope 

99 self._options = _PickerOptions(options=options) 

100 self._search_timer: Timer | None = None 

101 

102 def compose(self) -> ComposeResult: 

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

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

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

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

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

108 

109 def on_mount(self) -> None: 

110 self._refresh_list("") 

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

112 

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

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

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

116 

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

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

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

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

121 if self._search_timer is not None: 

122 self._search_timer.stop() 

123 search = event.value.strip() 

124 self._search_timer = self.set_timer( 

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

126 ) 

127 

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

129 def _on_search_submitted(self) -> None: 

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

131 if ml.option_count: 

132 ml.action_select() 

133 

134 @on(ModelList.Selected) 

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

136 event.stop() 

137 self.dismiss(event.row.ref) 

138 

139 def action_focus_search(self) -> None: 

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