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

shotgun/agents/history/token_counting/openai.py

@@ -0,0 +1,90 @@
+"""OpenAI token counting using tiktoken."""
+
+from pydantic_ai.messages import ModelMessage
+
+from shotgun.logging_config import get_logger
+
+from .base import TokenCounter, extract_text_from_messages
+
+logger = get_logger(__name__)
+
+
+class OpenAITokenCounter(TokenCounter):
+    """Token counter for OpenAI models using tiktoken."""
+
+    # Official encoding mappings for OpenAI models
+    ENCODING_MAP = {
+        "gpt-5": "o200k_base",
+        "gpt-4o": "o200k_base",
+        "gpt-4": "cl100k_base",
+        "gpt-3.5-turbo": "cl100k_base",
+    }
+
+    def __init__(self, model_name: str):
+        """Initialize OpenAI token counter.
+
+        Args:
+            model_name: OpenAI model name to get correct encoding for
+
+        Raises:
+            RuntimeError: If encoding initialization fails
+        """
+        self.model_name = model_name
+
+        import tiktoken
+
+        try:
+            # Get the appropriate encoding for this model
+            encoding_name = self.ENCODING_MAP.get(model_name, "o200k_base")
+            self.encoding = tiktoken.get_encoding(encoding_name)
+            logger.debug(
+                f"Initialized OpenAI token counter with {encoding_name} encoding"
+            )
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to initialize tiktoken encoding for {model_name}"
+            ) from e
+
+    async def count_tokens(self, text: str) -> int:
+        """Count tokens using tiktoken (async).
+
+        Args:
+            text: Text to count tokens for
+
+        Returns:
+            Exact token count using tiktoken
+
+        Raises:
+            RuntimeError: If token counting fails
+        """
+        # Handle empty text to avoid unnecessary encoding
+        if not text or not text.strip():
+            return 0
+
+        try:
+            return len(self.encoding.encode(text))
+        except BaseException as e:
+            # Must catch BaseException to handle PanicException from tiktoken's Rust layer
+            # which can occur with extremely long texts. Regular Exception won't catch it.
+            raise RuntimeError(
+                f"Failed to count tokens for OpenAI model {self.model_name}"
+            ) from e
+
+    async def count_message_tokens(self, messages: list[ModelMessage]) -> int:
+        """Count tokens across all messages using tiktoken (async).
+
+        Args:
+            messages: List of PydanticAI messages
+
+        Returns:
+            Total token count for all messages
+
+        Raises:
+            RuntimeError: If token counting fails
+        """
+        # Handle empty message list early
+        if not messages:
+            return 0
+
+        total_text = extract_text_from_messages(messages)
+        return await self.count_tokens(total_text)

shotgun/agents/history/token_counting/sentencepiece_counter.py

@@ -0,0 +1,127 @@
+"""Gemini token counting using official SentencePiece tokenizer.
+
+This implementation uses Google's official Gemini/Gemma tokenizer model
+for 100% accurate local token counting without API calls.
+
+Performance: 10-100x faster than API-based counting.
+Accuracy: 100% match with actual Gemini API usage.
+
+The tokenizer is downloaded on first use and cached locally for future use.
+"""
+
+from typing import Any
+
+from pydantic_ai.messages import ModelMessage
+
+from shotgun.logging_config import get_logger
+
+from .base import TokenCounter, extract_text_from_messages
+from .tokenizer_cache import download_gemini_tokenizer, get_gemini_tokenizer_path
+
+logger = get_logger(__name__)
+
+
+class SentencePieceTokenCounter(TokenCounter):
+    """Token counter for Gemini models using official SentencePiece tokenizer.
+
+    This counter provides 100% accurate token counting for Gemini models
+    using the official tokenizer model from Google's gemma_pytorch repository.
+    Token counting is performed locally without any API calls, resulting in
+    10-100x performance improvement over API-based methods.
+
+    The tokenizer is downloaded asynchronously on first use and cached locally.
+    """
+
+    def __init__(self, model_name: str):
+        """Initialize Gemini SentencePiece token counter.
+
+        The tokenizer is not loaded immediately - it will be downloaded and
+        loaded lazily on first use.
+
+        Args:
+            model_name: Gemini model name (used for logging)
+        """
+        self.model_name = model_name
+        self.sp: Any | None = None  # SentencePieceProcessor, loaded lazily
+
+    async def _ensure_tokenizer(self) -> None:
+        """Ensure tokenizer is downloaded and loaded.
+
+        This method downloads the tokenizer on first call (if not cached)
+        and loads it into memory. Subsequent calls reuse the loaded tokenizer.
+
+        Raises:
+            RuntimeError: If tokenizer download or loading fails
+        """
+        if self.sp is not None:
+            # Already loaded
+            return
+
+        import sentencepiece as spm  # type: ignore[import-untyped]
+
+        try:
+            # Check if already cached, otherwise download
+            tokenizer_path = get_gemini_tokenizer_path()
+            if not tokenizer_path.exists():
+                await download_gemini_tokenizer()
+
+            # Load the tokenizer
+            self.sp = spm.SentencePieceProcessor()
+            self.sp.load(str(tokenizer_path))
+            logger.debug(f"Loaded SentencePiece tokenizer for {self.model_name}")
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to load Gemini tokenizer for {self.model_name}"
+            ) from e
+
+    async def count_tokens(self, text: str) -> int:
+        """Count tokens using SentencePiece (async).
+
+        Downloads tokenizer on first call if not cached.
+
+        Args:
+            text: Text to count tokens for
+
+        Returns:
+            Exact token count using Gemini's tokenizer
+
+        Raises:
+            RuntimeError: If token counting fails
+        """
+        # Handle empty text to avoid unnecessary tokenization
+        if not text or not text.strip():
+            return 0
+
+        await self._ensure_tokenizer()
+
+        if self.sp is None:
+            raise RuntimeError(f"Tokenizer not initialized for {self.model_name}")
+
+        try:
+            tokens = self.sp.encode(text)
+            return len(tokens)
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to count tokens for Gemini model {self.model_name}"
+            ) from e
+
+    async def count_message_tokens(self, messages: list[ModelMessage]) -> int:
+        """Count tokens across all messages using SentencePiece (async).
+
+        Downloads tokenizer on first call if not cached.
+
+        Args:
+            messages: List of PydanticAI messages
+
+        Returns:
+            Total token count for all messages
+
+        Raises:
+            RuntimeError: If token counting fails
+        """
+        # Handle empty message list early
+        if not messages:
+            return 0
+
+        total_text = extract_text_from_messages(messages)
+        return await self.count_tokens(total_text)

shotgun/agents/history/token_counting/tokenizer_cache.py

@@ -0,0 +1,92 @@
+"""Async tokenizer download and caching utilities."""
+
+import hashlib
+from pathlib import Path
+
+import aiofiles
+import httpx
+
+from shotgun.logging_config import get_logger
+from shotgun.utils.file_system_utils import get_shotgun_home
+
+logger = get_logger(__name__)
+
+# Gemini tokenizer constants
+GEMINI_TOKENIZER_URL = "https://raw.githubusercontent.com/google/gemma_pytorch/main/tokenizer/tokenizer.model"
+GEMINI_TOKENIZER_SHA256 = (
+    "61a7b147390c64585d6c3543dd6fc636906c9af3865a5548f27f31aee1d4c8e2"
+)
+
+
+def get_tokenizer_cache_dir() -> Path:
+    """Get the directory for cached tokenizer models.
+
+    Returns:
+        Path to tokenizers cache directory
+    """
+    cache_dir = get_shotgun_home() / "tokenizers"
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
+
+
+def get_gemini_tokenizer_path() -> Path:
+    """Get the path where the Gemini tokenizer should be cached.
+
+    Returns:
+        Path to cached Gemini tokenizer
+    """
+    return get_tokenizer_cache_dir() / "gemini_tokenizer.model"
+
+
+async def download_gemini_tokenizer() -> Path:
+    """Download and cache the official Gemini tokenizer model.
+
+    This downloads Google's official Gemini/Gemma tokenizer from the
+    gemma_pytorch repository and caches it locally for future use.
+
+    The download is async and non-blocking, with SHA256 verification
+    for security.
+
+    Returns:
+        Path to the cached tokenizer file
+
+    Raises:
+        RuntimeError: If download fails or checksum verification fails
+    """
+    cache_path = get_gemini_tokenizer_path()
+
+    # Check if already cached
+    if cache_path.exists():
+        logger.debug(f"Gemini tokenizer already cached at {cache_path}")
+        return cache_path
+
+    logger.info("Downloading Gemini tokenizer (4MB, first time only)...")
+
+    try:
+        # Download with async httpx
+        async with httpx.AsyncClient(timeout=30.0) as client:
+            response = await client.get(GEMINI_TOKENIZER_URL, follow_redirects=True)
+            response.raise_for_status()
+            content = response.content
+
+        # Verify SHA256 checksum
+        actual_hash = hashlib.sha256(content).hexdigest()
+        if actual_hash != GEMINI_TOKENIZER_SHA256:
+            raise RuntimeError(
+                f"Gemini tokenizer checksum mismatch. "
+                f"Expected: {GEMINI_TOKENIZER_SHA256}, got: {actual_hash}"
+            )
+
+        # Atomic write: write to temp file first, then rename
+        temp_path = cache_path.with_suffix(".tmp")
+        async with aiofiles.open(temp_path, "wb") as f:
+            await f.write(content)
+        temp_path.rename(cache_path)
+
+        logger.info(f"Gemini tokenizer downloaded and cached at {cache_path}")
+        return cache_path
+
+    except httpx.HTTPError as e:
+        raise RuntimeError(f"Failed to download Gemini tokenizer: {e}") from e
+    except OSError as e:
+        raise RuntimeError(f"Failed to save Gemini tokenizer: {e}") from e

shotgun/agents/history/token_counting/utils.py

@@ -0,0 +1,144 @@
+"""Utility functions and cache for token counting."""
+
+from pydantic_ai.messages import ModelMessage
+
+from shotgun.agents.config.models import ModelConfig, ProviderType
+from shotgun.logging_config import get_logger
+
+from .anthropic import AnthropicTokenCounter
+from .base import TokenCounter
+from .openai import OpenAITokenCounter
+from .sentencepiece_counter import SentencePieceTokenCounter
+
+logger = get_logger(__name__)
+
+# Global cache for token counter instances (singleton pattern)
+_token_counter_cache: dict[tuple[str, str, str], TokenCounter] = {}
+
+
+def get_token_counter(model_config: ModelConfig) -> TokenCounter:
+    """Get appropriate token counter for the model provider (cached singleton).
+
+    This function ensures that every provider has a proper token counting
+    implementation without any fallbacks to estimation. Token counters are
+    cached to avoid repeated initialization overhead.
+
+    Args:
+        model_config: Model configuration with provider and credentials
+
+    Returns:
+        Cached provider-specific token counter
+
+    Raises:
+        ValueError: If provider is not supported for token counting
+        RuntimeError: If token counter initialization fails
+    """
+    # Create cache key from provider, model name, and API key
+    cache_key = (
+        model_config.provider.value,
+        model_config.name,
+        model_config.api_key[:10]
+        if model_config.api_key
+        else "no-key",  # Partial key for cache
+    )
+
+    # Return cached instance if available
+    if cache_key in _token_counter_cache:
+        return _token_counter_cache[cache_key]
+
+    # Create new instance and cache it
+    logger.debug(
+        f"Creating new token counter for {model_config.provider.value}:{model_config.name}"
+    )
+
+    counter: TokenCounter
+    if model_config.provider == ProviderType.OPENAI:
+        counter = OpenAITokenCounter(model_config.name)
+    elif model_config.provider == ProviderType.ANTHROPIC:
+        counter = AnthropicTokenCounter(
+            model_config.name, model_config.api_key, model_config.key_provider
+        )
+    elif model_config.provider == ProviderType.GOOGLE:
+        # Use local SentencePiece tokenizer (100% accurate, 10-100x faster than API)
+        counter = SentencePieceTokenCounter(model_config.name)
+    else:
+        raise ValueError(
+            f"Unsupported provider for token counting: {model_config.provider}. "
+            f"Supported providers: {[p.value for p in ProviderType]}"
+        )
+
+    # Cache the instance
+    _token_counter_cache[cache_key] = counter
+    logger.debug(
+        f"Cached token counter for {model_config.provider.value}:{model_config.name}"
+    )
+
+    return counter
+
+
+async def count_tokens_from_messages(
+    messages: list[ModelMessage], model_config: ModelConfig
+) -> int:
+    """Count actual tokens from messages using provider-specific methods (async).
+
+    This replaces the old estimation approach with accurate token counting
+    using each provider's official APIs and libraries.
+
+    Args:
+        messages: List of messages to count tokens for
+        model_config: Model configuration with provider info
+
+    Returns:
+        Exact token count for the messages
+
+    Raises:
+        ValueError: If provider is not supported
+        RuntimeError: If token counting fails
+    """
+    counter = get_token_counter(model_config)
+    return await counter.count_message_tokens(messages)
+
+
+async def count_post_summary_tokens(
+    messages: list[ModelMessage], summary_index: int, model_config: ModelConfig
+) -> int:
+    """Count actual tokens from summary onwards for incremental compaction decisions (async).
+
+    Args:
+        messages: Full message history
+        summary_index: Index of the last summary message
+        model_config: Model configuration with provider info
+
+    Returns:
+        Exact token count from summary onwards
+
+    Raises:
+        ValueError: If provider is not supported
+        RuntimeError: If token counting fails
+    """
+    if summary_index >= len(messages):
+        return 0
+
+    post_summary_messages = messages[summary_index:]
+    return await count_tokens_from_messages(post_summary_messages, model_config)
+
+
+async def count_tokens_from_message_parts(
+    messages: list[ModelMessage], model_config: ModelConfig
+) -> int:
+    """Count actual tokens from message parts for summarization requests (async).
+
+    Args:
+        messages: List of messages to count tokens for
+        model_config: Model configuration with provider info
+
+    Returns:
+        Exact token count from message parts
+
+    Raises:
+        ValueError: If provider is not supported
+        RuntimeError: If token counting fails
+    """
+    # For now, use the same logic as count_tokens_from_messages
+    # This can be optimized later if needed for different counting strategies
+    return await count_tokens_from_messages(messages, model_config)

shotgun/agents/history/token_estimation.py

@@ -19,10 +19,10 @@ from .constants import INPUT_BUFFER_TOKENS, MIN_SUMMARY_TOKENS
 from .token_counting import count_tokens_from_messages as _count_tokens_from_messages
 
 
-def estimate_tokens_from_messages(
+async def estimate_tokens_from_messages(
     messages: list[ModelMessage], model_config: ModelConfig
 ) -> int:
-    """Count actual tokens from current message list.
+    """Count actual tokens from current message list (async).
 
     This provides accurate token counting for compaction decisions using
     provider-specific token counting methods instead of rough estimation.
@@ -38,13 +38,13 @@ def estimate_tokens_from_messages(
         ValueError: If provider is not supported
         RuntimeError: If token counting fails
     """
-    return _count_tokens_from_messages(messages, model_config)
+    return await _count_tokens_from_messages(messages, model_config)
 
 
-def estimate_post_summary_tokens(
+async def estimate_post_summary_tokens(
     messages: list[ModelMessage], summary_index: int, model_config: ModelConfig
 ) -> int:
-    """Count actual tokens from summary onwards for incremental compaction decisions.
+    """Count actual tokens from summary onwards for incremental compaction decisions (async).
 
     This treats the summary as a reset point and only counts tokens from the summary
     message onwards. Used to determine if incremental compaction is needed.
@@ -65,13 +65,13 @@ def estimate_post_summary_tokens(
         return 0
 
     post_summary_messages = messages[summary_index:]
-    return estimate_tokens_from_messages(post_summary_messages, model_config)
+    return await estimate_tokens_from_messages(post_summary_messages, model_config)
 
 
-def estimate_tokens_from_message_parts(
+async def estimate_tokens_from_message_parts(
     messages: list[ModelMessage], model_config: ModelConfig
 ) -> int:
-    """Count actual tokens from message parts for summarization requests.
+    """Count actual tokens from message parts for summarization requests (async).
 
     This provides accurate token counting across the codebase using
     provider-specific methods instead of character estimation.
@@ -87,14 +87,14 @@ def estimate_tokens_from_message_parts(
         ValueError: If provider is not supported
         RuntimeError: If token counting fails
     """
-    return _count_tokens_from_messages(messages, model_config)
+    return await _count_tokens_from_messages(messages, model_config)
 
 
-def calculate_max_summarization_tokens(
+async def calculate_max_summarization_tokens(
     ctx_or_model_config: Union["RunContext[AgentDeps]", ModelConfig],
     request_messages: list[ModelMessage],
 ) -> int:
-    """Calculate maximum tokens available for summarization output.
+    """Calculate maximum tokens available for summarization output (async).
 
     This ensures we use the model's full capacity while leaving room for input tokens.
 
@@ -115,7 +115,7 @@ def calculate_max_summarization_tokens(
         return MIN_SUMMARY_TOKENS
 
     # Count actual input tokens using shared utility
-    estimated_input_tokens = estimate_tokens_from_message_parts(
+    estimated_input_tokens = await estimate_tokens_from_message_parts(
         request_messages, model_config
     )
 

shotgun/agents/llm.py

@@ -0,0 +1,62 @@
+"""LLM request utilities for Shotgun agents."""
+
+from typing import Any
+
+from pydantic_ai.direct import model_request
+from pydantic_ai.messages import ModelMessage, ModelResponse
+from pydantic_ai.settings import ModelSettings
+
+from shotgun.agents.config.models import ModelConfig
+from shotgun.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+async def shotgun_model_request(
+    model_config: ModelConfig,
+    messages: list[ModelMessage],
+    model_settings: ModelSettings | 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
+        model_settings: Optional ModelSettings. If None, creates default with max 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)
+
+        # With custom settings
+        response = await shotgun_model_request(model_config, messages, model_settings=ModelSettings(max_tokens=1000, temperature=0.7))
+    """
+    if kwargs.get("max_tokens") is not None:
+        logger.warning(
+            "⚠️ 'max_tokens' argument is ignored in shotgun_model_request. "
+            "Set 'model_settings.max_tokens' instead."
+        )
+
+    if not model_settings:
+        model_settings = ModelSettings()
+
+    if model_settings.get("max_tokens") is None:
+        model_settings["max_tokens"] = model_config.max_output_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/models.py

@@ -4,19 +4,45 @@ import os
 from asyncio import Future, Queue
 from collections.abc import Callable
 from datetime import datetime
-from enum import Enum, StrEnum
+from enum import StrEnum
 from pathlib import Path
 from typing import TYPE_CHECKING
 
 from pydantic import BaseModel, ConfigDict, Field
 from pydantic_ai import RunContext
 
+from shotgun.agents.usage_manager import SessionUsageManager, get_session_usage_manager
+
 from .config.models import ModelConfig
 
 if TYPE_CHECKING:
     from shotgun.codebase.service import CodebaseService
 
 
+class AgentResponse(BaseModel):
+    """Structured response from an agent with optional clarifying questions.
+
+    This model provides a consistent response format for all agents:
+    - response: The main response text (can be empty if only asking questions)
+    - clarifying_questions: Optional list of questions to ask the user
+
+    When clarifying_questions is provided, the agent expects to receive
+    answers before continuing its work. This replaces the ask_questions tool.
+    """
+
+    response: str = Field(
+        description="The agent's response text. Always respond with some text summarizing what happened, whats next, etc.",
+    )
+    clarifying_questions: list[str] | None = Field(
+        default=None,
+        description="""
+Optional list of clarifying questions to ask the user.
+- Single question: Shown as a non-blocking suggestion (user can answer or continue with other prompts)
+- Multiple questions (2+): Asked sequentially in Q&A mode (blocks input until all answered or cancelled)
+""",
+    )
+
+
 class AgentType(StrEnum):
     """Enumeration for available agent types."""
 
@@ -71,6 +97,30 @@ class UserQuestion(BaseModel):
     )
 
 
+class MultipleUserQuestions(BaseModel):
+    """Multiple questions to ask the user sequentially."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    questions: list[str] = Field(
+        description="List of questions to ask the user",
+    )
+    current_index: int = Field(
+        default=0,
+        description="Current question index being asked",
+    )
+    answers: list[str] = Field(
+        default_factory=list,
+        description="Accumulated answers from the user",
+    )
+    tool_call_id: str = Field(
+        description="Tool call id",
+    )
+    result: Future[UserAnswer] = Field(
+        description="Future that will contain all answers formatted as Q&A pairs"
+    )
+
+
 class AgentRuntimeOptions(BaseModel):
     """User interface options for agents."""
 
@@ -98,9 +148,9 @@ class AgentRuntimeOptions(BaseModel):
         description="Maximum number of iterations for agent loops",
     )
 
-    queue: Queue[UserQuestion] = Field(
+    queue: Queue[UserQuestion | MultipleUserQuestions] = Field(
         default_factory=Queue,
-        description="Queue for storing user responses",
+        description="Queue for storing user questions (single or multiple)",
     )
 
     tasks: list[Future[UserAnswer]] = Field(
@@ -108,8 +158,13 @@ class AgentRuntimeOptions(BaseModel):
         description="Tasks for storing deferred tool results",
     )
 
+    usage_manager: SessionUsageManager = Field(
+        default_factory=get_session_usage_manager,
+        description="Usage manager for tracking usage",
+    )
+
 
-class FileOperationType(str, Enum):
+class FileOperationType(StrEnum):
     """Types of file operations that can be tracked."""
 
     CREATED = "created"