shotgun-sh 0.1.16.dev2__py3-none-any.whl → 0.2.1__py3-none-any.whl

shotgun/tui/screens/model_picker.py

@@ -0,0 +1,327 @@
+"""Screen for selecting AI model."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+from textual import on
+from textual.app import ComposeResult
+from textual.containers import Horizontal, Vertical
+from textual.reactive import reactive
+from textual.screen import Screen
+from textual.widgets import Button, Label, ListItem, ListView, Static
+
+from shotgun.agents.config import ConfigManager
+from shotgun.agents.config.models import MODEL_SPECS, ModelName, ShotgunConfig
+from shotgun.logging_config import get_logger
+
+if TYPE_CHECKING:
+    from ..app import ShotgunApp
+
+logger = get_logger(__name__)
+
+
+# Available models for selection
+AVAILABLE_MODELS = list(ModelName)
+
+
+def _sanitize_model_name_for_id(model_name: ModelName) -> str:
+    """Convert model name to valid Textual ID by replacing dots with hyphens."""
+    return model_name.value.replace(".", "-")
+
+
+class ModelPickerScreen(Screen[None]):
+    """Select AI model to use."""
+
+    CSS = """
+        ModelPicker {
+            layout: vertical;
+        }
+
+        #titlebox {
+            height: auto;
+            margin: 2 0;
+            padding: 1;
+            border: hkey $border;
+            content-align: center middle;
+
+            & > * {
+                text-align: center;
+            }
+        }
+
+        #model-picker-title {
+            padding: 1 0;
+            text-style: bold;
+            color: $text-accent;
+        }
+
+        #model-list {
+            margin: 2 0;
+            height: auto;
+            padding: 1;
+            & > * {
+            padding: 1 0;
+            }
+        }
+        #model-actions {
+            padding: 1;
+        }
+        #model-actions > * {
+        margin-right: 2;
+        }
+    """
+
+    BINDINGS = [
+        ("escape", "done", "Back"),
+    ]
+
+    selected_model: reactive[ModelName] = reactive(ModelName.GPT_5)
+
+    def compose(self) -> ComposeResult:
+        with Vertical(id="titlebox"):
+            yield Static("Model selection", id="model-picker-title")
+            yield Static(
+                "Select the AI model you want to use for your tasks.",
+                id="model-picker-summary",
+            )
+        yield ListView(id="model-list")
+        with Horizontal(id="model-actions"):
+            yield Button("Select \\[ENTER]", variant="primary", id="select")
+            yield Button("Done \\[ESC]", id="done")
+
+    def _rebuild_model_list(self) -> None:
+        """Rebuild the model list from current config.
+
+        This method is called both on first show and when screen is resumed
+        to ensure the list always reflects the current configuration.
+        """
+        logger.debug("Rebuilding model list from current config")
+
+        # Load current config with force_reload to get latest API keys
+        config_manager = self.config_manager
+        config = config_manager.load(force_reload=True)
+
+        # Log provider key status
+        logger.debug(
+            "Provider keys: openai=%s, anthropic=%s, google=%s, shotgun=%s",
+            config_manager._provider_has_api_key(config.openai),
+            config_manager._provider_has_api_key(config.anthropic),
+            config_manager._provider_has_api_key(config.google),
+            config_manager._provider_has_api_key(config.shotgun),
+        )
+
+        current_model = config.selected_model or ModelName.CLAUDE_SONNET_4_5
+        self.selected_model = current_model
+        logger.debug("Current selected model: %s", current_model)
+
+        # Rebuild the model list with current available models
+        list_view = self.query_one(ListView)
+
+        # Remove all existing items
+        old_count = len(list(list_view.children))
+        for child in list(list_view.children):
+            child.remove()
+        logger.debug("Removed %d existing model items from list", old_count)
+
+        # Add new items (labels already have correct text including current indicator)
+        new_items = self._build_model_items(config)
+        for item in new_items:
+            list_view.append(item)
+        logger.debug("Added %d available model items to list", len(new_items))
+
+        # Find and highlight current selection (if it's in the filtered list)
+        if list_view.children:
+            for i, child in enumerate(list_view.children):
+                if isinstance(child, ListItem) and child.id:
+                    model_id = child.id.removeprefix("model-")
+                    # Find the model name
+                    for model_name in AVAILABLE_MODELS:
+                        if _sanitize_model_name_for_id(model_name) == model_id:
+                            if model_name == current_model:
+                                list_view.index = i
+                                break
+
+    def on_show(self) -> None:
+        """Rebuild model list when screen is first shown."""
+        logger.debug("ModelPickerScreen.on_show() called")
+        self._rebuild_model_list()
+
+    def on_screenresume(self) -> None:
+        """Rebuild model list when screen is resumed (subsequent visits).
+
+        This is called when returning to the screen after it was suspended,
+        ensuring the model list reflects any config changes made while away.
+        """
+        logger.debug("ModelPickerScreen.on_screenresume() called")
+        self._rebuild_model_list()
+
+    def action_done(self) -> None:
+        self.dismiss()
+
+    @on(ListView.Highlighted)
+    def _on_model_highlighted(self, event: ListView.Highlighted) -> None:
+        model_name = self._model_from_item(event.item)
+        if model_name:
+            self.selected_model = model_name
+
+    @on(ListView.Selected)
+    def _on_model_selected(self, event: ListView.Selected) -> None:
+        model_name = self._model_from_item(event.item)
+        if model_name:
+            self.selected_model = model_name
+            self._select_model()
+
+    @on(Button.Pressed, "#select")
+    def _on_select_pressed(self) -> None:
+        self._select_model()
+
+    @on(Button.Pressed, "#done")
+    def _on_done_pressed(self) -> None:
+        self.action_done()
+
+    @property
+    def config_manager(self) -> ConfigManager:
+        app = cast("ShotgunApp", self.app)
+        return app.config_manager
+
+    def refresh_model_labels(self) -> None:
+        """Update the list view entries to reflect current selection.
+
+        Note: This method only updates labels for currently displayed models.
+        To rebuild the entire list after provider changes, on_show() should be used.
+        """
+        # Load config once with force_reload
+        config = self.config_manager.load(force_reload=True)
+        current_model = config.selected_model or ModelName.CLAUDE_SONNET_4_5
+
+        # Update labels for available models only
+        for model_name in AVAILABLE_MODELS:
+            # Pass config to avoid multiple force reloads
+            if not self._is_model_available(model_name, config):
+                continue
+            label = self.query_one(
+                f"#label-{_sanitize_model_name_for_id(model_name)}", Label
+            )
+            label.update(
+                self._model_label(model_name, is_current=model_name == current_model)
+            )
+
+    def _build_model_items(self, config: ShotgunConfig | None = None) -> list[ListItem]:
+        if config is None:
+            config = self.config_manager.load(force_reload=True)
+
+        items: list[ListItem] = []
+        current_model = self.selected_model
+        for model_name in AVAILABLE_MODELS:
+            # Only add models that are available
+            if not self._is_model_available(model_name, config):
+                continue
+
+            label = Label(
+                self._model_label(model_name, is_current=model_name == current_model),
+                id=f"label-{_sanitize_model_name_for_id(model_name)}",
+            )
+            items.append(
+                ListItem(label, id=f"model-{_sanitize_model_name_for_id(model_name)}")
+            )
+        return items
+
+    def _model_from_item(self, item: ListItem | None) -> ModelName | None:
+        """Get ModelName from a ListItem."""
+        if item is None or item.id is None:
+            return None
+        sanitized_id = item.id.removeprefix("model-")
+        # Find the original model name by comparing sanitized versions
+        for model_name in AVAILABLE_MODELS:
+            if _sanitize_model_name_for_id(model_name) == sanitized_id:
+                return model_name
+        return None
+
+    def _is_model_available(
+        self, model_name: ModelName, config: ShotgunConfig | None = None
+    ) -> bool:
+        """Check if a model is available based on provider key configuration.
+
+        A model is available if:
+        1. Shotgun Account key is configured (provides access to all models), OR
+        2. The model's provider has an API key configured (BYOK mode)
+
+        Args:
+            model_name: The model to check availability for
+            config: Optional pre-loaded config to avoid multiple reloads
+
+        Returns:
+            True if the model can be used, False otherwise
+        """
+        if config is None:
+            config = self.config_manager.load(force_reload=True)
+
+        # If Shotgun Account is configured, all models are available
+        if self.config_manager._provider_has_api_key(config.shotgun):
+            logger.debug("Model %s available (Shotgun Account configured)", model_name)
+            return True
+
+        # In BYOK mode, check if the model's provider has a key
+        if model_name not in MODEL_SPECS:
+            logger.debug("Model %s not available (not in MODEL_SPECS)", model_name)
+            return False
+
+        spec = MODEL_SPECS[model_name]
+        # Check provider key directly using the loaded config to avoid stale cache
+        provider_config = self.config_manager._get_provider_config(
+            config, spec.provider
+        )
+        has_key = self.config_manager._provider_has_api_key(provider_config)
+        logger.debug(
+            "Model %s available=%s (provider=%s, has_key=%s)",
+            model_name,
+            has_key,
+            spec.provider,
+            has_key,
+        )
+        return has_key
+
+    def _model_label(self, model_name: ModelName, is_current: bool) -> str:
+        """Generate label for model with specs and current indicator."""
+        if model_name not in MODEL_SPECS:
+            return model_name.value
+
+        spec = MODEL_SPECS[model_name]
+        display_name = self._model_display_name(model_name)
+
+        # Format context/output tokens in readable format
+        input_k = spec.max_input_tokens // 1000
+        output_k = spec.max_output_tokens // 1000
+
+        label = f"{display_name} · {input_k}K context · {output_k}K output"
+
+        # Add cost indicator for expensive models
+        if model_name == ModelName.CLAUDE_OPUS_4_1:
+            label += " · Expensive"
+
+        if is_current:
+            label += " · Current"
+
+        return label
+
+    def _model_display_name(self, model_name: ModelName) -> str:
+        """Get human-readable model name."""
+        names = {
+            ModelName.GPT_5: "GPT-5 (OpenAI)",
+            ModelName.CLAUDE_OPUS_4_1: "Claude Opus 4.1 (Anthropic)",
+            ModelName.CLAUDE_SONNET_4_5: "Claude Sonnet 4.5 (Anthropic)",
+            ModelName.GEMINI_2_5_PRO: "Gemini 2.5 Pro (Google)",
+        }
+        return names.get(model_name, model_name.value)
+
+    def _select_model(self) -> None:
+        """Save the selected model."""
+        try:
+            self.config_manager.update_selected_model(self.selected_model)
+            self.refresh_model_labels()
+            self.notify(
+                f"Selected model: {self._model_display_name(self.selected_model)}"
+            )
+        except Exception as exc:  # pragma: no cover - defensive; textual path
+            self.notify(f"Failed to select model: {exc}", severity="error")

shotgun/tui/screens/provider_config.py

@@ -12,11 +12,25 @@ from textual.screen import Screen
 from textual.widgets import Button, Input, Label, ListItem, ListView, Markdown, Static
 
 from shotgun.agents.config import ConfigManager, ProviderType
+from shotgun.utils.env_utils import is_shotgun_account_enabled
 
 if TYPE_CHECKING:
     from ..app import ShotgunApp
 
 
+def get_configurable_providers() -> list[str]:
+    """Get list of configurable providers based on feature flags.
+
+    Returns:
+        List of provider identifiers that can be configured.
+        Includes shotgun only if SHOTGUN_ACCOUNT_ENABLED is set.
+    """
+    providers = ["openai", "anthropic", "google"]
+    if is_shotgun_account_enabled():
+        providers.append("shotgun")
+    return providers
+
+
 class ProviderConfigScreen(Screen[None]):
     """Collect API keys for available providers."""
 
@@ -71,9 +85,10 @@ class ProviderConfigScreen(Screen[None]):
 
     BINDINGS = [
         ("escape", "done", "Back"),
+        ("ctrl+c", "app.quit", "Quit"),
     ]
 
-    selected_provider: reactive[ProviderType] = reactive(ProviderType.OPENAI)
+    selected_provider: reactive[str] = reactive("openai")
 
     def compose(self) -> ComposeResult:
         with Vertical(id="titlebox"):
@@ -94,17 +109,30 @@ class ProviderConfigScreen(Screen[None]):
         )
         with Horizontal(id="provider-actions"):
             yield Button("Save key \\[ENTER]", variant="primary", id="save")
+            yield Button("Authenticate", variant="success", id="authenticate")
             yield Button("Clear key", id="clear", variant="warning")
             yield Button("Done \\[ESC]", id="done")
 
     def on_mount(self) -> None:
         self.refresh_provider_status()
+        self._update_done_button_visibility()
         list_view = self.query_one(ListView)
         if list_view.children:
             list_view.index = 0
-        self.selected_provider = ProviderType.OPENAI
+        self.selected_provider = "openai"
+
+        # Hide authenticate button by default (shown only for shotgun)
+        self.query_one("#authenticate", Button).display = False
         self.set_focus(self.query_one("#api-key", Input))
 
+    def on_screenresume(self) -> None:
+        """Refresh provider status when screen is resumed.
+
+        This ensures the UI reflects any provider changes made elsewhere.
+        """
+        self.refresh_provider_status()
+        self._update_done_button_visibility()
+
     def action_done(self) -> None:
         self.dismiss()
 
@@ -125,6 +153,10 @@ class ProviderConfigScreen(Screen[None]):
     def _on_save_pressed(self) -> None:
         self._save_api_key()
 
+    @on(Button.Pressed, "#authenticate")
+    def _on_authenticate_pressed(self) -> None:
+        self.run_worker(self._start_shotgun_auth(), exclusive=True)
+
     @on(Button.Pressed, "#clear")
     def _on_clear_pressed(self) -> None:
         self._clear_api_key()
@@ -141,9 +173,31 @@ class ProviderConfigScreen(Screen[None]):
     def watch_selected_provider(self, provider: ProviderType) -> None:
         if not self.is_mounted:
             return
+
+        # Show/hide UI elements based on provider type
+        is_shotgun = provider == "shotgun"
+
         input_widget = self.query_one("#api-key", Input)
-        input_widget.placeholder = self._input_placeholder(provider)
-        input_widget.value = ""
+        save_button = self.query_one("#save", Button)
+        auth_button = self.query_one("#authenticate", Button)
+
+        if is_shotgun:
+            # Hide API key input and save button
+            input_widget.display = False
+            save_button.display = False
+
+            # Only show Authenticate button if shotgun is NOT already configured
+            if self._has_provider_key("shotgun"):
+                auth_button.display = False
+            else:
+                auth_button.display = True
+        else:
+            # Show API key input and save button, hide authenticate button
+            input_widget.display = True
+            save_button.display = True
+            auth_button.display = False
+            input_widget.placeholder = self._input_placeholder(provider)
+            input_widget.value = ""
 
     @property
     def config_manager(self) -> ConfigManager:
@@ -152,45 +206,61 @@ class ProviderConfigScreen(Screen[None]):
 
     def refresh_provider_status(self) -> None:
         """Update the list view entries to reflect configured providers."""
-        for provider in ProviderType:
-            label = self.query_one(f"#label-{provider.value}", Label)
-            label.update(self._provider_label(provider))
+        for provider_id in get_configurable_providers():
+            label = self.query_one(f"#label-{provider_id}", Label)
+            label.update(self._provider_label(provider_id))
+
+    def _update_done_button_visibility(self) -> None:
+        """Show/hide Done button based on whether any provider keys are configured."""
+        done_button = self.query_one("#done", Button)
+        has_keys = self.config_manager.has_any_provider_key()
+        done_button.display = has_keys
 
     def _build_provider_items(self) -> list[ListItem]:
         items: list[ListItem] = []
-        for provider in ProviderType:
-            label = Label(self._provider_label(provider), id=f"label-{provider.value}")
-            items.append(ListItem(label, id=f"provider-{provider.value}"))
+        for provider_id in get_configurable_providers():
+            label = Label(self._provider_label(provider_id), id=f"label-{provider_id}")
+            items.append(ListItem(label, id=f"provider-{provider_id}"))
         return items
 
-    def _provider_from_item(self, item: ListItem | None) -> ProviderType | None:
+    def _provider_from_item(self, item: ListItem | None) -> str | None:
         if item is None or item.id is None:
             return None
         provider_id = item.id.removeprefix("provider-")
-        try:
-            return ProviderType(provider_id)
-        except ValueError:
-            return None
+        return provider_id if provider_id in get_configurable_providers() else None
 
-    def _provider_label(self, provider: ProviderType) -> str:
-        display = self._provider_display_name(provider)
+    def _provider_label(self, provider_id: str) -> str:
+        display = self._provider_display_name(provider_id)
         status = (
-            "Configured"
-            if self.config_manager.has_provider_key(provider)
-            else "Not configured"
+            "Configured" if self._has_provider_key(provider_id) else "Not configured"
         )
         return f"{display} · {status}"
 
-    def _provider_display_name(self, provider: ProviderType) -> str:
+    def _provider_display_name(self, provider_id: str) -> str:
         names = {
-            ProviderType.OPENAI: "OpenAI",
-            ProviderType.ANTHROPIC: "Anthropic",
-            ProviderType.GOOGLE: "Google Gemini",
+            "openai": "OpenAI",
+            "anthropic": "Anthropic",
+            "google": "Google Gemini",
+            "shotgun": "Shotgun Account",
         }
-        return names.get(provider, provider.value.title())
-
-    def _input_placeholder(self, provider: ProviderType) -> str:
-        return f"{self._provider_display_name(provider)} API key"
+        return names.get(provider_id, provider_id.title())
+
+    def _input_placeholder(self, provider_id: str) -> str:
+        return f"{self._provider_display_name(provider_id)} API key"
+
+    def _has_provider_key(self, provider_id: str) -> bool:
+        """Check if provider has a configured API key."""
+        if provider_id == "shotgun":
+            # Check shotgun key directly
+            config = self.config_manager.load()
+            return self.config_manager._provider_has_api_key(config.shotgun)
+        else:
+            # Check LLM provider key
+            try:
+                provider = ProviderType(provider_id)
+                return self.config_manager.has_provider_key(provider)
+            except ValueError:
+                return False
 
     def _save_api_key(self) -> None:
         input_widget = self.query_one("#api-key", Input)
@@ -211,6 +281,7 @@ class ProviderConfigScreen(Screen[None]):
 
         input_widget.value = ""
         self.refresh_provider_status()
+        self._update_done_button_visibility()
         self.notify(
             f"Saved API key for {self._provider_display_name(self.selected_provider)}."
         )
@@ -223,7 +294,26 @@ class ProviderConfigScreen(Screen[None]):
             return
 
         self.refresh_provider_status()
+        self._update_done_button_visibility()
         self.query_one("#api-key", Input).value = ""
+
+        # If we just cleared shotgun, show the Authenticate button
+        if self.selected_provider == "shotgun":
+            auth_button = self.query_one("#authenticate", Button)
+            auth_button.display = True
+
         self.notify(
             f"Cleared API key for {self._provider_display_name(self.selected_provider)}."
         )
+
+    async def _start_shotgun_auth(self) -> None:
+        """Launch Shotgun Account authentication flow."""
+        from .shotgun_auth import ShotgunAuthScreen
+
+        # Push the auth screen and wait for result
+        result = await self.app.push_screen_wait(ShotgunAuthScreen())
+
+        # Refresh provider status after auth completes
+        if result:
+            self.refresh_provider_status()
+            # Notify handled by auth screen