loader Public

Watch 0 Fork 0 Star 0
Python · 6207 bytes Raw Blame History
1 """Model selection modal with fuzzy filtering."""
2
3 from textual.app import ComposeResult
4 from textual.binding import Binding
5 from textual.containers import Vertical
6 from textual.screen import ModalScreen
7 from textual.widgets import Input, OptionList, Static
8 from textual.widgets.option_list import Option
9
10
11 class ModelSelectModal(ModalScreen[str | None]):
12 """Modal screen for selecting an Ollama model with fuzzy filtering."""
13
14 BINDINGS = [
15 Binding("escape", "cancel", "Cancel", show=True),
16 Binding("enter", "select", "Select", show=True),
17 Binding("up", "cursor_up", "Up", show=False),
18 Binding("down", "cursor_down", "Down", show=False),
19 Binding("k", "cursor_up", "Up", show=False),
20 Binding("j", "cursor_down", "Down", show=False),
21 ]
22
23 CSS = """
24 ModelSelectModal {
25 align: center middle;
26 }
27
28 #model-dialog {
29 width: 70;
30 height: 24;
31 border: heavy $primary;
32 background: $surface;
33 padding: 1 2;
34 }
35
36 #model-title {
37 text-style: bold;
38 color: $primary;
39 margin-bottom: 1;
40 text-align: center;
41 }
42
43 #model-filter {
44 margin-bottom: 1;
45 }
46
47 #model-list {
48 height: 14;
49 border: round $primary-darken-2;
50 }
51
52 #model-list > .option-list--option-highlighted {
53 background: $accent;
54 color: $text;
55 }
56
57 #model-hint {
58 text-align: center;
59 color: $text-muted;
60 margin-top: 1;
61 }
62
63 .model-current {
64 color: $success;
65 }
66 """
67
68 def __init__(
69 self,
70 models: list[dict],
71 current_model: str | None = None,
72 **kwargs,
73 ) -> None:
74 super().__init__(**kwargs)
75 self.models = models
76 self.current_model = current_model
77 self._all_options: list[tuple[str, str]] = [] # (display, value)
78
79 def compose(self) -> ComposeResult:
80 with Vertical(id="model-dialog"):
81 yield Static(
82 "[bold cyan]Select Model[/bold cyan]",
83 id="model-title",
84 )
85 yield Input(placeholder="Type to filter...", id="model-filter")
86 yield OptionList(id="model-list")
87 yield Static(
88 "[dim]↑↓/jk: navigate Enter: select Esc: cancel[/dim]",
89 id="model-hint",
90 )
91
92 def on_mount(self) -> None:
93 """Populate the model list on mount."""
94 self._build_options()
95 self._update_list("")
96 # Focus the filter input
97 self.query_one("#model-filter", Input).focus()
98
99 def _format_size(self, size_bytes: int) -> str:
100 """Format size in human-readable form."""
101 if size_bytes >= 1_000_000_000:
102 return f"{size_bytes / 1_000_000_000:.1f}GB"
103 elif size_bytes >= 1_000_000:
104 return f"{size_bytes / 1_000_000:.0f}MB"
105 else:
106 return f"{size_bytes / 1_000:.0f}KB"
107
108 def _build_options(self) -> None:
109 """Build the list of model options."""
110 # Sort by size (largest first)
111 sorted_models = sorted(self.models, key=lambda m: m.get("size", 0), reverse=True)
112
113 for model in sorted_models:
114 name = model.get("name", "unknown")
115 size = self._format_size(model.get("size", 0))
116
117 if name == self.current_model:
118 display = f"[green]● {name}[/green] [dim]({size})[/dim] [yellow]← current[/yellow]"
119 else:
120 display = f" {name} [dim]({size})[/dim]"
121
122 self._all_options.append((display, name))
123
124 def _update_list(self, filter_text: str) -> None:
125 """Update the option list based on filter."""
126 option_list = self.query_one("#model-list", OptionList)
127 option_list.clear_options()
128
129 filter_lower = filter_text.lower()
130 matched = []
131
132 for display, value in self._all_options:
133 # Fuzzy match: all filter chars must appear in order
134 if self._fuzzy_match(filter_lower, value.lower()):
135 matched.append((display, value))
136
137 for display, value in matched:
138 option_list.add_option(Option(display, id=value))
139
140 # Highlight current model if in list, otherwise first item
141 if matched:
142 # Try to find and highlight current model
143 for i, (_, value) in enumerate(matched):
144 if value == self.current_model:
145 option_list.highlighted = i
146 break
147 else:
148 option_list.highlighted = 0
149
150 def _fuzzy_match(self, pattern: str, text: str) -> bool:
151 """Simple fuzzy matching - all chars must appear in order."""
152 if not pattern:
153 return True
154 pattern_idx = 0
155 for char in text:
156 if char == pattern[pattern_idx]:
157 pattern_idx += 1
158 if pattern_idx == len(pattern):
159 return True
160 return False
161
162 def on_input_changed(self, event: Input.Changed) -> None:
163 """Filter the list as user types."""
164 if event.input.id == "model-filter":
165 self._update_list(event.value)
166
167 def on_option_list_option_selected(self, event: OptionList.OptionSelected) -> None:
168 """Handle option selection."""
169 if event.option.id:
170 self.dismiss(str(event.option.id))
171
172 def action_cancel(self) -> None:
173 """Cancel selection."""
174 self.dismiss(None)
175
176 def action_select(self) -> None:
177 """Select the highlighted option."""
178 option_list = self.query_one("#model-list", OptionList)
179 if option_list.highlighted is not None:
180 option = option_list.get_option_at_index(option_list.highlighted)
181 if option and option.id:
182 self.dismiss(str(option.id))
183
184 def action_cursor_up(self) -> None:
185 """Move cursor up in the list."""
186 option_list = self.query_one("#model-list", OptionList)
187 option_list.action_cursor_up()
188
189 def action_cursor_down(self) -> None:
190 """Move cursor down in the list."""
191 option_list = self.query_one("#model-list", OptionList)
192 option_list.action_cursor_down()