shotgun-sh 0.1.0.dev1__py3-none-any.whl

shotgun/agents/config/manager.py

@@ -0,0 +1,215 @@
+"""Configuration manager for Shotgun CLI."""
+
+import json
+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 .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, using defaults: %s", self.config_path
+            )
+            self._config = ShotgunConfig()
+            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)
+            return self._config
+
+        except Exception as e:
+            logger.error(
+                "Failed to load configuration from %s: %s", self.config_path, e
+            )
+            logger.info("Using default configuration")
+            self._config = ShotgunConfig()
+            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:
+            config = self._config or ShotgunConfig()
+
+        # 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" in kwargs:
+            api_key_value = kwargs["api_key"]
+            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"}
+        if unsupported_fields:
+            raise ValueError(f"Unsupported configuration fields: {unsupported_fields}")
+
+        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."""
+        config = self.load()
+        provider_enum = self._ensure_provider_enum(provider)
+        provider_config = self._get_provider_config(config, provider_enum)
+        return self._provider_has_api_key(provider_config)
+
+    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
+        """
+        config = ShotgunConfig()
+        self.save(config)
+        logger.info("Configuration initialized at %s", self.config_path)
+        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", "anthropic", "google"]:
+            if provider in data and isinstance(data[provider], dict):
+                if (
+                    "api_key" in data[provider]
+                    and data[provider]["api_key"] is not None
+                ):
+                    data[provider]["api_key"] = SecretStr(data[provider]["api_key"])
+
+    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", "anthropic", "google"]:
+            if provider in data and isinstance(data[provider], dict):
+                if (
+                    "api_key" in data[provider]
+                    and data[provider]["api_key"] is not None
+                ):
+                    if hasattr(data[provider]["api_key"], "get_secret_value"):
+                        data[provider]["api_key"] = data[provider][
+                            "api_key"
+                        ].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", 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_config_manager() -> ConfigManager:
+    """Get the global ConfigManager instance."""
+    return ConfigManager()

shotgun/agents/config/models.py

@@ -0,0 +1,120 @@
+"""Pydantic models for configuration."""
+
+from enum import Enum
+
+from pydantic import BaseModel, Field, SecretStr
+
+
+class ProviderType(str, Enum):
+    """Provider types for AI services."""
+
+    OPENAI = "openai"
+    ANTHROPIC = "anthropic"
+    GOOGLE = "google"
+
+
+class ModelConfig(BaseModel):
+    """Configuration for an LLM model."""
+
+    name: str  # Model identifier (e.g., "gpt-5", "claude-opus-4-1")
+    provider: ProviderType
+    max_input_tokens: int
+    max_output_tokens: int
+
+    @property
+    def pydantic_model_name(self) -> str:
+        """Compute the full Pydantic AI model identifier."""
+        provider_prefix = {
+            ProviderType.OPENAI: "openai",
+            ProviderType.ANTHROPIC: "anthropic",
+            ProviderType.GOOGLE: "google-gla",
+        }
+        return f"{provider_prefix[self.provider]}:{self.name}"
+
+
+# OpenAI Models
+GPT_5 = ModelConfig(
+    name="gpt-5",
+    provider=ProviderType.OPENAI,
+    max_input_tokens=400_000,
+    max_output_tokens=128_000,
+)
+
+GPT_4O = ModelConfig(
+    name="gpt-4o",
+    provider=ProviderType.OPENAI,
+    max_input_tokens=128_000,
+    max_output_tokens=16_000,
+)
+
+# Anthropic Models
+CLAUDE_OPUS_4_1 = ModelConfig(
+    name="claude-opus-4-1",
+    provider=ProviderType.ANTHROPIC,
+    max_input_tokens=200_000,
+    max_output_tokens=32_000,
+)
+
+CLAUDE_3_5_SONNET = ModelConfig(
+    name="claude-3-5-sonnet-latest",
+    provider=ProviderType.ANTHROPIC,
+    max_input_tokens=200_000,
+    max_output_tokens=20_000,
+)
+
+# Google Models
+GEMINI_2_5_PRO = ModelConfig(
+    name="gemini-2.5-pro",
+    provider=ProviderType.GOOGLE,
+    max_input_tokens=1_000_000,
+    max_output_tokens=64_000,
+)
+
+# List of all available models
+AVAILABLE_MODELS = [
+    GPT_5,
+    GPT_4O,
+    CLAUDE_OPUS_4_1,
+    CLAUDE_3_5_SONNET,
+    GEMINI_2_5_PRO,
+]
+
+
+def get_model_by_name(name: str) -> ModelConfig:
+    """Find a model configuration by name."""
+    for model in AVAILABLE_MODELS:
+        if model.name == name:
+            return model
+    raise ValueError(f"Model '{name}' not found")
+
+
+class OpenAIConfig(BaseModel):
+    """Configuration for OpenAI provider."""
+
+    api_key: SecretStr | None = None
+    model_name: str = "gpt-5"
+
+
+class AnthropicConfig(BaseModel):
+    """Configuration for Anthropic provider."""
+
+    api_key: SecretStr | None = None
+    model_name: str = "claude-opus-4-1"
+
+
+class GoogleConfig(BaseModel):
+    """Configuration for Google provider."""
+
+    api_key: SecretStr | None = None
+    model_name: str = "gemini-2.5-pro"
+
+
+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/agents/config/provider.py

@@ -0,0 +1,91 @@
+"""Provider management for LLM configuration."""
+
+import os
+
+from pydantic import SecretStr
+
+from shotgun.logging_config import get_logger
+
+from .manager import get_config_manager
+from .models import ModelConfig, ProviderType, get_model_by_name
+
+logger = get_logger(__name__)
+
+
+def get_provider_model(provider: ProviderType | None = None) -> ModelConfig:
+    """Get model configuration for the specified provider.
+
+    Args:
+        provider: Provider to get model for. If None, uses default provider
+
+    Returns:
+        ModelConfig with pydantic_model_name and token limits
+
+    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")
+        if not api_key:
+            raise ValueError(
+                "OpenAI API key not configured. Set via environment variable OPENAI_API_KEY or config."
+            )
+        # Set the API key in environment if not already there
+        if "OPENAI_API_KEY" not in os.environ:
+            os.environ["OPENAI_API_KEY"] = api_key
+
+        return get_model_by_name(config.openai.model_name)
+
+    elif provider_enum == ProviderType.ANTHROPIC:
+        api_key = _get_api_key(config.anthropic.api_key, "ANTHROPIC_API_KEY")
+        if not api_key:
+            raise ValueError(
+                "Anthropic API key not configured. Set via environment variable ANTHROPIC_API_KEY or config."
+            )
+        # Set the API key in environment if not already there
+        if "ANTHROPIC_API_KEY" not in os.environ:
+            os.environ["ANTHROPIC_API_KEY"] = api_key
+
+        return get_model_by_name(config.anthropic.model_name)
+
+    elif provider_enum == ProviderType.GOOGLE:
+        api_key = _get_api_key(config.google.api_key, "GOOGLE_API_KEY")
+        if not api_key:
+            raise ValueError(
+                "Google API key not configured. Set via environment variable GOOGLE_API_KEY or config."
+            )
+        # Set the API key in environment if not already there
+        if "GOOGLE_API_KEY" not in os.environ:
+            os.environ["GOOGLE_API_KEY"] = api_key
+
+        return get_model_by_name(config.google.model_name)
+
+    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)

shotgun/agents/history/__init__.py

@@ -0,0 +1,5 @@
+"""History management utilities for Shotgun agents."""
+
+from .history_processors import token_limit_compactor
+
+__all__ = ["token_limit_compactor"]

shotgun/agents/history/history_processors.py

@@ -0,0 +1,213 @@
+"""History processors for managing conversation history in Shotgun agents."""
+
+from pydantic_ai import RunContext
+from pydantic_ai.direct import model_request
+from pydantic_ai.messages import (
+    BuiltinToolCallPart,
+    BuiltinToolReturnPart,
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    ModelResponsePart,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ThinkingPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+
+from shotgun.agents.models import AgentDeps
+from shotgun.logging_config import get_logger
+from shotgun.prompts import PromptLoader
+
+logger = get_logger(__name__)
+
+# Global prompt loader instance
+prompt_loader = PromptLoader()
+
+
+async def token_limit_compactor(
+    ctx: RunContext[AgentDeps],
+    messages: list[ModelMessage],
+) -> list[ModelMessage]:
+    """Compact message history based on token limits.
+
+    This context-aware processor monitors token usage and removes older messages
+    when the conversation history becomes too large. It preserves system messages
+    and recent context while removing older user/assistant exchanges.
+
+    Args:
+        ctx: Run context with usage information and dependencies
+        messages: List of messages in the conversation history
+
+    Returns:
+        Compacted list of messages within token limits
+    """
+    # Get current token usage from context
+    current_tokens = ctx.usage.total_tokens if ctx.usage else 0
+
+    # Get token limit from model configuration or use fallback
+    model_max_tokens = ctx.deps.llm_model.max_input_tokens
+    max_tokens = int(
+        model_max_tokens * 0.8
+    )  # Use 80% of max to leave room for response
+    percentage_of_limit_used = (
+        (current_tokens / max_tokens) * 100 if max_tokens > 0 else 0
+    )
+    logger.debug(
+        "History compactor: current tokens=%d, limit=%d, percentage used=%.2f%%",
+        current_tokens,
+        max_tokens,
+        percentage_of_limit_used,
+    )
+
+    # If we're under the limit, return all messages
+    if current_tokens < max_tokens:
+        logger.debug("Under token limit, keeping all %d messages", len(messages))
+        return messages
+
+    # Get current token usage from context
+    current_tokens = ctx.usage.total_tokens if ctx.usage else 0
+
+    context = ""
+
+    # Separate system messages from conversation messages
+    for msg in messages:
+        if isinstance(msg, ModelResponse) or isinstance(msg, ModelRequest):
+            for part in msg.parts:
+                message_content = get_context_from_message(part)
+                if not message_content:
+                    continue
+                context += get_context_from_message(part) + "\n"
+        else:
+            # Handle whatever this is
+            pass
+
+    summarization_prompt = prompt_loader.render("history/summarization.j2")
+    summary_response = await model_request(
+        model=ctx.model,
+        messages=[
+            ModelRequest.user_text_prompt(context, instructions=summarization_prompt)
+        ],
+    )
+    # Usage before and after
+    summary_usage = summary_response.usage
+    reduction_percentage = (
+        (current_tokens - summary_usage.output_tokens) / current_tokens
+    ) * 100
+    logger.debug(
+        "Compacted %s tokens into %s tokens for a %.2f percent reduction",
+        current_tokens,
+        summary_usage.output_tokens,
+        reduction_percentage,
+    )
+
+    system_prompt = get_system_promt(messages) or ""
+    user_prompt = get_first_user_request(messages) or ""
+    # Extract content from the first response part safely
+    summarization_part = summary_response.parts[0]
+    return [
+        ModelRequest(
+            parts=[
+                SystemPromptPart(content=system_prompt),
+                UserPromptPart(content=user_prompt),
+            ]
+        ),
+        ModelResponse(
+            parts=[
+                summarization_part,
+            ]
+        ),
+    ]
+
+
+def get_first_user_request(messages: list[ModelMessage]) -> str | None:
+    """Extract first user request from messages.
+
+    Args:
+        messages: List of messages in the conversation history
+    Returns:
+        The first user request as a string.
+    """
+    for msg in messages:
+        if isinstance(msg, ModelRequest):
+            for part in msg.parts:
+                if isinstance(part, UserPromptPart):
+                    if isinstance(part.content, str):
+                        return part.content
+    return None
+
+
+def get_system_promt(messages: list[ModelMessage]) -> str | None:
+    """Extract system prompt from messages.
+
+    Args:
+        messages: List of messages in the conversation history
+
+    Returns:
+        The system prompt as a string.
+    """
+    for msg in messages:
+        if isinstance(msg, ModelRequest):
+            for part in msg.parts:
+                if isinstance(part, SystemPromptPart):
+                    return part.content
+    return None
+
+
+def get_context_from_message(
+    message_part: SystemPromptPart
+    | UserPromptPart
+    | ToolReturnPart
+    | RetryPromptPart
+    | ModelResponsePart,
+) -> str:
+    """Extract context from a message part.
+
+    Args:
+        message: The message part to extract context from.
+
+    Returns:
+        The extracted context as a string.
+    """
+
+    if isinstance(message_part, SystemPromptPart):
+        return ""  # We do not include system prompts in the summary
+    elif isinstance(message_part, UserPromptPart):
+        if isinstance(message_part.content, str):
+            return "<USER_PROMPT>\n" + message_part.content + "\n</USER_PROMPT>"
+        else:
+            return ""
+    elif isinstance(message_part, ToolReturnPart):
+        return "<TOOL_RETURN>\n" + str(message_part.content) + "\n</TOOL_RETURN>"
+    elif isinstance(message_part, RetryPromptPart):
+        if isinstance(message_part.content, str):
+            return "<RETRY_PROMPT>\n" + message_part.content + "\n</RETRY_PROMPT>"
+        return ""
+
+    # TextPart | ToolCallPart | BuiltinToolCallPart | BuiltinToolReturnPart | ThinkingPart
+    if isinstance(message_part, TextPart):
+        return "<ASSISTANT_TEXT>\n" + message_part.content + "\n</ASSISTANT_TEXT>"
+    elif isinstance(message_part, ToolCallPart):
+        if isinstance(message_part.args, dict):
+            args_str = ", ".join(f"{k}={repr(v)}" for k, v in message_part.args.items())
+            tool_call_str = f"{message_part.tool_name}({args_str})"
+        else:
+            tool_call_str = f"{message_part.tool_name}({message_part.args})"
+        return "<TOOL_CALL>\n" + tool_call_str + "\n</TOOL_CALL>"
+    elif isinstance(message_part, BuiltinToolCallPart):
+        return (
+            "<BUILTIN_TOOL_CALL>\n" + message_part.tool_name + "\n</BUILTIN_TOOL_CALL>"
+        )
+    elif isinstance(message_part, BuiltinToolReturnPart):
+        return (
+            "<BUILTIN_TOOL_RETURN>\n"
+            + message_part.tool_name
+            + "\n</BUILTIN_TOOL_RETURN>"
+        )
+    elif isinstance(message_part, ThinkingPart):
+        return "<THINKING>\n" + message_part.content + "\n</THINKING>"
+
+    return ""