Coverage for src / lilbee / cli / tui / widgets / model_bar.py: 100%
331 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-06-28 01:01 +0000
1"""Model bar: a horizontal band of the four role pickers plus the Search/Chat toggle."""
3from __future__ import annotations
5import contextlib
6import logging
7from pathlib import Path
8from typing import TYPE_CHECKING, ClassVar, NamedTuple
10if TYPE_CHECKING:
11 from lilbee.cli.tui.app import LilbeeApp
12 from lilbee.cli.tui.screens.model_picker import PickerScope
13 from lilbee.modelhub.registry import ModelRegistry
15from textual import events, work
16from textual.app import ComposeResult
17from textual.binding import Binding, BindingType
18from textual.containers import Horizontal
19from textual.widget import Widget
20from textual.widgets import Static
22from lilbee.app.services import get_services, reset_services
23from lilbee.catalog import clean_display_name, display_label_for_ref, extract_quant
24from lilbee.catalog.types import ModelTask
25from lilbee.cli.tui import messages as msg
26from lilbee.cli.tui.app import apply_setting
27from lilbee.cli.tui.pill import pill
28from lilbee.cli.tui.screens.settings_widgets import model_field_to_picker_scope
29from lilbee.cli.tui.thread_safe import call_from_thread
30from lilbee.cli.tui.widgets.model_pick import apply_model_pick, config_key_for_scope
31from lilbee.core.config import cfg
32from lilbee.core.config.enums import ChatMode
33from lilbee.providers.model_ref import format_remote_ref, parse_model_ref
34from lilbee.providers.sdk_backend import PROVIDER_KEYS
35from lilbee.retrieval.embedder import is_model_available
37log = logging.getLogger(__name__)
39_MMPROJ_MARKER = "mmproj"
41# Routing-name -> display-label map derived from PROVIDER_KEYS. Any new
42# entry added there lights up the warning without further changes here.
43_CLOUD_PROVIDER_LABELS: dict[str, str] = {name: label for name, _, _, label in PROVIDER_KEYS}
46def _cloud_provider_label(chat_model: str) -> str | None:
47 """Return the provider display label for cloud-routed models, else None."""
48 if not chat_model:
49 return None
50 ref = parse_model_ref(chat_model)
51 if not ref.is_api:
52 return None
53 return _CLOUD_PROVIDER_LABELS.get(ref.provider)
56class ModelOption(NamedTuple):
57 """A selectable model with display label and config ref."""
59 label: str # human-readable name for the dropdown
60 ref: str # canonical ref persisted to config
63def _is_mmproj(name: str) -> bool:
64 """Return True if a model name refers to an mmproj projection file."""
65 return _MMPROJ_MARKER in name.lower()
68def classify_installed_models_full() -> dict[ModelTask, list[ModelOption]]:
69 """Classify installed models into per-task lists, dropping mmproj entries."""
70 buckets: dict[ModelTask, list[ModelOption]] = {task: [] for task in ModelTask}
71 seen: set[str] = set()
73 _collect_native_models(buckets, seen)
74 _collect_remote_models(buckets, seen)
75 _collect_api_models(buckets, seen)
77 return {task: sorted(opts, key=lambda o: o.ref) for task, opts in buckets.items()}
80def _lookup_bucket(
81 buckets: dict[ModelTask, list[ModelOption]], task: str, ref: str
82) -> list[ModelOption] | None:
83 """Return the bucket for *task*, or None if it is not a known ModelTask."""
84 try:
85 key = ModelTask(task)
86 except ValueError:
87 log.debug("dropping %r with unknown task %r", ref, task)
88 return None
89 return buckets.get(key)
92def _native_label(hf_repo: str, gguf_filename: str, repo_count: int) -> str:
93 """Build the picker label, appending the quant suffix only on collision."""
94 base = clean_display_name(hf_repo)
95 if repo_count <= 1:
96 return base
97 quant = extract_quant(gguf_filename)
98 return f"{base} ({quant})" if quant else base
101def _has_vision_sidecar(registry: ModelRegistry, ref: str) -> bool:
102 """Return True if *ref* resolves to a model with an adjacent ``*mmproj*.gguf`` file.
104 Models like ``google/gemma-3-12b-it`` carry their vision capability in
105 a sibling ``mmproj`` GGUF; without checking the file system, the
106 ref's name alone gives no signal that the model is multimodal, so the
107 vision picker would silently miss it.
108 """
109 try:
110 path = registry.resolve(ref)
111 except (KeyError, ValueError):
112 return False
113 return any(path.parent.glob("*mmproj*.gguf"))
116def _collect_native_models(buckets: dict[ModelTask, list[ModelOption]], seen: set[str]) -> None:
117 """Add native registry models to buckets."""
118 try:
119 from lilbee.modelhub.registry import ModelRegistry
121 registry = ModelRegistry(cfg.models_dir)
122 manifests = registry.list_installed()
123 repo_counts: dict[str, int] = {}
124 for m in manifests:
125 repo_counts[m.hf_repo] = repo_counts.get(m.hf_repo, 0) + 1
127 from lilbee.modelhub.model_manager.discovery import reclassify_by_name
129 for manifest in manifests:
130 ref = manifest.ref
131 if _is_mmproj(manifest.gguf_filename) or ref in seen:
132 continue
133 task = reclassify_by_name(ref, manifest.task)
134 label = _native_label(
135 manifest.hf_repo, manifest.gguf_filename, repo_counts[manifest.hf_repo]
136 )
137 primary_bucket = _lookup_bucket(buckets, task, ref)
138 if primary_bucket is None:
139 continue
140 seen.add(ref)
141 primary_bucket.append(ModelOption(label=label, ref=ref))
142 # If the model has an mmproj sidecar it is also vision-capable.
143 # Surface it under the vision picker too without dropping its
144 # primary classification, so a chat model with vision (e.g.
145 # gemma-3 with mmproj) shows up in both pickers.
146 if task != ModelTask.VISION and _has_vision_sidecar(registry, ref):
147 buckets[ModelTask.VISION].append(ModelOption(label=label, ref=ref))
148 except Exception:
149 log.debug("Could not read native model registry", exc_info=True)
152def _collect_remote_models(buckets: dict[ModelTask, list[ModelOption]], seen: set[str]) -> None:
153 """Add remote (Ollama / OpenAI-compatible) models, prefixed for routing.
155 Skipped when the litellm extra is not installed -- surfacing a model
156 the SDK cannot route is a guaranteed runtime error.
157 """
158 from lilbee.providers.litellm_sdk import litellm_available
160 if not litellm_available():
161 return
162 try:
163 from lilbee.modelhub.model_manager import classify_all_remote_models
165 for model in classify_all_remote_models():
166 # Skip backend rows with a blank model name so the picker
167 # doesn't render an empty " (Ollama)" row.
168 if not model.name.strip():
169 continue
170 ref = format_remote_ref(model.name, model.provider)
171 if ref in seen or _is_mmproj(model.name):
172 continue
173 bucket = _lookup_bucket(buckets, model.task, ref)
174 if bucket is None:
175 continue
176 seen.add(ref)
177 label = f"{model.name} ({model.provider})"
178 bucket.append(ModelOption(label=label, ref=ref))
179 except Exception:
180 log.debug("Could not classify remote models", exc_info=True)
183def _collect_api_models(buckets: dict[ModelTask, list[ModelOption]], seen: set[str]) -> None:
184 """Add frontier API chat models. Skipped without litellm (cannot route)."""
185 from lilbee.providers.litellm_sdk import litellm_available
187 if not litellm_available():
188 return
189 try:
190 from lilbee.modelhub.model_manager import discover_api_models
192 # API discovery returns only chat-capable refs; revisit if providers
193 # expose embedding/vision/rerank.
194 for display_name, models in discover_api_models().items():
195 for model in models:
196 qualified = format_remote_ref(model.name, model.provider)
197 if qualified in seen:
198 continue
199 seen.add(qualified)
200 label = f"{model.name} ({display_name})"
201 buckets[ModelTask.CHAT].append(ModelOption(label=label, ref=qualified))
202 except Exception:
203 log.debug("Could not discover API models", exc_info=True)
206_CHAT_MODE_TOGGLE_ID = "chat-mode-toggle"
207_CHAT_MODE_SEARCH_PILL_ID = "chat-mode-search"
208_CHAT_MODE_CHAT_PILL_ID = "chat-mode-chat"
209_CHAT_MODE_PILL_CLASS = "chat-mode-pill"
210_CHAT_MODE_DISABLED_CLASS = "-disabled"
211_CHAT_MODE_ACTIVE_CLASS = "-active"
214_SCOPE_TO_TOOLTIP: dict[str, str] = {
215 "chat": msg.MODEL_PICKER_CHAT_TOOLTIP,
216 "embed": msg.MODEL_PICKER_EMBED_TOOLTIP,
217 "vision": msg.MODEL_PICKER_VISION_TOOLTIP,
218 "rerank": msg.MODEL_PICKER_RERANK_TOOLTIP,
219}
221_CSS_FILE = Path(__file__).parent / "model_bar.tcss"
223_CLOUD_WARNING_ID = "model-bar-cloud-warning"
225_SCOPE_TO_LABEL: dict[str, str] = {
226 "chat": msg.MODEL_BAR_CHAT_LABEL,
227 "embed": msg.MODEL_BAR_EMBED_LABEL,
228 "vision": msg.MODEL_BAR_VISION_LABEL,
229 "rerank": msg.MODEL_BAR_RERANK_LABEL,
230}
232# Per-role pill colors (background, foreground) when the role is active. Chat and
233# Embed mirror the original bar; Vision and Rerank get their own accent hues.
234_SCOPE_PILL_COLORS: dict[str, tuple[str, str]] = {
235 "chat": ("$primary", "$text"),
236 "embed": ("$secondary", "$text"),
237 "vision": ("#bc8cff", "$text"),
238 "rerank": ("#f0883e", "$text"),
239}
241# Muted pill for an optional role that is currently off.
242_OFF_PILL_COLORS: tuple[str, str] = ("$surface-lighten-2", "$text-muted")
245class ModelPickerButton(Static, can_focus=True):
246 """Pill button that opens a ModelPickerModal scoped to one of the four roles."""
248 BINDINGS: ClassVar[list[BindingType]] = [
249 Binding("enter", "open_picker", "Pick model", show=False),
250 Binding("space", "open_picker", "Pick model", show=False),
251 ]
253 def __init__(self, *, scope: PickerScope, button_id: str) -> None:
254 super().__init__(id=button_id)
255 self._scope: PickerScope = scope
256 self._key: str = config_key_for_scope(scope)
257 self._options: list[ModelOption] = []
258 self.tooltip = _SCOPE_TO_TOOLTIP[scope]
260 def on_mount(self) -> None:
261 self._refresh()
263 def set_options(self, options: list[ModelOption]) -> None:
264 """Update the options pool. Repaints the label from cfg."""
265 self._options = options
266 if self.is_mounted:
267 self._refresh()
269 def _refresh(self) -> None:
270 # Only optional roles (vision/rerank) can be empty; chat/embed are
271 # non-nullable, so an empty ref means the role is off, not a model
272 # called "(none)".
273 ref = getattr(cfg, self._key)
274 label = (display_label_for_ref(ref) or ref) if ref else msg.MODEL_BAR_DISABLED
275 self.update(label)
277 def repaint(self) -> None:
278 """Public entry for a parent container to repaint the label from cfg."""
279 self._refresh()
281 def on_click(self, event: events.Click) -> None:
282 event.stop()
283 self.open_picker()
285 def action_open_picker(self) -> None:
286 self.open_picker()
288 def _is_nullable(self) -> bool:
289 from lilbee.app.settings_map import SETTINGS_MAP
291 defn = SETTINGS_MAP.get(self._key)
292 return defn is not None and defn.nullable
294 def open_picker(self) -> None:
295 # Lazy import: model_picker imports ModelOption from this module.
296 from lilbee.cli.tui.screens.model_picker import ModelPickerModal
298 options = list(self._options)
299 # Optional role that's on: offer an explicit "turn off" action in the
300 # modal (ref "" disables it) so disabling isn't only a pill click.
301 if self._is_nullable() and getattr(cfg, self._key):
302 options.append(ModelOption(label=msg.MODEL_PICKER_TURN_OFF, ref=""))
303 modal = ModelPickerModal(scope=self._scope, options=options)
304 self.app.push_screen(modal, self._on_picker_dismissed)
306 def _on_picker_dismissed(self, ref: str | None) -> None:
307 if ref is not None and ref == getattr(cfg, self._key):
308 return
309 # Chat swaps reset services in _commit_after_change -> apply_model_change,
310 # so the helper must not also reload the chat worker (double teardown).
311 apply_model_pick(
312 self,
313 key=self._key,
314 ref=ref,
315 on_done=self._commit_after_change,
316 reload_worker=self._scope != "chat",
317 )
319 def _commit_after_change(self) -> None:
320 """Repaint the label, then run the chat-screen side effect for chat swaps.
322 ``apply_model_pick`` already persisted the ref and (for non-chat
323 scopes) reloaded the worker. Chat swaps cancel the in-flight stream
324 and reset services here so the new chat model takes over cleanly. Works
325 Works regardless of which container the button is mounted in.
326 """
327 self._refresh()
328 if self._scope != "chat":
329 return
330 from lilbee.cli.tui.screens.chat import ChatScreen
332 screen = self.app.screen
333 if isinstance(screen, ChatScreen):
334 screen.apply_model_change()
335 else:
336 reset_services()
339class ChatModePill(Static, can_focus=True):
340 """Single focusable mode pill; Enter / Space picks this pill's mode."""
342 BINDINGS: ClassVar[list[BindingType]] = [
343 Binding("enter", "select", "Pick mode", show=False),
344 Binding("space", "select", "Pick mode", show=False),
345 ]
347 def action_select(self) -> None:
348 toggle = next(
349 (n for n in self.ancestors_with_self if isinstance(n, ChatModeToggle)),
350 None,
351 )
352 if toggle is None:
353 return
354 target = (
355 ChatMode.SEARCH.value if self.id == _CHAT_MODE_SEARCH_PILL_ID else ChatMode.CHAT.value
356 )
357 toggle._set_mode(target)
360class ChatModeToggle(Widget, can_focus=False):
361 """Two-pill control toggling cfg.chat_mode between Search and Chat.
363 The toggle itself is not focusable; the inner pills are. Tab walks
364 Search then Chat, Enter / Space picks. The container keeps left /
365 right arrow handling so the legacy keyboard flow still works.
366 """
368 BINDINGS: ClassVar[list[BindingType]] = [
369 Binding("left", "select_search", "Search mode", show=False),
370 Binding("right", "select_chat", "Chat mode", show=False),
371 ]
373 def __init__(self) -> None:
374 super().__init__(id=_CHAT_MODE_TOGGLE_ID)
376 def compose(self) -> ComposeResult:
377 with Horizontal():
378 yield ChatModePill(
379 msg.CHAT_MODE_SEARCH_LABEL,
380 id=_CHAT_MODE_SEARCH_PILL_ID,
381 classes=_CHAT_MODE_PILL_CLASS,
382 )
383 yield ChatModePill(
384 msg.CHAT_MODE_CHAT_LABEL,
385 id=_CHAT_MODE_CHAT_PILL_ID,
386 classes=_CHAT_MODE_PILL_CLASS,
387 )
389 def on_mount(self) -> None:
390 self._refresh()
392 def refresh_state(self) -> None:
393 """Repaint label/state. Call after settings or embedding-model changes."""
394 if self.is_mounted:
395 self._refresh()
397 def _embedding_ready(self) -> bool:
398 return is_model_available(cfg.embedding_model, get_services().provider)
400 def _refresh(self) -> None:
401 ready = self._embedding_ready()
402 mode = cfg.chat_mode if ready else ChatMode.CHAT.value
403 active_search = mode == ChatMode.SEARCH.value
404 search_pill = self.query_one(f"#{_CHAT_MODE_SEARCH_PILL_ID}", ChatModePill)
405 chat_pill = self.query_one(f"#{_CHAT_MODE_CHAT_PILL_ID}", ChatModePill)
406 # Search half is disabled whenever embedding isn't ready; Chat is
407 # always reachable so it never carries the disabled class.
408 search_pill.set_class(active_search, _CHAT_MODE_ACTIVE_CLASS)
409 search_pill.set_class(not ready, _CHAT_MODE_DISABLED_CLASS)
410 chat_pill.set_class(not active_search, _CHAT_MODE_ACTIVE_CLASS)
411 chat_pill.set_class(False, _CHAT_MODE_DISABLED_CLASS)
412 # Parent carries the disabled class so external selectors can
413 # disable interaction on the whole toggle when search is gated.
414 self.set_class(not ready, _CHAT_MODE_DISABLED_CLASS)
415 self.tooltip = (
416 msg.CHAT_MODE_TOGGLE_DISABLED_TOOLTIP if not ready else msg.CHAT_MODE_TOGGLE_TOOLTIP
417 )
419 def _set_mode(self, target: str) -> bool:
420 """Apply *target* if it differs from the current mode and Search is allowed."""
421 if cfg.chat_mode == target:
422 return False
423 if target == ChatMode.SEARCH.value and not self._embedding_ready():
424 return False
425 apply_setting(self.app, "chat_mode", target)
426 self._refresh()
427 return True
429 def toggle(self) -> bool:
430 """Flip mode if embedding is ready. Returns True when the mode changed."""
431 target = (
432 ChatMode.CHAT.value if cfg.chat_mode == ChatMode.SEARCH.value else ChatMode.SEARCH.value
433 )
434 return self._set_mode(target)
436 def on_click(self, event: events.Click) -> None:
437 event.stop()
438 # Click on a specific pill picks that side; click on the container
439 # frame falls through to a toggle.
440 widget = event.widget
441 if widget is not None:
442 wid = widget.id
443 if wid == _CHAT_MODE_SEARCH_PILL_ID:
444 self._set_mode(ChatMode.SEARCH.value)
445 return
446 if wid == _CHAT_MODE_CHAT_PILL_ID:
447 self._set_mode(ChatMode.CHAT.value)
448 return
449 self.toggle()
451 def action_flip_mode(self) -> None:
452 self.toggle()
454 def action_select_search(self) -> None:
455 self._set_mode(ChatMode.SEARCH.value)
457 def action_select_chat(self) -> None:
458 self._set_mode(ChatMode.CHAT.value)
461class RoleRow(Widget, can_focus=False):
462 """One role unit in the bar: a colored role pill + its picker button."""
464 def __init__(self, *, scope: PickerScope) -> None:
465 super().__init__()
466 self.scope: PickerScope = scope
467 self._key: str = config_key_for_scope(scope)
469 def compose(self) -> ComposeResult:
470 yield Static("", classes="model-bar-pill")
471 yield ModelPickerButton(scope=self.scope, button_id=f"model-pick-{self.scope}")
473 def on_mount(self) -> None:
474 self.refresh_state()
476 @property
477 def is_active(self) -> bool:
478 return bool(getattr(cfg, self._key))
480 def _is_nullable(self) -> bool:
481 from lilbee.app.settings_map import SETTINGS_MAP
483 defn = SETTINGS_MAP.get(self._key)
484 return defn is not None and defn.nullable
486 def on_click(self, event: events.Click) -> None:
487 """Click the pill to toggle an optional role off; otherwise open the picker.
489 The picker button stops its own click events, so this handler only runs
490 for clicks on the role pill (or the row gutter).
491 """
492 event.stop()
493 if self._is_nullable() and self.is_active:
494 apply_model_pick(self, key=self._key, ref="", on_done=self.refresh_state)
495 else:
496 self.query_one(ModelPickerButton).open_picker()
498 def refresh_state(self) -> None:
499 """Repaint the role pill (colored when on, muted when off) and the picker label."""
500 active = self.is_active
501 self.set_class(active, "-active")
502 self.set_class(not active, "-off")
503 bg, fg = _SCOPE_PILL_COLORS[self.scope] if active else _OFF_PILL_COLORS
504 with contextlib.suppress(Exception):
505 self.query_one(".model-bar-pill", Static).update(
506 pill(_SCOPE_TO_LABEL[self.scope], bg, fg)
507 )
508 self.query_one(ModelPickerButton).repaint()
511class ModelBar(Widget, can_focus=False):
512 """Horizontal band of the four role pickers + the Search/Chat toggle, below the input."""
514 app: LilbeeApp # type: ignore[assignment]
515 DEFAULT_CSS: ClassVar[str] = _CSS_FILE.read_text(encoding="utf-8")
517 _SCOPES: ClassVar[tuple[PickerScope, ...]] = ("chat", "embed", "vision", "rerank")
519 def __init__(self, id: str | None = None) -> None:
520 super().__init__(id=id)
521 self._options_cache: dict[str, tuple[tuple[str, str], ...]] = {}
523 def compose(self) -> ComposeResult:
524 with Horizontal(classes="model-bar-roles"):
525 for scope in self._SCOPES:
526 yield RoleRow(scope=scope)
527 yield ChatModeToggle()
528 yield Static("", id=_CLOUD_WARNING_ID, classes="cloud-warning")
530 def on_mount(self) -> None:
531 self._refresh_cloud_warning()
532 self._scan_models()
533 self.app.settings_changed_signal.subscribe(self, self._on_settings_changed)
535 def _on_settings_changed(self, payload: tuple[str, object]) -> None:
536 key, _ = payload
537 scope = model_field_to_picker_scope().get(key)
538 if scope is None:
539 return
540 for row in self.query(RoleRow):
541 if row.scope == scope:
542 row.refresh_state()
543 if key == "chat_model":
544 self._refresh_cloud_warning()
546 @work(thread=True)
547 def _scan_models(self) -> None:
548 """Scan installed models off the UI thread and populate every role button."""
549 buckets = classify_installed_models_full()
550 scope_to_options: dict[str, list[ModelOption]] = {
551 "chat": list(buckets.get(ModelTask.CHAT, [])),
552 "embed": list(buckets.get(ModelTask.EMBEDDING, [])),
553 "vision": list(buckets.get(ModelTask.VISION, [])),
554 "rerank": list(buckets.get(ModelTask.RERANK, [])),
555 }
556 call_from_thread(self, self._populate, scope_to_options)
558 def _populate(self, scope_to_options: dict[str, list[ModelOption]]) -> None:
559 for row in self.query(RoleRow):
560 # Empty pool stays empty: the picker shows just its "Browse catalog"
561 # row rather than a pickable "(none)" pseudo-model.
562 opts = scope_to_options.get(row.scope, [])
563 fingerprint = tuple((o.label, o.ref) for o in opts)
564 if self._options_cache.get(row.scope) != fingerprint:
565 row.query_one(ModelPickerButton).set_options(opts)
566 self._options_cache[row.scope] = fingerprint
567 self._refresh_cloud_warning()
569 def _refresh_cloud_warning(self) -> None:
570 """Show a warning if the active chat model routes to a cloud provider."""
571 warning = self.query_one(f"#{_CLOUD_WARNING_ID}", Static)
572 label = _cloud_provider_label(cfg.chat_model)
573 if label is None:
574 warning.remove_class("-visible")
575 return
576 warning.update(msg.MODEL_BAR_CLOUD_PROVIDER_WARNING.format(provider=label))
577 warning.add_class("-visible")
579 def refresh_models(self) -> None:
580 """Re-scan installed models (called after downloads complete)."""
581 self._scan_models()