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

shotgun/agents/history/token_counting/utils.py

@@ -0,0 +1,147 @@
+"""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:
+        logger.debug(
+            f"Reusing cached token counter for {model_config.provider.value}:{model_config.name}"
+        )
+        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,7 +4,7 @@ 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
 
@@ -116,7 +116,7 @@ class AgentRuntimeOptions(BaseModel):
     )
 
 
-class FileOperationType(str, Enum):
+class FileOperationType(StrEnum):
     """Types of file operations that can be tracked."""
 
     CREATED = "created"

shotgun/agents/tools/web_search/__init__.py

@@ -1,13 +1,14 @@
 """Web search tools for Pydantic AI agents.
 
 Provides web search capabilities for multiple LLM providers:
-- OpenAI: Uses Responses API with web_search tool
-- Anthropic: Uses Messages API with web_search_20250305 tool
-- Gemini: Uses grounding with Google Search
+- OpenAI: Uses Responses API with web_search tool (BYOK only)
+- Anthropic: Uses Messages API with web_search_20250305 tool (BYOK only)
+- Gemini: Uses grounding with Google Search via Pydantic AI (works with Shotgun Account)
 """
 
-from collections.abc import Callable
+from collections.abc import Awaitable, Callable
 
+from shotgun.agents.config import get_config_manager
 from shotgun.agents.config.models import ProviderType
 from shotgun.logging_config import get_logger
 
@@ -18,29 +19,55 @@ from .utils import is_provider_available
 
 logger = get_logger(__name__)
 
-# Type alias for web search tools
-WebSearchTool = Callable[[str], str]
+# Type alias for web search tools (all now async)
+WebSearchTool = Callable[[str], Awaitable[str]]
 
 
 def get_available_web_search_tools() -> list[WebSearchTool]:
     """Get list of available web search tools based on configured API keys.
 
+    When using Shotgun Account (via LiteLLM proxy):
+        Only Gemini web search is available (others use provider-specific APIs)
+
+    When using BYOK (individual provider keys):
+        All provider tools are available based on their respective keys
+
     Returns:
         List of web search tool functions that have API keys configured
     """
     tools: list[WebSearchTool] = []
 
-    if is_provider_available(ProviderType.OPENAI):
-        logger.debug("✅ OpenAI web search tool available")
-        tools.append(openai_web_search_tool)
+    # Check if using Shotgun Account
+    config_manager = get_config_manager()
+    config = config_manager.load()
+    has_shotgun_key = config.shotgun.api_key is not None
+
+    if has_shotgun_key:
+        # Shotgun Account mode: Only Gemini supports web search via LiteLLM
+        if is_provider_available(ProviderType.GOOGLE):
+            logger.info("🔑 Shotgun Account detected - using Gemini web search only")
+            logger.debug("   OpenAI and Anthropic web search require direct API keys")
+            tools.append(gemini_web_search_tool)
+        else:
+            logger.warning(
+                "⚠️ Shotgun Account configured but no Gemini key - "
+                "web search unavailable"
+            )
+    else:
+        # BYOK mode: Load all available tools based on individual provider keys
+        logger.debug("🔑 BYOK mode - checking all provider web search tools")
+
+        if is_provider_available(ProviderType.OPENAI):
+            logger.debug("✅ OpenAI web search tool available")
+            tools.append(openai_web_search_tool)
 
-    if is_provider_available(ProviderType.ANTHROPIC):
-        logger.debug("✅ Anthropic web search tool available")
-        tools.append(anthropic_web_search_tool)
+        if is_provider_available(ProviderType.ANTHROPIC):
+            logger.debug("✅ Anthropic web search tool available")
+            tools.append(anthropic_web_search_tool)
 
-    if is_provider_available(ProviderType.GOOGLE):
-        logger.debug("✅ Gemini web search tool available")
-        tools.append(gemini_web_search_tool)
+        if is_provider_available(ProviderType.GOOGLE):
+            logger.debug("✅ Gemini web search tool available")
+            tools.append(gemini_web_search_tool)
 
     if not tools:
         logger.warning("⚠️ No web search tools available - no API keys configured")

shotgun/agents/tools/web_search/anthropic.py

@@ -1,20 +1,24 @@
 """Anthropic web search tool implementation."""
 
-import anthropic
 from opentelemetry import trace
+from pydantic_ai.messages import ModelMessage, ModelRequest, TextPart
+from pydantic_ai.settings import ModelSettings
 
 from shotgun.agents.config import get_provider_model
+from shotgun.agents.config.constants import MEDIUM_TEXT_8K_TOKENS
 from shotgun.agents.config.models import ProviderType
+from shotgun.agents.llm import shotgun_model_request
 from shotgun.logging_config import get_logger
 
 logger = get_logger(__name__)
 
 
-def anthropic_web_search_tool(query: str) -> str:
-    """Perform a web search using Anthropic's Claude API with streaming.
+async def anthropic_web_search_tool(query: str) -> str:
+    """Perform a web search using Anthropic's Claude API.
 
     This tool uses Anthropic's web search capabilities to find current information
-    about the given query. Results are streamed for faster response times.
+    about the given query. Works with both Shotgun API keys (via LiteLLM proxy)
+    and direct Anthropic API keys (BYOK).
 
     Args:
         query: The search query
@@ -27,49 +31,49 @@ def anthropic_web_search_tool(query: str) -> str:
     span = trace.get_current_span()
     span.set_attribute("input.value", f"**Query:** {query}\n")
 
-    logger.debug("📡 Executing Anthropic web search with streaming prompt: %s", query)
+    logger.debug("📡 Executing Anthropic web search with prompt: %s", query)
 
-    # Get API key from centralized configuration
+    # Get model configuration (supports both Shotgun and BYOK)
     try:
         model_config = get_provider_model(ProviderType.ANTHROPIC)
-        api_key = model_config.api_key
     except ValueError as e:
         error_msg = f"Anthropic API key not configured: {str(e)}"
         logger.error("❌ %s", error_msg)
         span.set_attribute("output.value", f"**Error:**\n {error_msg}\n")
         return error_msg
 
-    client = anthropic.Anthropic(api_key=api_key)
+    # Build the request messages
+    messages: list[ModelMessage] = [
+        ModelRequest.user_text_prompt(f"Search for: {query}")
+    ]
 
-    # Use the Messages API with web search tool and streaming
+    # Use the Messages API with web search tool
     try:
-        result_text = ""
-
-        with client.messages.stream(
-            model="claude-3-5-sonnet-latest",
-            max_tokens=8192,  # Maximum for Claude 3.5 Sonnet
-            messages=[{"role": "user", "content": f"Search for: {query}"}],
-            tools=[
-                {
-                    "type": "web_search_20250305",
-                    "name": "web_search",
-                }
-            ],
-            tool_choice={"type": "tool", "name": "web_search"},
-        ) as stream:
-            logger.debug("🌊 Started streaming Anthropic web search response")
-
-            for event in stream:
-                if event.type == "content_block_delta":
-                    if hasattr(event.delta, "text"):
-                        result_text += event.delta.text
-                elif event.type == "message_start":
-                    logger.debug("🚀 Streaming started")
-                elif event.type == "message_stop":
-                    logger.debug("✅ Streaming completed")
-
-        if not result_text:
-            result_text = "No content returned from search"
+        response = await shotgun_model_request(
+            model_config=model_config,
+            messages=messages,
+            model_settings=ModelSettings(
+                max_tokens=MEDIUM_TEXT_8K_TOKENS,
+                # Enable Anthropic web search tool
+                extra_body={
+                    "tools": [
+                        {
+                            "type": "web_search_20250305",
+                            "name": "web_search",
+                        }
+                    ],
+                    "tool_choice": {"type": "tool", "name": "web_search"},
+                },
+            ),
+        )
+
+        # Extract text from response
+        result_text = "No content returned from search"
+        if response.parts:
+            for part in response.parts:
+                if isinstance(part, TextPart):
+                    result_text = part.content
+                    break
 
         logger.debug("📄 Anthropic web search result: %d characters", len(result_text))
         logger.debug(
@@ -88,9 +92,8 @@ def anthropic_web_search_tool(query: str) -> str:
         return error_msg
 
 
-def main() -> None:
+async def main() -> None:
     """Main function for testing the Anthropic web search tool."""
-    import os
     import sys
 
     from shotgun.logging_config import setup_logger
@@ -110,24 +113,23 @@ def main() -> None:
     # Join all arguments as the search query
     query = " ".join(sys.argv[1:])
 
-    print("🔍 Testing Anthropic Web Search with streaming")
+    print("🔍 Testing Anthropic Web Search")
     print(f"📝 Query: {query}")
     print("=" * 60)
 
     # Check if API key is available
-    if not (
-        os.getenv("ANTHROPIC_API_KEY")
-        or (
-            callable(get_provider_model)
-            and get_provider_model(ProviderType.ANTHROPIC).api_key
-        )
-    ):
-        print("❌ Error: ANTHROPIC_API_KEY environment variable not set")
-        print("   Please set it with: export ANTHROPIC_API_KEY=your_key_here")
+    try:
+        if callable(get_provider_model):
+            model_config = get_provider_model(ProviderType.ANTHROPIC)
+            if not model_config.api_key:
+                raise ValueError("No API key configured")
+    except (ValueError, Exception):
+        print("❌ Error: Anthropic API key not configured")
+        print("   Please set it in your config file")
         sys.exit(1)
 
     try:
-        result = anthropic_web_search_tool(query)
+        result = await anthropic_web_search_tool(query)
         print(f"✅ Search completed! Result length: {len(result)} characters")
         print("=" * 60)
         print("📄 RESULTS:")
@@ -141,4 +143,6 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    import asyncio
+
+    asyncio.run(main())

shotgun/agents/tools/web_search/gemini.py

@@ -1,20 +1,24 @@
 """Gemini web search tool implementation."""
 
-import google.generativeai as genai
 from opentelemetry import trace
+from pydantic_ai.messages import ModelMessage, ModelRequest
+from pydantic_ai.settings import ModelSettings
 
 from shotgun.agents.config import get_provider_model
-from shotgun.agents.config.models import ProviderType
+from shotgun.agents.config.constants import MEDIUM_TEXT_8K_TOKENS
+from shotgun.agents.config.models import ModelName
+from shotgun.agents.llm import shotgun_model_request
 from shotgun.logging_config import get_logger
 
 logger = get_logger(__name__)
 
 
-def gemini_web_search_tool(query: str) -> str:
+async def gemini_web_search_tool(query: str) -> str:
     """Perform a web search using Google's Gemini API with grounding.
 
     This tool uses Gemini's Google Search grounding to find current information
-    about the given query.
+    about the given query. Works with both Shotgun API keys (via LiteLLM proxy)
+    and direct Gemini API keys (BYOK).
 
     Args:
         query: The search query
@@ -29,23 +33,16 @@ def gemini_web_search_tool(query: str) -> str:
 
     logger.debug("📡 Executing Gemini web search with prompt: %s", query)
 
-    # Get API key from centralized configuration
+    # Get model configuration (supports both Shotgun and BYOK)
     try:
-        model_config = get_provider_model(ProviderType.GOOGLE)
-        api_key = model_config.api_key
+        model_config = get_provider_model(ModelName.GEMINI_2_5_FLASH)
     except ValueError as e:
         error_msg = f"Gemini API key not configured: {str(e)}"
         logger.error("❌ %s", error_msg)
         span.set_attribute("output.value", f"**Error:**\n {error_msg}\n")
         return error_msg
 
-    genai.configure(api_key=api_key)  # type: ignore[attr-defined]
-
-    # Create model without built-in tools to avoid conflict with Pydantic AI
-    # Using prompt-based search approach instead
-    model = genai.GenerativeModel("gemini-2.5-pro")  # type: ignore[attr-defined]
-
-    # Create a search-optimized prompt that leverages Gemini's knowledge
+    # Create a search-optimized prompt
     search_prompt = f"""Please provide current and accurate information about the following query:
 
 Query: {query}
@@ -56,17 +53,31 @@ Instructions:
 - Focus on current and recent information
 - Be specific and accurate in your response"""
 
-    # Generate response using the model's knowledge
+    # Build the request messages
+    messages: list[ModelMessage] = [ModelRequest.user_text_prompt(search_prompt)]
+
+    # Generate response using Pydantic AI with Google Search grounding
     try:
-        response = model.generate_content(
-            search_prompt,
-            generation_config=genai.GenerationConfig(  # type: ignore[attr-defined]
+        response = await shotgun_model_request(
+            model_config=model_config,
+            messages=messages,
+            model_settings=ModelSettings(
                 temperature=0.3,
-                max_output_tokens=8192,
+                max_tokens=MEDIUM_TEXT_8K_TOKENS,
+                # Enable Google Search grounding for Gemini
+                extra_body={"tools": [{"googleSearch": {}}]},
             ),
         )
 
-        result_text = response.text or "No content returned from search"
+        # Extract text from response
+        from pydantic_ai.messages import TextPart
+
+        result_text = "No content returned from search"
+        if response.parts:
+            for part in response.parts:
+                if isinstance(part, TextPart):
+                    result_text = part.content
+                    break
 
         logger.debug("📄 Gemini web search result: %d characters", len(result_text))
         logger.debug(

shotgun/agents/tools/web_search/openai.py

@@ -1,6 +1,6 @@
 """OpenAI web search tool implementation."""
 
-from openai import OpenAI
+from openai import AsyncOpenAI
 from opentelemetry import trace
 
 from shotgun.agents.config import get_provider_model
@@ -10,7 +10,7 @@ from shotgun.logging_config import get_logger
 logger = get_logger(__name__)
 
 
-def openai_web_search_tool(query: str) -> str:
+async def openai_web_search_tool(query: str) -> str:
     """Perform a web search and return results.
 
     This tool uses OpenAI's web search capabilities to find current information
@@ -54,8 +54,8 @@ Instructions:
 ALWAYS PROVIDE THE SOURCES (urls) TO BACK UP THE INFORMATION YOU PROVIDE.
 """
 
-        client = OpenAI(api_key=api_key)
-        response = client.responses.create(  # type: ignore[call-overload]
+        client = AsyncOpenAI(api_key=api_key)
+        response = await client.responses.create(  # type: ignore[call-overload]
             model="gpt-5-mini",
             input=[
                 {"role": "user", "content": [{"type": "input_text", "text": prompt}]}