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

shotgun/agents/tools/registry.py

@@ -0,0 +1,217 @@
+"""Tool category registry using decorators for automatic registration.
+
+This module provides a decorator-based system for categorizing tools used by agents.
+Tools can be decorated with @register_tool to automatically register their category,
+which is then used by the context analyzer to break down token usage by tool type.
+
+It also provides a display registry system for tool formatting in the TUI, allowing
+tools to declare how they should be displayed when streaming.
+"""
+
+from collections.abc import Callable
+from enum import StrEnum
+from typing import TypeVar, overload
+
+import sentry_sdk
+from pydantic import BaseModel
+
+from shotgun.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+# Type variable for decorated functions
+F = TypeVar("F", bound=Callable[..., object])
+
+
+class ToolCategory(StrEnum):
+    """Categories for agent tools used in context analysis."""
+
+    CODEBASE_UNDERSTANDING = "codebase_understanding"
+    ARTIFACT_MANAGEMENT = "artifact_management"
+    WEB_RESEARCH = "web_research"
+    AGENT_RESPONSE = "agent_response"
+    UNKNOWN = "unknown"
+
+
+class ToolDisplayConfig(BaseModel):
+    """Configuration for how a tool should be displayed in the TUI.
+
+    Attributes:
+        display_text: Text to show (e.g., "Reading file", "Querying code")
+        key_arg: Primary argument to extract from tool args for display
+        hide: Whether to completely hide this tool call from the UI
+    """
+
+    display_text: str
+    key_arg: str
+    hide: bool = False
+
+
+# Global registry mapping tool names to categories
+_TOOL_REGISTRY: dict[str, ToolCategory] = {}
+
+# Global registry mapping tool names to display configs
+_TOOL_DISPLAY_REGISTRY: dict[str, ToolDisplayConfig] = {}
+
+
+@overload
+def register_tool(
+    category: ToolCategory,
+    display_text: str,
+    key_arg: str,
+) -> Callable[[F], F]: ...
+
+
+@overload
+def register_tool(
+    category: ToolCategory,
+    display_text: str,
+    key_arg: str,
+    *,
+    hide: bool,
+) -> Callable[[F], F]: ...
+
+
+def register_tool(
+    category: ToolCategory,
+    display_text: str,
+    key_arg: str,
+    *,
+    hide: bool = False,
+) -> Callable[[F], F]:
+    """Decorator to register a tool's category and display configuration.
+
+    Args:
+        category: The ToolCategory enum value for this tool
+        display_text: Text to show (e.g., "Reading file", "Querying code")
+        key_arg: Primary argument name to extract for display (e.g., "query", "filename")
+        hide: Whether to hide this tool call completely from the UI (default: False)
+
+    Returns:
+        Decorator function that registers the tool and returns it unchanged
+
+    Display Format:
+        - When key_arg value is missing: Shows just display_text (e.g., "Reading file")
+        - When key_arg value is present: Shows "display_text: key_arg_value" (e.g., "Reading file: foo.py")
+
+    Example:
+        @register_tool(
+            category=ToolCategory.CODEBASE_UNDERSTANDING,
+            display_text="Querying code",
+            key_arg="query",
+        )
+        async def query_graph(ctx: RunContext[AgentDeps], query: str) -> str:
+            ...
+    """
+
+    def decorator(func: F) -> F:
+        tool_name = func.__name__
+        _TOOL_REGISTRY[tool_name] = category
+        logger.debug(f"Registered tool '{tool_name}' as category '{category.value}'")
+
+        # Register display config
+        config = ToolDisplayConfig(
+            display_text=display_text,
+            key_arg=key_arg,
+            hide=hide,
+        )
+        _TOOL_DISPLAY_REGISTRY[tool_name] = config
+        logger.debug(f"Registered display config for tool '{tool_name}'")
+
+        return func
+
+    return decorator
+
+
+# Backwards compatibility alias
+tool_category = register_tool
+
+
+def get_tool_category(tool_name: str) -> ToolCategory:
+    """Get category for a tool, logging unknown tools to Sentry.
+
+    Args:
+        tool_name: Name of the tool to look up
+
+    Returns:
+        ToolCategory enum value for the tool, or UNKNOWN if not registered
+    """
+    category = _TOOL_REGISTRY.get(tool_name)
+
+    if category is None:
+        logger.warning(f"Unknown tool encountered in context analysis: {tool_name}")
+        sentry_sdk.capture_message(
+            f"Unknown tool in context analysis: {tool_name}",
+            level="warning",
+            extras={"tool_name": tool_name},
+        )
+        return ToolCategory.UNKNOWN
+
+    return category
+
+
+def register_special_tool(tool_name: str, category: ToolCategory) -> None:
+    """Register a special tool that doesn't have a decorator.
+
+    Used for tools like 'final_result' that aren't actual Python functions
+    but need to be categorized.
+
+    Args:
+        tool_name: Name of the special tool
+        category: Category to assign to this tool
+    """
+    _TOOL_REGISTRY[tool_name] = category
+    logger.debug(
+        f"Registered special tool '{tool_name}' as category '{category.value}'"
+    )
+
+
+def get_tool_display_config(tool_name: str) -> ToolDisplayConfig | None:
+    """Get display configuration for a tool.
+
+    Args:
+        tool_name: Name of the tool to look up
+
+    Returns:
+        ToolDisplayConfig for the tool, or None if not registered
+    """
+    return _TOOL_DISPLAY_REGISTRY.get(tool_name)
+
+
+def register_tool_display(
+    tool_name: str,
+    display_text: str,
+    key_arg: str,
+    *,
+    hide: bool = False,
+) -> None:
+    """Register a display config for a special tool that doesn't have a decorator.
+
+    Used for tools like 'final_result' or builtin tools that aren't actual Python functions.
+
+    Args:
+        tool_name: Name of the special tool
+        display_text: Text to show (e.g., "Reading file", "Querying code")
+        key_arg: Primary argument name to extract for display
+        hide: Whether to hide this tool call completely
+    """
+    config = ToolDisplayConfig(
+        display_text=display_text,
+        key_arg=key_arg,
+        hide=hide,
+    )
+    _TOOL_DISPLAY_REGISTRY[tool_name] = config
+    logger.debug(f"Registered display config for special tool '{tool_name}'")
+
+
+# Register special tools that don't have decorators
+register_special_tool("final_result", ToolCategory.AGENT_RESPONSE)
+register_tool_display("final_result", display_text="", key_arg="", hide=True)
+
+# Register builtin tools (tools that come from Pydantic AI or model providers)
+# These don't have Python function definitions but need display formatting
+register_tool_display(
+    "web_search",
+    display_text="Searching",
+    key_arg="query",
+)

shotgun/agents/tools/web_search/__init__.py

@@ -1,13 +1,17 @@
 """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 (Shotgun Account and BYOK)
+
+Shotgun Account: Only Gemini web search is available
+BYOK: All tools work with direct provider API keys
 """
 
-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 +22,64 @@ 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]:
+async def get_available_web_search_tools() -> list[WebSearchTool]:
     """Get list of available web search tools based on configured API keys.
 
+    Works with both Shotgun Account (via LiteLLM proxy) and BYOK (individual provider keys).
+
+    Available tools:
+    - Gemini: Available for both Shotgun Account and BYOK
+    - Anthropic: BYOK only (uses Messages API with web search)
+    - OpenAI: BYOK only (uses Responses API not compatible with LiteLLM proxy)
+
     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 = await config_manager.load()
+    has_shotgun_key = config.shotgun.api_key is not None
+
+    if has_shotgun_key:
+        logger.debug("🔑 Shotgun Account - only Gemini web search available")
+
+        # Gemini: Only search tool available for Shotgun Account
+        if await is_provider_available(ProviderType.GOOGLE):
+            logger.debug("✅ Gemini web search tool available")
+            tools.append(gemini_web_search_tool)
+
+        # Anthropic: Not available for Shotgun Account (Gemini-only for Shotgun)
+        if await is_provider_available(ProviderType.ANTHROPIC):
+            logger.debug(
+                "⚠️  Anthropic web search requires BYOK (Shotgun Account uses Gemini only)"
+            )
+
+        # OpenAI: Not available for Shotgun Account (Responses API incompatible with proxy)
+        if await is_provider_available(ProviderType.OPENAI):
+            logger.debug(
+                "⚠️  OpenAI web search requires BYOK (Responses API not supported via proxy)"
+            )
+    else:
+        # BYOK mode: Load all available tools based on individual provider keys
+        logger.debug("🔑 BYOK mode - checking all provider web search tools")
+
+        if await 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 await 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 await 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,35 @@
 """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.agents.tools.registry import ToolCategory, register_tool
 from shotgun.logging_config import get_logger
+from shotgun.prompts import PromptLoader
+from shotgun.utils.datetime_utils import get_datetime_context
 
 logger = get_logger(__name__)
 
+# Global prompt loader instance
+prompt_loader = PromptLoader()
 
-def anthropic_web_search_tool(query: str) -> str:
-    """Perform a web search using Anthropic's Claude API with streaming.
+
+@register_tool(
+    category=ToolCategory.WEB_RESEARCH,
+    display_text="Searching web",
+    key_arg="query",
+)
+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 +42,59 @@ 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
+        model_config = await get_provider_model(ProviderType.ANTHROPIC)
     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)
+    # Get datetime context for the search prompt
+    dt_context = get_datetime_context()
+
+    # Render search prompt from template
+    search_prompt = prompt_loader.render(
+        "tools/web_search.j2",
+        query=query,
+        current_datetime=dt_context.datetime_formatted,
+        timezone_name=dt_context.timezone_name,
+        utc_offset=dt_context.utc_offset,
+    )
 
-    # Use the Messages API with web search tool and streaming
+    # Build the request messages
+    messages: list[ModelMessage] = [ModelRequest.user_text_prompt(search_prompt)]
+
+    # 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 +113,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 +134,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 = await 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 +164,6 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    import asyncio
+
+    asyncio.run(main())

shotgun/agents/tools/web_search/gemini.py

@@ -1,20 +1,35 @@
 """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.agents.tools.registry import ToolCategory, register_tool
 from shotgun.logging_config import get_logger
+from shotgun.prompts import PromptLoader
+from shotgun.utils.datetime_utils import get_datetime_context
 
 logger = get_logger(__name__)
 
+# Global prompt loader instance
+prompt_loader = PromptLoader()
 
-def gemini_web_search_tool(query: str) -> str:
+
+@register_tool(
+    category=ToolCategory.WEB_RESEARCH,
+    display_text="Searching web",
+    key_arg="query",
+)
+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,44 +44,52 @@ 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 = await 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]
+    # Get datetime context for the search prompt
+    dt_context = get_datetime_context()
 
-    # Create a search-optimized prompt that leverages Gemini's knowledge
-    search_prompt = f"""Please provide current and accurate information about the following query:
+    # Render search prompt from template
+    search_prompt = prompt_loader.render(
+        "tools/web_search.j2",
+        query=query,
+        current_datetime=dt_context.datetime_formatted,
+        timezone_name=dt_context.timezone_name,
+        utc_offset=dt_context.utc_offset,
+    )
 
-Query: {query}
+    # Build the request messages
+    messages: list[ModelMessage] = [ModelRequest.user_text_prompt(search_prompt)]
 
-Instructions:
-- Provide comprehensive, factual information
-- Include relevant details and context
-- Focus on current and recent information
-- Be specific and accurate in your response"""
-
-    # Generate response using the model's knowledge
+    # 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(