shotgun-sh 0.2.8.dev2__py3-none-any.whl → 0.3.3.dev1__py3-none-any.whl

shotgun/agents/config/models.py

@@ -1,5 +1,6 @@
 """Pydantic models for configuration."""
 
+from datetime import datetime
 from enum import StrEnum
 
 from pydantic import BaseModel, Field, PrivateAttr, SecretStr
@@ -24,12 +25,17 @@ class KeyProvider(StrEnum):
 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"
+    GPT_5_1 = "gpt-5.1"
+    GPT_5_1_CODEX = "gpt-5.1-codex"
+    GPT_5_1_CODEX_MINI = "gpt-5.1-codex-mini"
+    CLAUDE_OPUS_4_5 = "claude-opus-4-5"
+    CLAUDE_SONNET_4 = "claude-sonnet-4"
     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"
+    GEMINI_2_5_FLASH_LITE = "gemini-2.5-flash-lite"
+    GEMINI_3_PRO_PREVIEW = "gemini-3-pro-preview"
 
 
 class ModelSpec(BaseModel):
@@ -42,6 +48,7 @@ class ModelSpec(BaseModel):
     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):
@@ -53,6 +60,10 @@ class ModelConfig(BaseModel):
     max_input_tokens: int
     max_output_tokens: int
     api_key: str
+    supports_streaming: bool = Field(
+        default=True,
+        description="Whether this model configuration supports streaming. False only for BYOK GPT-5 models without streaming enabled.",
+    )
     _model_instance: Model | None = PrivateAttr(default=None)
 
     class Config:
@@ -79,29 +90,41 @@ class ModelConfig(BaseModel):
         }
         return f"{provider_prefix[self.provider]}:{self.name}"
 
+    @property
+    def is_shotgun_account(self) -> bool:
+        """Check if this model is using Shotgun Account authentication.
+
+        Returns:
+            True if using Shotgun Account, False if BYOK
+        """
+        return self.key_provider == KeyProvider.SHOTGUN
+
 
 # Model specifications registry (static metadata)
 MODEL_SPECS: dict[ModelName, ModelSpec] = {
-    ModelName.GPT_5: ModelSpec(
-        name=ModelName.GPT_5,
+    ModelName.GPT_5_1: ModelSpec(
+        name=ModelName.GPT_5_1,
         provider=ProviderType.OPENAI,
-        max_input_tokens=400_000,
+        max_input_tokens=272_000,
         max_output_tokens=128_000,
-        litellm_proxy_model_name="openai/gpt-5",
+        litellm_proxy_model_name="openai/gpt-5.1",
+        short_name="GPT-5.1",
     ),
-    ModelName.GPT_5_MINI: ModelSpec(
-        name=ModelName.GPT_5_MINI,
+    ModelName.GPT_5_1_CODEX: ModelSpec(
+        name=ModelName.GPT_5_1_CODEX,
         provider=ProviderType.OPENAI,
-        max_input_tokens=400_000,
+        max_input_tokens=272_000,
         max_output_tokens=128_000,
-        litellm_proxy_model_name="openai/gpt-5-mini",
+        litellm_proxy_model_name="openai/gpt-5.1-codex",
+        short_name="GPT-5.1 Codex",
     ),
-    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",
+    ModelName.GPT_5_1_CODEX_MINI: ModelSpec(
+        name=ModelName.GPT_5_1_CODEX_MINI,
+        provider=ProviderType.OPENAI,
+        max_input_tokens=272_000,
+        max_output_tokens=128_000,
+        litellm_proxy_model_name="openai/gpt-5.1-codex-mini",
+        short_name="GPT-5.1 Codex Mini",
     ),
     ModelName.CLAUDE_SONNET_4_5: ModelSpec(
         name=ModelName.CLAUDE_SONNET_4_5,
@@ -109,6 +132,15 @@ MODEL_SPECS: dict[ModelName, ModelSpec] = {
         max_input_tokens=200_000,
         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",
     ),
     ModelName.GEMINI_2_5_PRO: ModelSpec(
         name=ModelName.GEMINI_2_5_PRO,
@@ -116,6 +148,7 @@ MODEL_SPECS: dict[ModelName, ModelSpec] = {
         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,
@@ -123,6 +156,39 @@ MODEL_SPECS: dict[ModelName, ModelSpec] = {
         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",
+    ),
+    ModelName.CLAUDE_OPUS_4_5: ModelSpec(
+        name=ModelName.CLAUDE_OPUS_4_5,
+        provider=ProviderType.ANTHROPIC,
+        max_input_tokens=200_000,
+        max_output_tokens=64_000,
+        litellm_proxy_model_name="anthropic/claude-opus-4-5",
+        short_name="Opus 4.5",
+    ),
+    ModelName.CLAUDE_SONNET_4: ModelSpec(
+        name=ModelName.CLAUDE_SONNET_4,
+        provider=ProviderType.ANTHROPIC,
+        max_input_tokens=200_000,
+        max_output_tokens=64_000,
+        litellm_proxy_model_name="anthropic/claude-sonnet-4",
+        short_name="Sonnet 4",
+    ),
+    ModelName.GEMINI_2_5_FLASH_LITE: ModelSpec(
+        name=ModelName.GEMINI_2_5_FLASH_LITE,
+        provider=ProviderType.GOOGLE,
+        max_input_tokens=1_048_576,
+        max_output_tokens=65_536,
+        litellm_proxy_model_name="gemini/gemini-2.5-flash-lite",
+        short_name="Gemini 2.5 Flash Lite",
+    ),
+    ModelName.GEMINI_3_PRO_PREVIEW: ModelSpec(
+        name=ModelName.GEMINI_3_PRO_PREVIEW,
+        provider=ProviderType.GOOGLE,
+        max_input_tokens=1_048_576,
+        max_output_tokens=65_536,
+        litellm_proxy_model_name="gemini/gemini-3-pro-preview",
+        short_name="Gemini 3 Pro",
     ),
 }
 
@@ -131,6 +197,10 @@ class OpenAIConfig(BaseModel):
     """Configuration for OpenAI provider."""
 
     api_key: SecretStr | None = None
+    supports_streaming: bool | None = Field(
+        default=None,
+        description="Whether streaming is supported for this API key. None = not tested yet",
+    )
 
 
 class AnthropicConfig(BaseModel):
@@ -152,6 +222,24 @@ class ShotgunAccountConfig(BaseModel):
     supabase_jwt: SecretStr | None = Field(
         default=None, description="Supabase authentication JWT"
     )
+    workspace_id: str | None = Field(
+        default=None, description="Default workspace ID for shared specs"
+    )
+
+
+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):
@@ -168,8 +256,24 @@ class ShotgunConfig(BaseModel):
     shotgun_instance_id: str = Field(
         description="Unique shotgun instance identifier (also used for anonymous telemetry)",
     )
-    config_version: int = Field(default=3, description="Configuration schema version")
+    config_version: int = Field(default=5, 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)",
+    )
+    marketing: MarketingConfig = Field(
+        default_factory=MarketingConfig,
+        description="Marketing messages configuration and tracking",
+    )
+    migration_failed: bool = Field(
+        default=False,
+        description="Whether the last config migration failed (cleared after user configures a provider)",
+    )
+    migration_backup_path: str | None = Field(
+        default=None,
+        description="Path to the backup file created when migration failed",
+    )

shotgun/agents/config/provider.py

@@ -25,6 +25,7 @@ from .models import (
     ProviderType,
     ShotgunConfig,
 )
+from .streaming_test import check_streaming_capability
 
 logger = get_logger(__name__)
 
@@ -32,6 +33,34 @@ logger = get_logger(__name__)
 _model_cache: dict[tuple[ProviderType, KeyProvider, ModelName, str], Model] = {}
 
 
+def get_default_model_for_provider(config: ShotgunConfig) -> ModelName:
+    """Get the default model based on which provider/account is configured.
+
+    Checks API keys in priority order and returns appropriate default model.
+    Treats Shotgun Account as a provider context.
+
+    Args:
+        config: Shotgun configuration containing API keys
+
+    Returns:
+        Default ModelName for the configured provider/account
+    """
+    # Priority 1: Shotgun Account
+    if _get_api_key(config.shotgun.api_key):
+        return ModelName.GPT_5_1
+
+    # Priority 2: Individual provider keys
+    if _get_api_key(config.anthropic.api_key):
+        return ModelName.CLAUDE_HAIKU_4_5
+    if _get_api_key(config.openai.api_key):
+        return ModelName.GPT_5_1
+    if _get_api_key(config.google.api_key):
+        return ModelName.GEMINI_2_5_PRO
+
+    # Fallback: system-wide default
+    return ModelName.CLAUDE_HAIKU_4_5
+
+
 def get_or_create_model(
     provider: ProviderType,
     key_provider: "KeyProvider",
@@ -142,7 +171,7 @@ def get_or_create_model(
     return _model_cache[cache_key]
 
 
-def get_provider_model(
+async def get_provider_model(
     provider_or_model: ProviderType | ModelName | None = None,
 ) -> ModelConfig:
     """Get a fully configured ModelConfig with API key and Model instance.
@@ -161,7 +190,7 @@ def get_provider_model(
     """
     config_manager = get_config_manager()
     # Use cached config for read-only access (performance)
-    config = config_manager.load(force_reload=False)
+    config = await config_manager.load(force_reload=False)
 
     # Priority 1: Check if Shotgun key exists - if so, use it for ANY model
     shotgun_api_key = _get_api_key(config.shotgun.api_key)
@@ -172,13 +201,16 @@ def get_provider_model(
             model_name = provider_or_model
         else:
             # No specific model requested - use selected or default
-            model_name = config.selected_model or ModelName.CLAUDE_SONNET_4_5
+            model_name = config.selected_model or get_default_model_for_provider(config)
 
+        # Gracefully fall back if the selected model doesn't exist (backwards compatibility)
         if model_name not in MODEL_SPECS:
-            raise ValueError(f"Model '{model_name.value}' not found")
+            model_name = get_default_model_for_provider(config)
+
         spec = MODEL_SPECS[model_name]
 
         # Use Shotgun Account with determined model (provider = actual LLM provider)
+        # Shotgun accounts always support streaming (via LiteLLM proxy)
         return ModelConfig(
             name=spec.name,
             provider=spec.provider,  # Actual LLM provider (OPENAI/ANTHROPIC/GOOGLE)
@@ -186,6 +218,7 @@ def get_provider_model(
             max_input_tokens=spec.max_input_tokens,
             max_output_tokens=spec.max_output_tokens,
             api_key=shotgun_api_key,
+            supports_streaming=True,  # Shotgun accounts always support streaming
         )
 
     # Priority 2: Fall back to individual provider keys
@@ -194,10 +227,12 @@ def get_provider_model(
     if isinstance(provider_or_model, ModelName):
         # Look up the model spec
         if provider_or_model not in MODEL_SPECS:
-            raise ValueError(f"Model '{provider_or_model.value}' not found")
-        spec = MODEL_SPECS[provider_or_model]
-        provider_enum = spec.provider
-        requested_model = provider_or_model
+            requested_model = None  # Fall back to provider default
+            provider_enum = None  # Will be determined below
+        else:
+            spec = MODEL_SPECS[provider_or_model]
+            provider_enum = spec.provider
+            requested_model = provider_or_model
     else:
         # Convert string to ProviderType enum if needed (backward compatible)
         if provider_or_model:
@@ -226,12 +261,40 @@ def get_provider_model(
         if not api_key:
             raise ValueError("OpenAI API key not configured. Set via config.")
 
-        # Use requested model or default to gpt-5
-        model_name = requested_model if requested_model else ModelName.GPT_5
+        # Use requested model or default to gpt-5.1
+        model_name = requested_model if requested_model else ModelName.GPT_5_1
+        # Gracefully fall back if model doesn't exist
         if model_name not in MODEL_SPECS:
-            raise ValueError(f"Model '{model_name.value}' not found")
+            model_name = ModelName.GPT_5_1
         spec = MODEL_SPECS[model_name]
 
+        # Check and test streaming capability for GPT-5 family models
+        supports_streaming = True  # Default to True for all models
+        if model_name in (
+            ModelName.GPT_5_1,
+            ModelName.GPT_5_1_CODEX,
+            ModelName.GPT_5_1_CODEX_MINI,
+        ):
+            # Check if streaming capability has been tested
+            streaming_capability = config.openai.supports_streaming
+
+            if streaming_capability is None:
+                # Not tested yet - run streaming test (test once for all GPT-5 models)
+                logger.info("Testing streaming capability for OpenAI GPT-5 family...")
+                streaming_capability = await check_streaming_capability(
+                    api_key, model_name.value
+                )
+
+                # Save result to config (applies to all OpenAI models)
+                config.openai.supports_streaming = streaming_capability
+                await config_manager.save(config)
+                logger.info(
+                    f"Streaming test result: "
+                    f"{'enabled' if streaming_capability else 'disabled'}"
+                )
+
+            supports_streaming = streaming_capability
+
         # Create fully configured ModelConfig
         return ModelConfig(
             name=spec.name,
@@ -240,6 +303,7 @@ def get_provider_model(
             max_input_tokens=spec.max_input_tokens,
             max_output_tokens=spec.max_output_tokens,
             api_key=api_key,
+            supports_streaming=supports_streaming,
         )
 
     elif provider_enum == ProviderType.ANTHROPIC:
@@ -247,10 +311,11 @@ def get_provider_model(
         if not api_key:
             raise ValueError("Anthropic API key not configured. Set via config.")
 
-        # Use requested model or default to claude-sonnet-4-5
-        model_name = requested_model if requested_model else ModelName.CLAUDE_SONNET_4_5
+        # Use requested model or default to claude-haiku-4-5
+        model_name = requested_model if requested_model else ModelName.CLAUDE_HAIKU_4_5
+        # Gracefully fall back if model doesn't exist
         if model_name not in MODEL_SPECS:
-            raise ValueError(f"Model '{model_name.value}' not found")
+            model_name = ModelName.CLAUDE_HAIKU_4_5
         spec = MODEL_SPECS[model_name]
 
         # Create fully configured ModelConfig
@@ -270,8 +335,9 @@ def get_provider_model(
 
         # Use requested model or default to gemini-2.5-pro
         model_name = requested_model if requested_model else ModelName.GEMINI_2_5_PRO
+        # Gracefully fall back if model doesn't exist
         if model_name not in MODEL_SPECS:
-            raise ValueError(f"Model '{model_name.value}' not found")
+            model_name = ModelName.GEMINI_2_5_PRO
         spec = MODEL_SPECS[model_name]
 
         # Create fully configured ModelConfig

shotgun/agents/config/streaming_test.py

@@ -0,0 +1,119 @@
+"""Utility for testing streaming capability of OpenAI models."""
+
+import logging
+
+import httpx
+
+logger = logging.getLogger(__name__)
+
+# Maximum number of attempts to test streaming capability
+MAX_STREAMING_TEST_ATTEMPTS = 3
+
+# Timeout for each streaming test attempt (in seconds)
+STREAMING_TEST_TIMEOUT = 10.0
+
+
+async def check_streaming_capability(
+    api_key: str, model_name: str, max_attempts: int = MAX_STREAMING_TEST_ATTEMPTS
+) -> bool:
+    """Check if the given OpenAI model supports streaming with this API key.
+
+    Retries multiple times to handle transient network issues. Only returns False
+    if streaming definitively fails after all retry attempts.
+
+    Args:
+        api_key: The OpenAI API key to test
+        model_name: The model name (e.g., "gpt-5", "gpt-5-mini")
+        max_attempts: Maximum number of attempts (default: 3)
+
+    Returns:
+        True if streaming is supported, False if it definitively fails
+    """
+    url = "https://api.openai.com/v1/chat/completions"
+    headers = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+    }
+    # GPT-5 family uses max_completion_tokens instead of max_tokens
+    payload = {
+        "model": model_name,
+        "messages": [{"role": "user", "content": "test"}],
+        "stream": True,
+        "max_completion_tokens": 10,
+    }
+
+    last_error = None
+
+    for attempt in range(1, max_attempts + 1):
+        logger.debug(
+            f"Streaming test attempt {attempt}/{max_attempts} for {model_name}"
+        )
+
+        try:
+            async with httpx.AsyncClient(timeout=STREAMING_TEST_TIMEOUT) as client:
+                async with client.stream(
+                    "POST", url, json=payload, headers=headers
+                ) as response:
+                    # Check if we get a successful response
+                    if response.status_code != 200:
+                        last_error = f"HTTP {response.status_code}"
+                        logger.warning(
+                            f"Streaming test attempt {attempt} failed for {model_name}: {last_error}"
+                        )
+
+                        # For definitive errors (403 Forbidden, 404 Not Found), don't retry
+                        if response.status_code in (403, 404):
+                            logger.info(
+                                f"Streaming definitively unsupported for {model_name} (HTTP {response.status_code})"
+                            )
+                            return False
+
+                        # For other errors, retry
+                        continue
+
+                    # Try to read at least one chunk from the stream
+                    try:
+                        async for _ in response.aiter_bytes():
+                            # Successfully received streaming data
+                            logger.info(
+                                f"Streaming test passed for {model_name} (attempt {attempt})"
+                            )
+                            return True
+                    except Exception as e:
+                        last_error = str(e)
+                        logger.warning(
+                            f"Streaming test attempt {attempt} failed for {model_name} while reading stream: {e}"
+                        )
+                        continue
+
+        except httpx.TimeoutException:
+            last_error = "timeout"
+            logger.warning(
+                f"Streaming test attempt {attempt} timed out for {model_name}"
+            )
+            continue
+        except httpx.HTTPStatusError as e:
+            last_error = str(e)
+            logger.warning(
+                f"Streaming test attempt {attempt} failed for {model_name}: {e}"
+            )
+            continue
+        except Exception as e:
+            last_error = str(e)
+            logger.warning(
+                f"Streaming test attempt {attempt} failed for {model_name} with unexpected error: {e}"
+            )
+            continue
+
+        # If we got here without reading any chunks, streaming didn't work
+        last_error = "no data received"
+        logger.warning(
+            f"Streaming test attempt {attempt} failed for {model_name}: no data received"
+        )
+
+    # All attempts exhausted
+    logger.error(
+        f"Streaming test failed for {model_name} after {max_attempts} attempts. "
+        f"Last error: {last_error}. Assuming streaming is NOT supported."
+    )
+    return False

shotgun/agents/context_analyzer/__init__.py

@@ -0,0 +1,28 @@
+"""Context analysis module for conversation composition statistics.
+
+This module provides tools for analyzing conversation context usage, breaking down
+token consumption by message type and tool category.
+"""
+
+from .analyzer import ContextAnalyzer
+from .constants import ToolCategory, get_tool_category
+from .formatter import ContextFormatter
+from .models import (
+    ContextAnalysis,
+    ContextAnalysisOutput,
+    ContextCompositionTelemetry,
+    MessageTypeStats,
+    TokenAllocation,
+)
+
+__all__ = [
+    "ContextAnalyzer",
+    "ContextAnalysis",
+    "ContextAnalysisOutput",
+    "ContextCompositionTelemetry",
+    "ContextFormatter",
+    "MessageTypeStats",
+    "TokenAllocation",
+    "ToolCategory",
+    "get_tool_category",
+]