shotgun-sh 0.1.16.dev2__py3-none-any.whl → 0.2.1.dev2__py3-none-any.whl

shotgun/agents/config/provider.py

@@ -1,7 +1,5 @@
 """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
@@ -12,27 +10,36 @@ from pydantic_ai.providers.google import GoogleProvider
 from pydantic_ai.providers.openai import OpenAIProvider
 from pydantic_ai.settings import ModelSettings
 
+from shotgun.llm_proxy import create_litellm_provider
 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
+from .models import (
+    MODEL_SPECS,
+    KeyProvider,
+    ModelConfig,
+    ModelName,
+    ProviderType,
+    ShotgunConfig,
+)
 
 logger = get_logger(__name__)
 
 # Global cache for Model instances (singleton pattern)
-_model_cache: dict[tuple[ProviderType, str, str], Model] = {}
+_model_cache: dict[tuple[ProviderType, KeyProvider, ModelName, str], Model] = {}
 
 
-def get_or_create_model(provider: ProviderType, model_name: str, api_key: str) -> Model:
+def get_or_create_model(
+    provider: ProviderType,
+    key_provider: "KeyProvider",
+    model_name: ModelName,
+    api_key: str,
+) -> Model:
     """Get or create a singleton Model instance.
 
     Args:
-        provider: Provider type
+        provider: Actual LLM provider (openai, anthropic, google)
+        key_provider: Authentication method (byok or shotgun)
         model_name: Name of the model
         api_key: API key for the provider
 
@@ -42,66 +49,88 @@ def get_or_create_model(provider: ProviderType, model_name: str, api_key: str) -
     Raises:
         ValueError: If provider is not supported
     """
-    cache_key = (provider, model_name, api_key)
+    cache_key = (provider, 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)
+        logger.debug(
+            "Creating new %s model instance via %s: %s",
+            provider.value,
+            key_provider.value,
+            model_name,
+        )
 
-        if provider == ProviderType.OPENAI:
-            # Get max_tokens from MODEL_SPECS to use full capacity
+        # Get max_tokens from MODEL_SPECS
+        if model_name in MODEL_SPECS:
+            max_tokens = MODEL_SPECS[model_name].max_output_tokens
+        else:
+            # Fallback defaults based on provider
+            max_tokens = {
+                ProviderType.OPENAI: 16_000,
+                ProviderType.ANTHROPIC: 32_000,
+                ProviderType.GOOGLE: 64_000,
+            }.get(provider, 16_000)
+
+        # Use LiteLLM proxy for Shotgun Account, native providers for BYOK
+        if key_provider == KeyProvider.SHOTGUN:
+            # Shotgun Account uses LiteLLM proxy for any model
             if model_name in MODEL_SPECS:
-                max_tokens = MODEL_SPECS[model_name].max_output_tokens
+                litellm_model_name = MODEL_SPECS[model_name].litellm_proxy_model_name
             else:
-                max_tokens = 16_000  # Default for GPT models
+                # Fallback for unmapped models
+                litellm_model_name = f"openai/{model_name.value}"
 
-            openai_provider = OpenAIProvider(api_key=api_key)
+            litellm_provider = create_litellm_provider(api_key)
             _model_cache[cache_key] = OpenAIChatModel(
-                model_name,
-                provider=openai_provider,
+                litellm_model_name,
+                provider=litellm_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
+        elif key_provider == KeyProvider.BYOK:
+            # Use native provider implementations with user's API keys
+            if provider == ProviderType.OPENAI:
+                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:
+                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:
+                google_provider = GoogleProvider(api_key=api_key)
+                _model_cache[cache_key] = GoogleModel(
+                    model_name,
+                    provider=google_provider,
+                    settings=ModelSettings(max_tokens=max_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),
-            )
+                raise ValueError(f"Unsupported provider: {provider}")
         else:
-            raise ValueError(f"Unsupported provider: {provider}")
+            raise ValueError(f"Unsupported key provider: {key_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:
+def get_provider_model(
+    provider_or_model: ProviderType | ModelName | 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
+        provider_or_model: Either a ProviderType, ModelName, or None.
+            - If ModelName: returns that specific model with appropriate API key
+            - If ProviderType: returns default model for that provider (backward compatible)
+            - If None: uses default provider with its default model
 
     Returns:
         ModelConfig with API key configured and lazy Model instance
@@ -111,76 +140,117 @@ def get_provider_model(provider: ProviderType | None = None) -> ModelConfig:
     """
     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
-    )
+
+    # Priority 1: Check if Shotgun key exists - if so, use it for ANY model
+    shotgun_api_key = _get_api_key(config.shotgun.api_key)
+    if shotgun_api_key:
+        # Use selected model or default to claude-sonnet-4-5
+        model_name = config.selected_model or ModelName.CLAUDE_SONNET_4_5
+        if model_name not in MODEL_SPECS:
+            raise ValueError(f"Model '{model_name.value}' not found")
+        spec = MODEL_SPECS[model_name]
+
+        # Use Shotgun Account with selected model (provider = actual LLM provider)
+        return ModelConfig(
+            name=spec.name,
+            provider=spec.provider,  # Actual LLM provider (OPENAI/ANTHROPIC/GOOGLE)
+            key_provider=KeyProvider.SHOTGUN,  # Authenticated via Shotgun Account
+            max_input_tokens=spec.max_input_tokens,
+            max_output_tokens=spec.max_output_tokens,
+            api_key=shotgun_api_key,
+        )
+
+    # Priority 2: Fall back to individual provider keys
+
+    # Check if a specific model was requested
+    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
+    else:
+        # Convert string to ProviderType enum if needed (backward compatible)
+        if provider_or_model:
+            provider_enum = (
+                provider_or_model
+                if isinstance(provider_or_model, ProviderType)
+                else ProviderType(provider_or_model)
+            )
+        else:
+            # No provider specified - find first available provider with a key
+            provider_enum = None
+            for provider in ProviderType:
+                if _has_provider_key(config, provider):
+                    provider_enum = provider
+                    break
+
+            if provider_enum is None:
+                raise ValueError(
+                    "No provider keys configured. Set via environment variables or config."
+                )
+
+        requested_model = None  # Will use provider's default model
 
     if provider_enum == ProviderType.OPENAI:
-        api_key = _get_api_key(config.openai.api_key, OPENAI_API_KEY_ENV)
+        api_key = _get_api_key(config.openai.api_key)
         if not api_key:
-            raise ValueError(
-                f"OpenAI API key not configured. Set via environment variable {OPENAI_API_KEY_ENV} or config."
-            )
+            raise ValueError("OpenAI API key not configured. Set via config.")
 
-        # Get model spec - hardcoded to gpt-5
-        model_name = "gpt-5"
+        # Use requested model or default to gpt-5
+        model_name = requested_model if requested_model else ModelName.GPT_5
         if model_name not in MODEL_SPECS:
-            raise ValueError(f"Model '{model_name}' not found")
+            raise ValueError(f"Model '{model_name.value}' not found")
         spec = MODEL_SPECS[model_name]
 
         # Create fully configured ModelConfig
         return ModelConfig(
             name=spec.name,
             provider=spec.provider,
+            key_provider=KeyProvider.BYOK,
             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)
+        api_key = _get_api_key(config.anthropic.api_key)
         if not api_key:
-            raise ValueError(
-                f"Anthropic API key not configured. Set via environment variable {ANTHROPIC_API_KEY_ENV} or config."
-            )
+            raise ValueError("Anthropic API key not configured. Set via config.")
 
-        # Get model spec - hardcoded to claude-opus-4-1
-        model_name = "claude-opus-4-1"
+        # Use requested model or default to claude-opus-4-1
+        model_name = requested_model if requested_model else ModelName.CLAUDE_OPUS_4_1
         if model_name not in MODEL_SPECS:
-            raise ValueError(f"Model '{model_name}' not found")
+            raise ValueError(f"Model '{model_name.value}' not found")
         spec = MODEL_SPECS[model_name]
 
         # Create fully configured ModelConfig
         return ModelConfig(
             name=spec.name,
             provider=spec.provider,
+            key_provider=KeyProvider.BYOK,
             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)
+        api_key = _get_api_key(config.google.api_key)
         if not api_key:
-            raise ValueError(
-                f"Gemini API key not configured. Set via environment variable {GEMINI_API_KEY_ENV} or config."
-            )
+            raise ValueError("Gemini API key not configured. Set via config.")
 
-        # Get model spec - hardcoded to gemini-2.5-pro
-        model_name = "gemini-2.5-pro"
+        # Use requested model or default to gemini-2.5-pro
+        model_name = requested_model if requested_model else ModelName.GEMINI_2_5_PRO
         if model_name not in MODEL_SPECS:
-            raise ValueError(f"Model '{model_name}' not found")
+            raise ValueError(f"Model '{model_name.value}' not found")
         spec = MODEL_SPECS[model_name]
 
         # Create fully configured ModelConfig
         return ModelConfig(
             name=spec.name,
             provider=spec.provider,
+            key_provider=KeyProvider.BYOK,
             max_input_tokens=spec.max_input_tokens,
             max_output_tokens=spec.max_output_tokens,
             api_key=api_key,
@@ -190,12 +260,30 @@ def get_provider_model(provider: ProviderType | None = None) -> ModelConfig:
         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.
+def _has_provider_key(config: "ShotgunConfig", provider: ProviderType) -> bool:
+    """Check if a provider has a configured API key.
+
+    Args:
+        config: Shotgun configuration
+        provider: Provider to check
+
+    Returns:
+        True if provider has a configured API key
+    """
+    if provider == ProviderType.OPENAI:
+        return bool(_get_api_key(config.openai.api_key))
+    elif provider == ProviderType.ANTHROPIC:
+        return bool(_get_api_key(config.anthropic.api_key))
+    elif provider == ProviderType.GOOGLE:
+        return bool(_get_api_key(config.google.api_key))
+    return False
+
+
+def _get_api_key(config_key: SecretStr | None) -> str | None:
+    """Get API key from config.
 
     Args:
         config_key: API key from configuration
-        env_var: Environment variable name to check
 
     Returns:
         API key string or None
@@ -203,4 +291,4 @@ def _get_api_key(config_key: SecretStr | None, env_var: str) -> str | None:
     if config_key is not None:
         return config_key.get_secret_value()
 
-    return os.getenv(env_var)
+    return None

shotgun/agents/history/compaction.py

@@ -31,7 +31,7 @@ async def apply_persistent_compaction(
 
     try:
         # Count actual token usage using shared utility
-        estimated_tokens = estimate_tokens_from_messages(messages, deps.llm_model)
+        estimated_tokens = await estimate_tokens_from_messages(messages, deps.llm_model)
 
         # Create minimal usage info for compaction check
         usage = RequestUsage(

shotgun/agents/history/history_processors.py

@@ -2,6 +2,7 @@
 
 from typing import TYPE_CHECKING, Any, Protocol
 
+from pydantic_ai import ModelSettings
 from pydantic_ai.messages import (
     ModelMessage,
     ModelRequest,
@@ -10,7 +11,7 @@ from pydantic_ai.messages import (
     UserPromptPart,
 )
 
-from shotgun.agents.config.models import shotgun_model_request
+from shotgun.agents.llm import shotgun_model_request
 from shotgun.agents.messages import AgentSystemPrompt, SystemStatusPrompt
 from shotgun.agents.models import AgentDeps
 from shotgun.logging_config import get_logger
@@ -154,7 +155,7 @@ async def token_limit_compactor(
 
     if last_summary_index is not None:
         # Check if post-summary conversation exceeds threshold for incremental compaction
-        post_summary_tokens = estimate_post_summary_tokens(
+        post_summary_tokens = await estimate_post_summary_tokens(
             messages, last_summary_index, deps.llm_model
         )
         post_summary_percentage = (
@@ -248,7 +249,7 @@ async def token_limit_compactor(
         ]
 
         # Calculate optimal max_tokens for summarization
-        max_tokens = calculate_max_summarization_tokens(
+        max_tokens = await calculate_max_summarization_tokens(
             deps.llm_model, request_messages
         )
 
@@ -261,7 +262,9 @@ async def token_limit_compactor(
         summary_response = await shotgun_model_request(
             model_config=deps.llm_model,
             messages=request_messages,
-            max_tokens=max_tokens,  # Use calculated optimal tokens for summarization
+            model_settings=ModelSettings(
+                max_tokens=max_tokens  # Use calculated optimal tokens for summarization
+            ),
         )
 
         log_summarization_response(summary_response, "INCREMENTAL")
@@ -328,7 +331,9 @@ async def token_limit_compactor(
 
         # Track compaction completion
         messages_after = len(compacted_messages)
-        tokens_after = estimate_tokens_from_messages(compacted_messages, deps.llm_model)
+        tokens_after = await estimate_tokens_from_messages(
+            compacted_messages, deps.llm_model
+        )
         reduction_percentage = (
             ((messages_before - messages_after) / messages_before * 100)
             if messages_before > 0
@@ -354,7 +359,7 @@ async def token_limit_compactor(
 
     else:
         # Check if total conversation exceeds threshold for full compaction
-        total_tokens = estimate_tokens_from_messages(messages, deps.llm_model)
+        total_tokens = await estimate_tokens_from_messages(messages, deps.llm_model)
         total_percentage = (total_tokens / max_tokens) * 100 if max_tokens > 0 else 0
 
         logger.debug(
@@ -392,7 +397,9 @@ async def _full_compaction(
     ]
 
     # Calculate optimal max_tokens for summarization
-    max_tokens = calculate_max_summarization_tokens(deps.llm_model, request_messages)
+    max_tokens = await calculate_max_summarization_tokens(
+        deps.llm_model, request_messages
+    )
 
     # Debug logging using shared utilities
     log_summarization_request(
@@ -403,11 +410,13 @@ async def _full_compaction(
     summary_response = await shotgun_model_request(
         model_config=deps.llm_model,
         messages=request_messages,
-        max_tokens=max_tokens,  # Use calculated optimal tokens for summarization
+        model_settings=ModelSettings(
+            max_tokens=max_tokens  # Use calculated optimal tokens for summarization
+        ),
     )
 
     # Calculate token reduction
-    current_tokens = estimate_tokens_from_messages(messages, deps.llm_model)
+    current_tokens = await estimate_tokens_from_messages(messages, deps.llm_model)
     summary_usage = summary_response.usage
     reduction_percentage = (
         ((current_tokens - summary_usage.output_tokens) / current_tokens) * 100

shotgun/agents/history/token_counting/__init__.py

@@ -0,0 +1,31 @@
+"""Real token counting for all supported providers.
+
+This module provides accurate token counting using each provider's official
+APIs and libraries, eliminating the need for rough character-based estimation.
+"""
+
+from .anthropic import AnthropicTokenCounter
+from .base import TokenCounter, extract_text_from_messages
+from .openai import OpenAITokenCounter
+from .sentencepiece_counter import SentencePieceTokenCounter
+from .utils import (
+    count_post_summary_tokens,
+    count_tokens_from_message_parts,
+    count_tokens_from_messages,
+    get_token_counter,
+)
+
+__all__ = [
+    # Base classes
+    "TokenCounter",
+    # Counter implementations
+    "OpenAITokenCounter",
+    "AnthropicTokenCounter",
+    "SentencePieceTokenCounter",
+    # Utility functions
+    "get_token_counter",
+    "count_tokens_from_messages",
+    "count_post_summary_tokens",
+    "count_tokens_from_message_parts",
+    "extract_text_from_messages",
+]

shotgun/agents/history/token_counting/anthropic.py

@@ -0,0 +1,89 @@
+"""Anthropic token counting using official client."""
+
+from pydantic_ai.messages import ModelMessage
+
+from shotgun.agents.config.models import KeyProvider
+from shotgun.llm_proxy import create_anthropic_proxy_client
+from shotgun.logging_config import get_logger
+
+from .base import TokenCounter, extract_text_from_messages
+
+logger = get_logger(__name__)
+
+
+class AnthropicTokenCounter(TokenCounter):
+    """Token counter for Anthropic models using official client."""
+
+    def __init__(
+        self,
+        model_name: str,
+        api_key: str,
+        key_provider: KeyProvider = KeyProvider.BYOK,
+    ):
+        """Initialize Anthropic token counter.
+
+        Args:
+            model_name: Anthropic model name for token counting
+            api_key: API key (Anthropic for BYOK, Shotgun for proxy)
+            key_provider: Key provider type (BYOK or SHOTGUN)
+
+        Raises:
+            RuntimeError: If client initialization fails
+        """
+        self.model_name = model_name
+        import anthropic
+
+        try:
+            if key_provider == KeyProvider.SHOTGUN:
+                # Use LiteLLM proxy for Shotgun Account
+                # Proxies to Anthropic's token counting API
+                self.client = create_anthropic_proxy_client(api_key)
+                logger.debug(
+                    f"Initialized Anthropic token counter for {model_name} via LiteLLM proxy"
+                )
+            else:
+                # Direct Anthropic API for BYOK
+                self.client = anthropic.Anthropic(api_key=api_key)
+                logger.debug(
+                    f"Initialized Anthropic token counter for {model_name} via direct API"
+                )
+        except Exception as e:
+            raise RuntimeError("Failed to initialize Anthropic client") from e
+
+    async def count_tokens(self, text: str) -> int:
+        """Count tokens using Anthropic's official API (async).
+
+        Args:
+            text: Text to count tokens for
+
+        Returns:
+            Exact token count from Anthropic API
+
+        Raises:
+            RuntimeError: If API call fails
+        """
+        try:
+            # Anthropic API expects messages format and model parameter
+            result = self.client.messages.count_tokens(
+                messages=[{"role": "user", "content": text}], model=self.model_name
+            )
+            return result.input_tokens
+        except Exception as e:
+            raise RuntimeError(
+                f"Anthropic token counting API failed for {self.model_name}"
+            ) from e
+
+    async def count_message_tokens(self, messages: list[ModelMessage]) -> int:
+        """Count tokens across all messages using Anthropic API (async).
+
+        Args:
+            messages: List of PydanticAI messages
+
+        Returns:
+            Total token count for all messages
+
+        Raises:
+            RuntimeError: If token counting fails
+        """
+        total_text = extract_text_from_messages(messages)
+        return await self.count_tokens(total_text)

shotgun/agents/history/token_counting/base.py

@@ -0,0 +1,67 @@
+"""Base classes and shared utilities for token counting."""
+
+from abc import ABC, abstractmethod
+
+from pydantic_ai.messages import ModelMessage
+
+
+class TokenCounter(ABC):
+    """Abstract base class for provider-specific token counting.
+
+    All methods are async to support non-blocking operations like
+    downloading tokenizer models or making API calls.
+    """
+
+    @abstractmethod
+    async def count_tokens(self, text: str) -> int:
+        """Count tokens in text using provider-specific method (async).
+
+        Args:
+            text: Text to count tokens for
+
+        Returns:
+            Exact token count as determined by the provider
+
+        Raises:
+            RuntimeError: If token counting fails
+        """
+
+    @abstractmethod
+    async def count_message_tokens(self, messages: list[ModelMessage]) -> int:
+        """Count tokens in PydanticAI message structures (async).
+
+        Args:
+            messages: List of messages to count tokens for
+
+        Returns:
+            Total token count across all messages
+
+        Raises:
+            RuntimeError: If token counting fails
+        """
+
+
+def extract_text_from_messages(messages: list[ModelMessage]) -> str:
+    """Extract all text content from messages for token counting.
+
+    Args:
+        messages: List of PydanticAI messages
+
+    Returns:
+        Combined text content from all messages
+    """
+    text_parts = []
+
+    for message in messages:
+        if hasattr(message, "parts"):
+            for part in message.parts:
+                if hasattr(part, "content") and isinstance(part.content, str):
+                    text_parts.append(part.content)
+                else:
+                    # Handle non-text parts (tool calls, etc.)
+                    text_parts.append(str(part))
+        else:
+            # Handle messages without parts
+            text_parts.append(str(message))
+
+    return "\n".join(text_parts)