shotgun-sh 0.1.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

shotgun/agents/config/models.py

@@ -1,16 +1,13 @@
 """Pydantic models for configuration."""
 
-from enum import Enum
-from typing import Any
+from datetime import datetime
+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 +15,44 @@ 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"
+    CLAUDE_SONNET_4_5 = "claude-sonnet-4-5"
+    CLAUDE_HAIKU_4_5 = "claude-haiku-4-5"
+    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")
+    )
+    short_name: str  # Display name for UI (e.g., "Sonnet 4.5", "GPT-5")
 
 
 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 +68,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 +82,64 @@ 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",
+        short_name="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",
+        short_name="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",
+        short_name="Opus 4.1",
     ),
-    "claude-3-5-sonnet-latest": ModelSpec(
-        name="claude-3-5-sonnet-latest",
+    ModelName.CLAUDE_SONNET_4_5: ModelSpec(
+        name=ModelName.CLAUDE_SONNET_4_5,
         provider=ProviderType.ANTHROPIC,
         max_input_tokens=200_000,
-        max_output_tokens=8_192,
+        max_output_tokens=16_000,
+        litellm_proxy_model_name="anthropic/claude-sonnet-4-5",
+        short_name="Sonnet 4.5",
+    ),
+    ModelName.CLAUDE_HAIKU_4_5: ModelSpec(
+        name=ModelName.CLAUDE_HAIKU_4_5,
+        provider=ProviderType.ANTHROPIC,
+        max_input_tokens=200_000,
+        max_output_tokens=64_000,
+        litellm_proxy_model_name="anthropic/claude-haiku-4-5",
+        short_name="Haiku 4.5",
     ),
-    "gemini-2.5-pro": ModelSpec(
-        name="gemini-2.5-pro",
+    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",
+        short_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",
+        short_name="Gemini 2.5 Flash",
     ),
 }
 
@@ -131,55 +162,54 @@ class GoogleConfig(BaseModel):
     api_key: SecretStr | None = None
 
 
+class ShotgunAccountConfig(BaseModel):
+    """Configuration for Shotgun Account (LiteLLM proxy)."""
+
+    api_key: SecretStr | None = None
+    supabase_jwt: SecretStr | None = Field(
+        default=None, description="Supabase authentication JWT"
+    )
+
+
+class MarketingMessageRecord(BaseModel):
+    """Record of when a marketing message was shown to the user."""
+
+    shown_at: datetime = Field(description="Timestamp when the message was shown")
+
+
+class MarketingConfig(BaseModel):
+    """Configuration for marketing messages shown to users."""
+
+    messages: dict[str, MarketingMessageRecord] = Field(
+        default_factory=dict,
+        description="Tracking which marketing messages have been shown. Key is message ID (e.g., 'github_star_v1')",
+    )
+
+
 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",
+    )
+    shotgun_instance_id: str = Field(
+        description="Unique shotgun instance identifier (also used for anonymous telemetry)",
+    )
+    config_version: int = Field(default=4, description="Configuration schema version")
+    shown_welcome_screen: bool = Field(
+        default=False,
+        description="Whether the welcome screen has been shown to the user",
+    )
+    shown_onboarding_popup: datetime | None = Field(
+        default=None,
+        description="Timestamp when the onboarding popup was shown to the user (ISO8601 format)",
     )
-    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,
+    marketing: MarketingConfig = Field(
+        default_factory=MarketingConfig,
+        description="Marketing messages configuration and tracking",
     )