shotgun-sh 0.1.0__py3-none-any.whl

shotgun/agents/config/manager.py

@@ -0,0 +1,294 @@
+"""Configuration manager for Shotgun CLI."""
+
+import json
+import os
+import uuid
+from pathlib import Path
+from typing import Any
+
+from pydantic import SecretStr
+
+from shotgun.logging_config import get_logger
+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,
+)
+from .models import ProviderType, ShotgunConfig
+
+logger = get_logger(__name__)
+
+
+class ConfigManager:
+    """Manager for Shotgun configuration."""
+
+    def __init__(self, config_path: Path | None = None):
+        """Initialize ConfigManager.
+
+        Args:
+            config_path: Path to config file. If None, uses default ~/.shotgun-sh/config.json
+        """
+        if config_path is None:
+            self.config_path = get_shotgun_home() / "config.json"
+        else:
+            self.config_path = config_path
+
+        self._config: ShotgunConfig | None = None
+
+    def load(self) -> ShotgunConfig:
+        """Load configuration from file.
+
+        Returns:
+            ShotgunConfig: Loaded configuration or default config if file doesn't exist
+        """
+        if self._config is not None:
+            return self._config
+
+        if not self.config_path.exists():
+            logger.info(
+                "Configuration file not found, creating new config with user_id: %s",
+                self.config_path,
+            )
+            # Create new config with generated user_id
+            self._config = self.initialize()
+            return self._config
+
+        try:
+            with open(self.config_path, encoding="utf-8") as f:
+                data = json.load(f)
+
+            # Convert plain text secrets to SecretStr objects
+            self._convert_secrets_to_secretstr(data)
+
+            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):
+                        logger.info(
+                            "Default provider %s has no API key, updating to %s",
+                            original_default.value,
+                            provider.value,
+                        )
+                        self._config.default_provider = provider
+                        self.save(self._config)
+                        break
+
+            return self._config
+
+        except Exception as e:
+            logger.error(
+                "Failed to load configuration from %s: %s", self.config_path, e
+            )
+            logger.info("Creating new configuration with generated user_id")
+            self._config = self.initialize()
+            return self._config
+
+    def save(self, config: ShotgunConfig | None = None) -> None:
+        """Save configuration to file.
+
+        Args:
+            config: Configuration to save. If None, saves current loaded config
+        """
+        if config is None:
+            if self._config:
+                config = self._config
+            else:
+                # Create a new config with generated user_id
+                config = ShotgunConfig(
+                    user_id=str(uuid.uuid4()),
+                    config_version=1,
+                )
+
+        # Ensure directory exists
+        self.config_path.parent.mkdir(parents=True, exist_ok=True)
+
+        try:
+            # Convert SecretStr to plain text for JSON serialization
+            data = config.model_dump()
+            self._convert_secretstr_to_plain(data)
+
+            with open(self.config_path, "w", encoding="utf-8") as f:
+                json.dump(data, f, indent=2, ensure_ascii=False)
+
+            logger.debug("Configuration saved to %s", self.config_path)
+            self._config = config
+
+        except Exception as e:
+            logger.error("Failed to save configuration to %s: %s", self.config_path, e)
+            raise
+
+    def update_provider(self, provider: ProviderType | str, **kwargs: Any) -> None:
+        """Update provider configuration.
+
+        Args:
+            provider: Provider to update
+            **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)
+
+        # Only support api_key updates
+        if API_KEY_FIELD in kwargs:
+            api_key_value = kwargs[API_KEY_FIELD]
+            provider_config.api_key = (
+                SecretStr(api_key_value) if api_key_value is not None else None
+            )
+
+        # Reject other fields
+        unsupported_fields = set(kwargs.keys()) - {API_KEY_FIELD}
+        if unsupported_fields:
+            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:
+            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
+
+        self.save(config)
+
+    def clear_provider_key(self, provider: ProviderType | str) -> None:
+        """Remove the API key for the given provider."""
+        config = self.load()
+        provider_enum = self._ensure_provider_enum(provider)
+        provider_config = self._get_provider_config(config, provider_enum)
+        provider_config.api_key = None
+        self.save(config)
+
+    def has_provider_key(self, provider: ProviderType | str) -> bool:
+        """Check if the given provider has a non-empty API key configured.
+
+        This checks both the configuration file and environment variables.
+        """
+        config = self.load()
+        provider_enum = self._ensure_provider_enum(provider)
+        provider_config = self._get_provider_config(config, provider_enum)
+
+        # Check config first
+        if self._provider_has_api_key(provider_config):
+            return True
+
+        # Check environment variable
+        if provider_enum == ProviderType.OPENAI:
+            return bool(os.getenv(OPENAI_API_KEY_ENV))
+        elif provider_enum == ProviderType.ANTHROPIC:
+            return bool(os.getenv(ANTHROPIC_API_KEY_ENV))
+        elif provider_enum == ProviderType.GOOGLE:
+            return bool(os.getenv(GEMINI_API_KEY_ENV))
+
+        return False
+
+    def has_any_provider_key(self) -> bool:
+        """Determine whether any provider has a configured API key."""
+        config = self.load()
+        return any(
+            self._provider_has_api_key(self._get_provider_config(config, provider))
+            for provider in (
+                ProviderType.OPENAI,
+                ProviderType.ANTHROPIC,
+                ProviderType.GOOGLE,
+            )
+        )
+
+    def initialize(self) -> ShotgunConfig:
+        """Initialize configuration with defaults and save to file.
+
+        Returns:
+            Default ShotgunConfig
+        """
+        # Generate unique user ID for new config
+        config = ShotgunConfig(
+            user_id=str(uuid.uuid4()),
+            config_version=1,
+        )
+        self.save(config)
+        logger.info(
+            "Configuration initialized at %s with user_id: %s",
+            self.config_path,
+            config.user_id,
+        )
+        return config
+
+    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):
+                if (
+                    API_KEY_FIELD in data[provider]
+                    and data[provider][API_KEY_FIELD] is not None
+                ):
+                    data[provider][API_KEY_FIELD] = SecretStr(
+                        data[provider][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):
+                if (
+                    API_KEY_FIELD in data[provider]
+                    and data[provider][API_KEY_FIELD] is not None
+                ):
+                    if hasattr(data[provider][API_KEY_FIELD], "get_secret_value"):
+                        data[provider][API_KEY_FIELD] = data[provider][
+                            API_KEY_FIELD
+                        ].get_secret_value()
+
+    def _ensure_provider_enum(self, provider: ProviderType | str) -> ProviderType:
+        """Normalize provider values to ProviderType enum."""
+        return (
+            provider if isinstance(provider, ProviderType) else ProviderType(provider)
+        )
+
+    def _get_provider_config(
+        self, config: ShotgunConfig, provider: ProviderType
+    ) -> Any:
+        """Retrieve the provider-specific configuration section."""
+        if provider == ProviderType.OPENAI:
+            return config.openai
+        if provider == ProviderType.ANTHROPIC:
+            return config.anthropic
+        if provider == ProviderType.GOOGLE:
+            return config.google
+        raise ValueError(f"Unsupported provider: {provider}")
+
+    def _provider_has_api_key(self, provider_config: Any) -> bool:
+        """Return True if the provider config contains a usable API key."""
+        api_key = getattr(provider_config, API_KEY_FIELD, None)
+        if api_key is None:
+            return False
+
+        if isinstance(api_key, SecretStr):
+            value = api_key.get_secret_value()
+        else:
+            value = str(api_key)
+
+        return bool(value.strip())
+
+    def get_user_id(self) -> str:
+        """Get the user ID from configuration.
+
+        Returns:
+            The unique user ID string
+        """
+        config = self.load()
+        return config.user_id
+
+
+def get_config_manager() -> ConfigManager:
+    """Get the global ConfigManager instance."""
+    return ConfigManager()

shotgun/agents/config/models.py

@@ -0,0 +1,185 @@
+"""Pydantic models for configuration."""
+
+from enum import Enum
+from typing import Any
+
+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):
+    """Provider types for AI services."""
+
+    OPENAI = "openai"
+    ANTHROPIC = "anthropic"
+    GOOGLE = "google"
+
+
+class ModelSpec(BaseModel):
+    """Static specification for a model - just metadata."""
+
+    name: str  # Model identifier (e.g., "gpt-5", "claude-opus-4-1")
+    provider: ProviderType
+    max_input_tokens: int
+    max_output_tokens: int
+
+
+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
+    max_input_tokens: int
+    max_output_tokens: int
+    api_key: str
+    _model_instance: Model | None = PrivateAttr(default=None)
+
+    class Config:
+        arbitrary_types_allowed = True
+
+    @property
+    def model_instance(self) -> Model:
+        """Lazy load the Model instance."""
+        if self._model_instance is None:
+            from .provider import get_or_create_model
+
+            self._model_instance = get_or_create_model(
+                self.provider, self.name, self.api_key
+            )
+        return self._model_instance
+
+    @property
+    def pydantic_model_name(self) -> str:
+        """Compute the full Pydantic AI model identifier. For backward compatibility."""
+        provider_prefix = {
+            ProviderType.OPENAI: "openai",
+            ProviderType.ANTHROPIC: "anthropic",
+            ProviderType.GOOGLE: "google-gla",
+        }
+        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",
+        provider=ProviderType.OPENAI,
+        max_input_tokens=400_000,
+        max_output_tokens=128_000,
+    ),
+    "gpt-4o": ModelSpec(
+        name="gpt-4o",
+        provider=ProviderType.OPENAI,
+        max_input_tokens=128_000,
+        max_output_tokens=16_000,
+    ),
+    "claude-opus-4-1": ModelSpec(
+        name="claude-opus-4-1",
+        provider=ProviderType.ANTHROPIC,
+        max_input_tokens=200_000,
+        max_output_tokens=32_000,
+    ),
+    "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,
+    ),
+    "gemini-2.5-pro": ModelSpec(
+        name="gemini-2.5-pro",
+        provider=ProviderType.GOOGLE,
+        max_input_tokens=1_000_000,
+        max_output_tokens=64_000,
+    ),
+}
+
+
+class OpenAIConfig(BaseModel):
+    """Configuration for OpenAI provider."""
+
+    api_key: SecretStr | None = None
+
+
+class AnthropicConfig(BaseModel):
+    """Configuration for Anthropic provider."""
+
+    api_key: SecretStr | None = None
+
+
+class GoogleConfig(BaseModel):
+    """Configuration for Google provider."""
+
+    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"
+    )
+    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,
+    )

shotgun/agents/config/provider.py

@@ -0,0 +1,206 @@
+"""Provider management for LLM configuration."""
+
+import os
+
+from pydantic import SecretStr
+from pydantic_ai.models import Model
+from pydantic_ai.models.anthropic import AnthropicModel, AnthropicModelSettings
+from pydantic_ai.models.google import GoogleModel
+from pydantic_ai.models.openai import OpenAIChatModel
+from pydantic_ai.providers.anthropic import AnthropicProvider
+from pydantic_ai.providers.google import GoogleProvider
+from pydantic_ai.providers.openai import OpenAIProvider
+from pydantic_ai.settings import ModelSettings
+
+from shotgun.logging_config import get_logger
+
+from .constants import (
+    ANTHROPIC_API_KEY_ENV,
+    GEMINI_API_KEY_ENV,
+    OPENAI_API_KEY_ENV,
+)
+from .manager import get_config_manager
+from .models import MODEL_SPECS, ModelConfig, ProviderType
+
+logger = get_logger(__name__)
+
+# Global cache for Model instances (singleton pattern)
+_model_cache: dict[tuple[ProviderType, str, str], Model] = {}
+
+
+def get_or_create_model(provider: ProviderType, model_name: str, api_key: str) -> Model:
+    """Get or create a singleton Model instance.
+
+    Args:
+        provider: Provider type
+        model_name: Name of the model
+        api_key: API key for the provider
+
+    Returns:
+        Cached or newly created Model instance
+
+    Raises:
+        ValueError: If provider is not supported
+    """
+    cache_key = (provider, model_name, api_key)
+
+    if cache_key not in _model_cache:
+        logger.debug("Creating new %s model instance: %s", provider.value, model_name)
+
+        if provider == ProviderType.OPENAI:
+            # Get max_tokens from MODEL_SPECS to use full capacity
+            if model_name in MODEL_SPECS:
+                max_tokens = MODEL_SPECS[model_name].max_output_tokens
+            else:
+                max_tokens = 16_000  # Default for GPT models
+
+            openai_provider = OpenAIProvider(api_key=api_key)
+            _model_cache[cache_key] = OpenAIChatModel(
+                model_name,
+                provider=openai_provider,
+                settings=ModelSettings(max_tokens=max_tokens),
+            )
+        elif provider == ProviderType.ANTHROPIC:
+            # Get max_tokens from MODEL_SPECS to use full capacity
+            if model_name in MODEL_SPECS:
+                max_tokens = MODEL_SPECS[model_name].max_output_tokens
+            else:
+                max_tokens = 32_000  # Default for Claude models
+
+            anthropic_provider = AnthropicProvider(api_key=api_key)
+            _model_cache[cache_key] = AnthropicModel(
+                model_name,
+                provider=anthropic_provider,
+                settings=AnthropicModelSettings(
+                    max_tokens=max_tokens,
+                    timeout=600,  # 10 minutes timeout for large responses
+                ),
+            )
+        elif provider == ProviderType.GOOGLE:
+            # Get max_tokens from MODEL_SPECS to use full capacity
+            if model_name in MODEL_SPECS:
+                max_tokens = MODEL_SPECS[model_name].max_output_tokens
+            else:
+                max_tokens = 64_000  # Default for Gemini models
+
+            google_provider = GoogleProvider(api_key=api_key)
+            _model_cache[cache_key] = GoogleModel(
+                model_name,
+                provider=google_provider,
+                settings=ModelSettings(max_tokens=max_tokens),
+            )
+        else:
+            raise ValueError(f"Unsupported provider: {provider}")
+    else:
+        logger.debug("Reusing cached %s model instance: %s", provider.value, model_name)
+
+    return _model_cache[cache_key]
+
+
+def get_provider_model(provider: ProviderType | None = None) -> ModelConfig:
+    """Get a fully configured ModelConfig with API key and Model instance.
+
+    Args:
+        provider: Provider to get model for. If None, uses default provider
+
+    Returns:
+        ModelConfig with API key configured and lazy Model instance
+
+    Raises:
+        ValueError: If provider is not configured properly or model not found
+    """
+    config_manager = get_config_manager()
+    config = config_manager.load()
+    # Convert string to ProviderType enum if needed
+    provider_enum = (
+        provider
+        if isinstance(provider, ProviderType)
+        else ProviderType(provider)
+        if provider
+        else config.default_provider
+    )
+
+    if provider_enum == ProviderType.OPENAI:
+        api_key = _get_api_key(config.openai.api_key, OPENAI_API_KEY_ENV)
+        if not api_key:
+            raise ValueError(
+                f"OpenAI API key not configured. Set via environment variable {OPENAI_API_KEY_ENV} or config."
+            )
+
+        # Get model spec - hardcoded to gpt-5
+        model_name = "gpt-5"
+        if model_name not in MODEL_SPECS:
+            raise ValueError(f"Model '{model_name}' not found")
+        spec = MODEL_SPECS[model_name]
+
+        # Create fully configured ModelConfig
+        return ModelConfig(
+            name=spec.name,
+            provider=spec.provider,
+            max_input_tokens=spec.max_input_tokens,
+            max_output_tokens=spec.max_output_tokens,
+            api_key=api_key,
+        )
+
+    elif provider_enum == ProviderType.ANTHROPIC:
+        api_key = _get_api_key(config.anthropic.api_key, ANTHROPIC_API_KEY_ENV)
+        if not api_key:
+            raise ValueError(
+                f"Anthropic API key not configured. Set via environment variable {ANTHROPIC_API_KEY_ENV} or config."
+            )
+
+        # Get model spec - hardcoded to claude-opus-4-1
+        model_name = "claude-opus-4-1"
+        if model_name not in MODEL_SPECS:
+            raise ValueError(f"Model '{model_name}' not found")
+        spec = MODEL_SPECS[model_name]
+
+        # Create fully configured ModelConfig
+        return ModelConfig(
+            name=spec.name,
+            provider=spec.provider,
+            max_input_tokens=spec.max_input_tokens,
+            max_output_tokens=spec.max_output_tokens,
+            api_key=api_key,
+        )
+
+    elif provider_enum == ProviderType.GOOGLE:
+        api_key = _get_api_key(config.google.api_key, GEMINI_API_KEY_ENV)
+        if not api_key:
+            raise ValueError(
+                f"Gemini API key not configured. Set via environment variable {GEMINI_API_KEY_ENV} or config."
+            )
+
+        # Get model spec - hardcoded to gemini-2.5-pro
+        model_name = "gemini-2.5-pro"
+        if model_name not in MODEL_SPECS:
+            raise ValueError(f"Model '{model_name}' not found")
+        spec = MODEL_SPECS[model_name]
+
+        # Create fully configured ModelConfig
+        return ModelConfig(
+            name=spec.name,
+            provider=spec.provider,
+            max_input_tokens=spec.max_input_tokens,
+            max_output_tokens=spec.max_output_tokens,
+            api_key=api_key,
+        )
+
+    else:
+        raise ValueError(f"Unsupported provider: {provider_enum}")
+
+
+def _get_api_key(config_key: SecretStr | None, env_var: str) -> str | None:
+    """Get API key from config or environment variable.
+
+    Args:
+        config_key: API key from configuration
+        env_var: Environment variable name to check
+
+    Returns:
+        API key string or None
+    """
+    if config_key is not None:
+        return config_key.get_secret_value()
+
+    return os.getenv(env_var)