shotgun-sh 0.1.15.dev1__py3-none-any.whl → 0.2.1.dev1__py3-none-any.whl

shotgun/agents/common.py

@@ -18,7 +18,7 @@ from pydantic_ai.messages import (
     ModelRequest,
 )
 
-from shotgun.agents.config import ProviderType, get_config_manager, get_provider_model
+from shotgun.agents.config import ProviderType, get_provider_model
 from shotgun.agents.models import AgentType
 from shotgun.logging_config import get_logger
 from shotgun.prompts import PromptLoader
@@ -115,14 +115,13 @@ def create_base_agent(
     """
     ensure_shotgun_directory_exists()
 
-    # Get configured model or fall back to hardcoded default
+    # Get configured model or fall back to first available provider
     try:
         model_config = get_provider_model(provider)
-        config_manager = get_config_manager()
-        provider_name = provider or config_manager.load().default_provider
+        provider_name = model_config.provider
         logger.debug(
             "🤖 Creating agent with configured %s model: %s",
-            provider_name.upper(),
+            provider_name.value.upper(),
             model_config.name,
         )
         # Use the Model instance directly (has API key baked in)

shotgun/agents/config/constants.py

@@ -1,17 +1,33 @@
 """Configuration constants for Shotgun agents."""
 
+from enum import StrEnum, auto
+
 # Field names
 API_KEY_FIELD = "api_key"
-DEFAULT_PROVIDER_FIELD = "default_provider"
 USER_ID_FIELD = "user_id"
 CONFIG_VERSION_FIELD = "config_version"
 
-# Provider names (for consistency with data dict keys)
-OPENAI_PROVIDER = "openai"
-ANTHROPIC_PROVIDER = "anthropic"
-GOOGLE_PROVIDER = "google"
+
+class ConfigSection(StrEnum):
+    """Configuration file section names (JSON keys)."""
+
+    OPENAI = auto()
+    ANTHROPIC = auto()
+    GOOGLE = auto()
+    SHOTGUN = auto()
+
+
+# Backwards compatibility - deprecated
+OPENAI_PROVIDER = ConfigSection.OPENAI.value
+ANTHROPIC_PROVIDER = ConfigSection.ANTHROPIC.value
+GOOGLE_PROVIDER = ConfigSection.GOOGLE.value
+SHOTGUN_PROVIDER = ConfigSection.SHOTGUN.value
 
 # Environment variable names
 OPENAI_API_KEY_ENV = "OPENAI_API_KEY"
 ANTHROPIC_API_KEY_ENV = "ANTHROPIC_API_KEY"
 GEMINI_API_KEY_ENV = "GEMINI_API_KEY"
+SHOTGUN_API_KEY_ENV = "SHOTGUN_API_KEY"
+
+# Token limits
+MEDIUM_TEXT_8K_TOKENS = 8192  # Default max_tokens for web search requests

shotgun/agents/config/manager.py

@@ -13,17 +13,26 @@ from shotgun.utils import get_shotgun_home
 
 from .constants import (
     ANTHROPIC_API_KEY_ENV,
-    ANTHROPIC_PROVIDER,
     API_KEY_FIELD,
     GEMINI_API_KEY_ENV,
-    GOOGLE_PROVIDER,
     OPENAI_API_KEY_ENV,
-    OPENAI_PROVIDER,
+    ConfigSection,
+)
+from .models import (
+    AnthropicConfig,
+    GoogleConfig,
+    ModelName,
+    OpenAIConfig,
+    ProviderType,
+    ShotgunAccountConfig,
+    ShotgunConfig,
 )
-from .models import ProviderType, ShotgunConfig
 
 logger = get_logger(__name__)
 
+# Type alias for provider configuration objects
+ProviderConfig = OpenAIConfig | AnthropicConfig | GoogleConfig | ShotgunAccountConfig
+
 
 class ConfigManager:
     """Manager for Shotgun configuration."""
@@ -69,20 +78,56 @@ class ConfigManager:
             self._config = ShotgunConfig.model_validate(data)
             logger.debug("Configuration loaded successfully from %s", self.config_path)
 
-            # Check if the default provider has a key, if not find one that does
-            if not self.has_provider_key(self._config.default_provider):
-                original_default = self._config.default_provider
-                # Find first provider with a configured key
-                for provider in ProviderType:
-                    if self.has_provider_key(provider):
+            # Validate selected_model if in BYOK mode (no Shotgun key)
+            if not self._provider_has_api_key(self._config.shotgun):
+                should_save = False
+
+                # If selected_model is set, verify its provider has a key
+                if self._config.selected_model:
+                    from .models import MODEL_SPECS
+
+                    if self._config.selected_model in MODEL_SPECS:
+                        spec = MODEL_SPECS[self._config.selected_model]
+                        if not self.has_provider_key(spec.provider):
+                            logger.info(
+                                "Selected model %s provider has no API key, finding available model",
+                                self._config.selected_model.value,
+                            )
+                            self._config.selected_model = None
+                            should_save = True
+                    else:
                         logger.info(
-                            "Default provider %s has no API key, updating to %s",
-                            original_default.value,
-                            provider.value,
+                            "Selected model %s not found in MODEL_SPECS, resetting",
+                            self._config.selected_model.value,
                         )
-                        self._config.default_provider = provider
-                        self.save(self._config)
-                        break
+                        self._config.selected_model = None
+                        should_save = True
+
+                # If no selected_model or it was invalid, find first available model
+                if not self._config.selected_model:
+                    for provider in ProviderType:
+                        if self.has_provider_key(provider):
+                            # Set to that provider's default model
+                            from .models import MODEL_SPECS, ModelName
+
+                            # Find default model for this provider
+                            provider_models = {
+                                ProviderType.OPENAI: ModelName.GPT_5,
+                                ProviderType.ANTHROPIC: ModelName.CLAUDE_OPUS_4_1,
+                                ProviderType.GOOGLE: ModelName.GEMINI_2_5_PRO,
+                            }
+
+                            if provider in provider_models:
+                                self._config.selected_model = provider_models[provider]
+                                logger.info(
+                                    "Set selected_model to %s (first available provider)",
+                                    self._config.selected_model.value,
+                                )
+                                should_save = True
+                                break
+
+                if should_save:
+                    self.save(self._config)
 
             return self._config
 
@@ -107,7 +152,6 @@ class ConfigManager:
                 # Create a new config with generated user_id
                 config = ShotgunConfig(
                     user_id=str(uuid.uuid4()),
-                    config_version=1,
                 )
 
         # Ensure directory exists
@@ -136,8 +180,13 @@ class ConfigManager:
             **kwargs: Configuration fields to update (only api_key supported)
         """
         config = self.load()
-        provider_enum = self._ensure_provider_enum(provider)
-        provider_config = self._get_provider_config(config, provider_enum)
+
+        # Get provider config and check if it's shotgun
+        provider_config, is_shotgun = self._get_provider_config_and_type(
+            config, provider
+        )
+        # For non-shotgun providers, we need the enum for default provider logic
+        provider_enum = None if is_shotgun else self._ensure_provider_enum(provider)
 
         # Only support api_key updates
         if API_KEY_FIELD in kwargs:
@@ -152,23 +201,47 @@ class ConfigManager:
             raise ValueError(f"Unsupported configuration fields: {unsupported_fields}")
 
         # If no other providers have keys configured and we just added one,
-        # set this provider as the default
-        if API_KEY_FIELD in kwargs and api_key_value is not None:
+        # set selected_model to that provider's default model (only for LLM providers, not shotgun)
+        if not is_shotgun and API_KEY_FIELD in kwargs and api_key_value is not None:
+            # provider_enum is guaranteed to be non-None here since is_shotgun is False
+            if provider_enum is None:
+                raise RuntimeError("Provider enum should not be None for LLM providers")
             other_providers = [p for p in ProviderType if p != provider_enum]
             has_other_keys = any(self.has_provider_key(p) for p in other_providers)
             if not has_other_keys:
-                config.default_provider = provider_enum
+                # Set selected_model to this provider's default model
+                from .models import ModelName
+
+                provider_models = {
+                    ProviderType.OPENAI: ModelName.GPT_5,
+                    ProviderType.ANTHROPIC: ModelName.CLAUDE_OPUS_4_1,
+                    ProviderType.GOOGLE: ModelName.GEMINI_2_5_PRO,
+                }
+                if provider_enum in provider_models:
+                    config.selected_model = provider_models[provider_enum]
 
         self.save(config)
 
     def clear_provider_key(self, provider: ProviderType | str) -> None:
-        """Remove the API key for the given provider."""
+        """Remove the API key for the given provider (LLM provider or shotgun)."""
         config = self.load()
-        provider_enum = self._ensure_provider_enum(provider)
-        provider_config = self._get_provider_config(config, provider_enum)
+
+        # Get provider config (shotgun or LLM provider)
+        provider_config, _ = self._get_provider_config_and_type(config, provider)
+
         provider_config.api_key = None
         self.save(config)
 
+    def update_selected_model(self, model_name: "ModelName") -> None:
+        """Update the selected model.
+
+        Args:
+            model_name: Model to select
+        """
+        config = self.load()
+        config.selected_model = model_name
+        self.save(config)
+
     def has_provider_key(self, provider: ProviderType | str) -> bool:
         """Check if the given provider has a non-empty API key configured.
 
@@ -195,7 +268,8 @@ class ConfigManager:
     def has_any_provider_key(self) -> bool:
         """Determine whether any provider has a configured API key."""
         config = self.load()
-        return any(
+        # Check LLM provider keys (BYOK)
+        has_llm_key = any(
             self._provider_has_api_key(self._get_provider_config(config, provider))
             for provider in (
                 ProviderType.OPENAI,
@@ -203,6 +277,9 @@ class ConfigManager:
                 ProviderType.GOOGLE,
             )
         )
+        # Also check Shotgun Account key
+        has_shotgun_key = self._provider_has_api_key(config.shotgun)
+        return has_llm_key or has_shotgun_key
 
     def initialize(self) -> ShotgunConfig:
         """Initialize configuration with defaults and save to file.
@@ -213,7 +290,6 @@ class ConfigManager:
         # Generate unique user ID for new config
         config = ShotgunConfig(
             user_id=str(uuid.uuid4()),
-            config_version=1,
         )
         self.save(config)
         logger.info(
@@ -225,26 +301,26 @@ class ConfigManager:
 
     def _convert_secrets_to_secretstr(self, data: dict[str, Any]) -> None:
         """Convert plain text secrets in data to SecretStr objects."""
-        for provider in [OPENAI_PROVIDER, ANTHROPIC_PROVIDER, GOOGLE_PROVIDER]:
-            if provider in data and isinstance(data[provider], dict):
+        for section in ConfigSection:
+            if section.value in data and isinstance(data[section.value], dict):
                 if (
-                    API_KEY_FIELD in data[provider]
-                    and data[provider][API_KEY_FIELD] is not None
+                    API_KEY_FIELD in data[section.value]
+                    and data[section.value][API_KEY_FIELD] is not None
                 ):
-                    data[provider][API_KEY_FIELD] = SecretStr(
-                        data[provider][API_KEY_FIELD]
+                    data[section.value][API_KEY_FIELD] = SecretStr(
+                        data[section.value][API_KEY_FIELD]
                     )
 
     def _convert_secretstr_to_plain(self, data: dict[str, Any]) -> None:
         """Convert SecretStr objects in data to plain text for JSON serialization."""
-        for provider in [OPENAI_PROVIDER, ANTHROPIC_PROVIDER, GOOGLE_PROVIDER]:
-            if provider in data and isinstance(data[provider], dict):
+        for section in ConfigSection:
+            if section.value in data and isinstance(data[section.value], dict):
                 if (
-                    API_KEY_FIELD in data[provider]
-                    and data[provider][API_KEY_FIELD] is not None
+                    API_KEY_FIELD in data[section.value]
+                    and data[section.value][API_KEY_FIELD] is not None
                 ):
-                    if hasattr(data[provider][API_KEY_FIELD], "get_secret_value"):
-                        data[provider][API_KEY_FIELD] = data[provider][
+                    if hasattr(data[section.value][API_KEY_FIELD], "get_secret_value"):
+                        data[section.value][API_KEY_FIELD] = data[section.value][
                             API_KEY_FIELD
                         ].get_secret_value()
 
@@ -279,6 +355,38 @@ class ConfigManager:
 
         return bool(value.strip())
 
+    def _is_shotgun_provider(self, provider: ProviderType | str) -> bool:
+        """Check if provider string represents Shotgun Account.
+
+        Args:
+            provider: Provider type or string
+
+        Returns:
+            True if provider is shotgun account
+        """
+        return (
+            isinstance(provider, str)
+            and provider.lower() == ConfigSection.SHOTGUN.value
+        )
+
+    def _get_provider_config_and_type(
+        self, config: ShotgunConfig, provider: ProviderType | str
+    ) -> tuple[ProviderConfig, bool]:
+        """Get provider config, handling shotgun as special case.
+
+        Args:
+            config: Shotgun configuration
+            provider: Provider type or string
+
+        Returns:
+            Tuple of (provider_config, is_shotgun)
+        """
+        if self._is_shotgun_provider(provider):
+            return (config.shotgun, True)
+
+        provider_enum = self._ensure_provider_enum(provider)
+        return (self._get_provider_config(config, provider_enum), False)
+
     def get_user_id(self) -> str:
         """Get the user ID from configuration.
 

shotgun/agents/config/models.py

@@ -1,16 +1,12 @@
 """Pydantic models for configuration."""
 
-from enum import Enum
-from typing import Any
+from enum import StrEnum
 
 from pydantic import BaseModel, Field, PrivateAttr, SecretStr
-from pydantic_ai.direct import model_request
-from pydantic_ai.messages import ModelMessage, ModelResponse
 from pydantic_ai.models import Model
-from pydantic_ai.settings import ModelSettings
 
 
-class ProviderType(str, Enum):
+class ProviderType(StrEnum):
     """Provider types for AI services."""
 
     OPENAI = "openai"
@@ -18,20 +14,41 @@ class ProviderType(str, Enum):
     GOOGLE = "google"
 
 
+class KeyProvider(StrEnum):
+    """Authentication method for accessing AI models."""
+
+    BYOK = "byok"  # Bring Your Own Key (individual provider keys)
+    SHOTGUN = "shotgun"  # Shotgun Account (unified LiteLLM proxy)
+
+
+class ModelName(StrEnum):
+    """Available AI model names."""
+
+    GPT_5 = "gpt-5"
+    GPT_5_MINI = "gpt-5-mini"
+    CLAUDE_OPUS_4_1 = "claude-opus-4-1"
+    GEMINI_2_5_PRO = "gemini-2.5-pro"
+    GEMINI_2_5_FLASH = "gemini-2.5-flash"
+
+
 class ModelSpec(BaseModel):
     """Static specification for a model - just metadata."""
 
-    name: str  # Model identifier (e.g., "gpt-5", "claude-opus-4-1")
+    name: ModelName  # Model identifier
     provider: ProviderType
     max_input_tokens: int
     max_output_tokens: int
+    litellm_proxy_model_name: (
+        str  # LiteLLM format (e.g., "openai/gpt-5", "gemini/gemini-2-pro")
+    )
 
 
 class ModelConfig(BaseModel):
     """A fully configured model with API key and settings."""
 
-    name: str  # Model identifier (e.g., "gpt-5", "claude-opus-4-1")
-    provider: ProviderType
+    name: ModelName  # Model identifier
+    provider: ProviderType  # Actual LLM provider (openai, anthropic, google)
+    key_provider: KeyProvider  # Authentication method (byok or shotgun)
     max_input_tokens: int
     max_output_tokens: int
     api_key: str
@@ -47,7 +64,7 @@ class ModelConfig(BaseModel):
             from .provider import get_or_create_model
 
             self._model_instance = get_or_create_model(
-                self.provider, self.name, self.api_key
+                self.provider, self.key_provider, self.name, self.api_key
             )
         return self._model_instance
 
@@ -61,54 +78,43 @@ class ModelConfig(BaseModel):
         }
         return f"{provider_prefix[self.provider]}:{self.name}"
 
-    def get_model_settings(self, max_tokens: int | None = None) -> ModelSettings:
-        """Get ModelSettings with optional token override.
-
-        This provides flexibility for specific use cases that need different
-        token limits while defaulting to maximum utilization.
-
-        Args:
-            max_tokens: Optional override for max_tokens. If None, uses max_output_tokens
-
-        Returns:
-            ModelSettings configured with specified or maximum tokens
-        """
-        return ModelSettings(
-            max_tokens=max_tokens if max_tokens is not None else self.max_output_tokens
-        )
-
 
 # Model specifications registry (static metadata)
-MODEL_SPECS: dict[str, ModelSpec] = {
-    "gpt-5": ModelSpec(
-        name="gpt-5",
+MODEL_SPECS: dict[ModelName, ModelSpec] = {
+    ModelName.GPT_5: ModelSpec(
+        name=ModelName.GPT_5,
         provider=ProviderType.OPENAI,
         max_input_tokens=400_000,
         max_output_tokens=128_000,
+        litellm_proxy_model_name="openai/gpt-5",
     ),
-    "gpt-4o": ModelSpec(
-        name="gpt-4o",
+    ModelName.GPT_5_MINI: ModelSpec(
+        name=ModelName.GPT_5_MINI,
         provider=ProviderType.OPENAI,
-        max_input_tokens=128_000,
-        max_output_tokens=16_000,
+        max_input_tokens=400_000,
+        max_output_tokens=128_000,
+        litellm_proxy_model_name="openai/gpt-5-mini",
     ),
-    "claude-opus-4-1": ModelSpec(
-        name="claude-opus-4-1",
+    ModelName.CLAUDE_OPUS_4_1: ModelSpec(
+        name=ModelName.CLAUDE_OPUS_4_1,
         provider=ProviderType.ANTHROPIC,
         max_input_tokens=200_000,
         max_output_tokens=32_000,
+        litellm_proxy_model_name="anthropic/claude-opus-4-1",
     ),
-    "claude-3-5-sonnet-latest": ModelSpec(
-        name="claude-3-5-sonnet-latest",
-        provider=ProviderType.ANTHROPIC,
-        max_input_tokens=200_000,
-        max_output_tokens=8_192,
+    ModelName.GEMINI_2_5_PRO: ModelSpec(
+        name=ModelName.GEMINI_2_5_PRO,
+        provider=ProviderType.GOOGLE,
+        max_input_tokens=1_000_000,
+        max_output_tokens=64_000,
+        litellm_proxy_model_name="gemini/gemini-2.5-pro",
     ),
-    "gemini-2.5-pro": ModelSpec(
-        name="gemini-2.5-pro",
+    ModelName.GEMINI_2_5_FLASH: ModelSpec(
+        name=ModelName.GEMINI_2_5_FLASH,
         provider=ProviderType.GOOGLE,
         max_input_tokens=1_000_000,
         max_output_tokens=64_000,
+        litellm_proxy_model_name="gemini/gemini-2.5-flash",
     ),
 }
 
@@ -131,55 +137,22 @@ class GoogleConfig(BaseModel):
     api_key: SecretStr | None = None
 
 
+class ShotgunAccountConfig(BaseModel):
+    """Configuration for Shotgun Account (LiteLLM proxy)."""
+
+    api_key: SecretStr | None = None
+
+
 class ShotgunConfig(BaseModel):
     """Main configuration for Shotgun CLI."""
 
     openai: OpenAIConfig = Field(default_factory=OpenAIConfig)
     anthropic: AnthropicConfig = Field(default_factory=AnthropicConfig)
     google: GoogleConfig = Field(default_factory=GoogleConfig)
-    default_provider: ProviderType = Field(
-        default=ProviderType.OPENAI, description="Default AI provider to use"
+    shotgun: ShotgunAccountConfig = Field(default_factory=ShotgunAccountConfig)
+    selected_model: ModelName | None = Field(
+        default=None,
+        description="User-selected model",
     )
     user_id: str = Field(description="Unique anonymous user identifier")
-    config_version: int = Field(default=1, description="Configuration schema version")
-
-
-async def shotgun_model_request(
-    model_config: ModelConfig,
-    messages: list[ModelMessage],
-    max_tokens: int | None = None,
-    **kwargs: Any,
-) -> ModelResponse:
-    """Model request wrapper that uses full token capacity by default.
-
-    This wrapper ensures all LLM calls in Shotgun use the maximum available
-    token capacity of each model, improving response quality and completeness.
-    The most common issue this fixes is truncated summaries that were cut off
-    at default token limits (e.g., 4096 for Claude models).
-
-    Args:
-        model_config: ModelConfig instance with model settings and API key
-        messages: Messages to send to the model
-        max_tokens: Optional override for max_tokens. If None, uses model's max_output_tokens
-        **kwargs: Additional arguments passed to model_request
-
-    Returns:
-        ModelResponse from the model
-
-    Example:
-        # Uses full token capacity (e.g., 4096 for Claude, 128k for GPT-5)
-        response = await shotgun_model_request(model_config, messages)
-
-        # Override for specific use case
-        response = await shotgun_model_request(model_config, messages, max_tokens=1000)
-    """
-    # Get properly configured ModelSettings with maximum or overridden token limit
-    model_settings = model_config.get_model_settings(max_tokens)
-
-    # Make the model request with full token utilization
-    return await model_request(
-        model=model_config.model_instance,
-        messages=messages,
-        model_settings=model_settings,
-        **kwargs,
-    )
+    config_version: int = Field(default=2, description="Configuration schema version")