shotgun-sh 0.2.3.dev2__py3-none-any.whl → 0.2.11.dev1__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

@@ -3,7 +3,10 @@
 Provides web search capabilities for multiple LLM providers:
 - 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)
+- 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 Awaitable, Callable
@@ -26,11 +29,12 @@ 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)
+    Works with both Shotgun Account (via LiteLLM proxy) and BYOK (individual provider keys).
 
-    When using BYOK (individual provider keys):
-        All provider tools are available based on their respective 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
@@ -43,15 +47,23 @@ def get_available_web_search_tools() -> list[WebSearchTool]:
     has_shotgun_key = config.shotgun.api_key is not None
 
     if has_shotgun_key:
-        # Shotgun Account mode: Only Gemini supports web search via LiteLLM
+        logger.debug("🔑 Shotgun Account - only Gemini web search available")
+
+        # Gemini: Only search tool available for Shotgun Account
         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")
+            logger.debug("✅ Gemini web search tool available")
             tools.append(gemini_web_search_tool)
-        else:
-            logger.warning(
-                "⚠️ Shotgun Account configured but no Gemini key - "
-                "web search unavailable"
+
+        # Anthropic: Not available for Shotgun Account (Gemini-only for Shotgun)
+        if 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 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

shotgun/agents/tools/web_search/anthropic.py

@@ -8,11 +8,22 @@ 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()
 
+
+@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.
 
@@ -42,10 +53,20 @@ async def anthropic_web_search_tool(query: str) -> str:
         span.set_attribute("output.value", f"**Error:**\n {error_msg}\n")
         return error_msg
 
+    # 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,
+    )
+
     # Build the request messages
-    messages: list[ModelMessage] = [
-        ModelRequest.user_text_prompt(f"Search for: {query}")
-    ]
+    messages: list[ModelMessage] = [ModelRequest.user_text_prompt(search_prompt)]
 
     # Use the Messages API with web search tool
     try:

shotgun/agents/tools/web_search/gemini.py

@@ -8,11 +8,22 @@ from shotgun.agents.config import get_provider_model
 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()
 
+
+@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.
 
@@ -42,16 +53,17 @@ async def gemini_web_search_tool(query: str) -> str:
         span.set_attribute("output.value", f"**Error:**\n {error_msg}\n")
         return error_msg
 
-    # Create a search-optimized prompt
-    search_prompt = f"""Please provide current and accurate information about the following query:
-
-Query: {query}
-
-Instructions:
-- Provide comprehensive, factual information
-- Include relevant details and context
-- Focus on current and recent information
-- Be specific and accurate in your response"""
+    # 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,
+    )
 
     # Build the request messages
     messages: list[ModelMessage] = [ModelRequest.user_text_prompt(search_prompt)]

shotgun/agents/tools/web_search/openai.py

@@ -5,11 +5,22 @@ from opentelemetry import trace
 
 from shotgun.agents.config import get_provider_model
 from shotgun.agents.config.models import ProviderType
+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()
 
+
+@register_tool(
+    category=ToolCategory.WEB_RESEARCH,
+    display_text="Searching web",
+    key_arg="query",
+)
 async def openai_web_search_tool(query: str) -> str:
     """Perform a web search and return results.
 
@@ -40,19 +51,17 @@ async def openai_web_search_tool(query: str) -> str:
             span.set_attribute("output.value", f"**Error:**\n {error_msg}\n")
             return error_msg
 
-        prompt = f"""Please provide current and accurate information about the following query:
-
-Query: {query}
+        # Get datetime context for the search prompt
+        dt_context = get_datetime_context()
 
-Instructions:
-- Provide comprehensive, factual information
-- Include relevant details and context
-- Focus on current and recent information
-- Be specific and accurate in your response
-- You can't ask the user for details, so assume the most relevant details for the query
-
-ALWAYS PROVIDE THE SOURCES (urls) TO BACK UP THE INFORMATION YOU PROVIDE.
-"""
+        # Render search prompt from template
+        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,
+        )
 
         client = AsyncOpenAI(api_key=api_key)
         response = await client.responses.create(  # type: ignore[call-overload]

shotgun/api_endpoints.py

@@ -1,10 +1,14 @@
 """Shotgun backend service API endpoints and URLs."""
 
+from shotgun.settings import settings
+
 # Shotgun Web API base URL (for authentication/subscription)
-# Can be overridden with environment variable
-SHOTGUN_WEB_BASE_URL = "https://api-219702594231.us-east4.run.app"
+# Can be overridden with SHOTGUN_WEB_BASE_URL environment variable
+SHOTGUN_WEB_BASE_URL = settings.api.web_base_url
+
 # Shotgun's LiteLLM proxy base URL (for AI model requests)
-LITELLM_PROXY_BASE_URL = "https://litellm-219702594231.us-east4.run.app"
+# Can be overridden with SHOTGUN_ACCOUNT_LLM_BASE_URL environment variable
+LITELLM_PROXY_BASE_URL = settings.api.account_llm_base_url
 
 # Provider-specific LiteLLM proxy endpoints
 LITELLM_PROXY_ANTHROPIC_BASE = f"{LITELLM_PROXY_BASE_URL}/anthropic"

shotgun/build_constants.py

@@ -13,7 +13,7 @@ POSTHOG_PROJECT_ID = '191396'
 
 # Logfire configuration embedded at build time (only for dev builds)
 LOGFIRE_ENABLED = 'true'
-LOGFIRE_TOKEN = 'pylf_v1_us_KZ5NM1pP3NwgJkbBJt6Ftdzk8mMhmrXcGJHQQgDJ1LfK'
+LOGFIRE_TOKEN = 'pylf_v1_us_RwZMlJm1tX6j0PL5RWWbmZpzK2hLBNtFWStNKlySfjh8'
 
 # Build metadata
 BUILD_TIME_ENV = "production" if SENTRY_DSN else "development"

shotgun/cli/clear.py

@@ -0,0 +1,52 @@
+"""Clear command for shotgun CLI."""
+
+from pathlib import Path
+
+import typer
+from rich.console import Console
+
+from shotgun.agents.conversation_manager import ConversationManager
+from shotgun.logging_config import get_logger
+
+app = typer.Typer(
+    name="clear", help="Clear the conversation history", no_args_is_help=False
+)
+logger = get_logger(__name__)
+console = Console()
+
+
+@app.callback(invoke_without_command=True)
+def clear() -> None:
+    """Clear the current conversation history.
+
+    This command deletes the conversation file at ~/.shotgun-sh/conversation.json,
+    removing all conversation history. Other files in ~/.shotgun-sh/ (config, usage,
+    codebases, logs) are preserved.
+    """
+    try:
+        # Get conversation file path
+        conversation_file = Path.home() / ".shotgun-sh" / "conversation.json"
+
+        # Check if file exists
+        if not conversation_file.exists():
+            console.print(
+                "[yellow]No conversation file found.[/yellow] Nothing to clear.",
+                style="bold",
+            )
+            return
+
+        # Clear the conversation
+        manager = ConversationManager(conversation_file)
+        manager.clear()
+
+        console.print(
+            "[green]✓[/green] Conversation cleared successfully", style="bold"
+        )
+        logger.info("Conversation cleared successfully")
+
+    except Exception as e:
+        console.print(
+            f"[red]Error:[/red] Failed to clear conversation: {e}", style="bold"
+        )
+        logger.debug("Full traceback:", exc_info=True)
+        raise typer.Exit(code=1) from e