shotgun-sh 0.2.6.dev1__py3-none-any.whl → 0.2.17__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

@@ -26,7 +26,7 @@ logger = get_logger(__name__)
 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).
@@ -43,25 +43,25 @@ def get_available_web_search_tools() -> list[WebSearchTool]:
 
     # Check if using Shotgun Account
     config_manager = get_config_manager()
-    config = config_manager.load()
+    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 is_provider_available(ProviderType.GOOGLE):
+        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 is_provider_available(ProviderType.ANTHROPIC):
+        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 is_provider_available(ProviderType.OPENAI):
+        if await is_provider_available(ProviderType.OPENAI):
             logger.debug(
                 "⚠️  OpenAI web search requires BYOK (Responses API not supported via proxy)"
             )
@@ -69,15 +69,15 @@ def get_available_web_search_tools() -> list[WebSearchTool]:
         # 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):
+        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):
+        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):
+        if await is_provider_available(ProviderType.GOOGLE):
             logger.debug("✅ Gemini web search tool available")
             tools.append(gemini_web_search_tool)
 

shotgun/agents/tools/web_search/anthropic.py

@@ -8,6 +8,7 @@ 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
@@ -18,6 +19,11 @@ logger = get_logger(__name__)
 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.
 
@@ -40,7 +46,7 @@ async def anthropic_web_search_tool(query: str) -> str:
 
     # Get model configuration (supports both Shotgun and BYOK)
     try:
-        model_config = get_provider_model(ProviderType.ANTHROPIC)
+        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)
@@ -135,7 +141,7 @@ async def main() -> None:
     # Check if API key is available
     try:
         if callable(get_provider_model):
-            model_config = get_provider_model(ProviderType.ANTHROPIC)
+            model_config = await get_provider_model(ProviderType.ANTHROPIC)
             if not model_config.api_key:
                 raise ValueError("No API key configured")
     except (ValueError, Exception):

shotgun/agents/tools/web_search/gemini.py

@@ -8,6 +8,7 @@ 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
@@ -18,6 +19,11 @@ logger = get_logger(__name__)
 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.
 
@@ -40,7 +46,7 @@ async def gemini_web_search_tool(query: str) -> str:
 
     # Get model configuration (supports both Shotgun and BYOK)
     try:
-        model_config = get_provider_model(ModelName.GEMINI_2_5_FLASH)
+        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)

shotgun/agents/tools/web_search/openai.py

@@ -5,6 +5,7 @@ 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
@@ -15,6 +16,11 @@ logger = get_logger(__name__)
 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.
 
@@ -37,7 +43,7 @@ async def openai_web_search_tool(query: str) -> str:
 
         # Get API key from centralized configuration
         try:
-            model_config = get_provider_model(ProviderType.OPENAI)
+            model_config = await get_provider_model(ProviderType.OPENAI)
             api_key = model_config.api_key
         except ValueError as e:
             error_msg = f"OpenAI API key not configured: {str(e)}"

shotgun/agents/tools/web_search/utils.py

@@ -4,7 +4,7 @@ from shotgun.agents.config import get_provider_model
 from shotgun.agents.config.models import ProviderType
 
 
-def is_provider_available(provider: ProviderType) -> bool:
+async def is_provider_available(provider: ProviderType) -> bool:
     """Check if a provider has API key configured.
 
     Args:
@@ -14,7 +14,7 @@ def is_provider_available(provider: ProviderType) -> bool:
         True if the provider has valid credentials configured (from config or env)
     """
     try:
-        get_provider_model(provider)
+        await get_provider_model(provider)
         return True
     except ValueError:
         return False

shotgun/agents/usage_manager.py

@@ -6,6 +6,8 @@ from logging import getLogger
 from pathlib import Path
 from typing import TypeAlias
 
+import aiofiles
+import aiofiles.os
 from genai_prices import calc_price
 from pydantic import BaseModel, Field
 from pydantic_ai import RunUsage
@@ -48,9 +50,10 @@ class SessionUsageManager:
         self._model_providers: dict[ModelName, ProviderType] = {}
         self._usage_log: list[UsageLogEntry] = []
         self._usage_path: Path = get_shotgun_home() / "usage.json"
-        self.restore_usage_state()
+        # Note: restore_usage_state needs to be called asynchronously after init
+        # Caller should use: manager = SessionUsageManager(); await manager.restore_usage_state()
 
-    def add_usage(
+    async def add_usage(
         self, usage: RunUsage, *, model_name: ModelName, provider: ProviderType
     ) -> None:
         self.usage[model_name] += usage
@@ -58,7 +61,7 @@ class SessionUsageManager:
         self._usage_log.append(
             UsageLogEntry(model_name=model_name, usage=usage, provider=provider)
         )
-        self.persist_usage_state()
+        await self.persist_usage_state()
 
     def get_usage_report(self) -> dict[ModelName, RunUsage]:
         return self.usage.copy()
@@ -78,7 +81,7 @@ class SessionUsageManager:
     def build_usage_hint(self) -> str | None:
         return format_usage_hint(self.get_usage_breakdown())
 
-    def persist_usage_state(self) -> None:
+    async def persist_usage_state(self) -> None:
         state = UsageState(
             usage=dict(self.usage.items()),
             model_providers=self._model_providers.copy(),
@@ -86,23 +89,25 @@ class SessionUsageManager:
         )
 
         try:
-            self._usage_path.parent.mkdir(parents=True, exist_ok=True)
-            with self._usage_path.open("w", encoding="utf-8") as f:
-                json.dump(state.model_dump(mode="json"), f, indent=2)
+            await aiofiles.os.makedirs(self._usage_path.parent, exist_ok=True)
+            json_content = json.dumps(state.model_dump(mode="json"), indent=2)
+            async with aiofiles.open(self._usage_path, "w", encoding="utf-8") as f:
+                await f.write(json_content)
             logger.debug("Usage state persisted to %s", self._usage_path)
         except Exception as exc:
             logger.error(
                 "Failed to persist usage state to %s: %s", self._usage_path, exc
             )
 
-    def restore_usage_state(self) -> None:
-        if not self._usage_path.exists():
+    async def restore_usage_state(self) -> None:
+        if not await aiofiles.os.path.exists(self._usage_path):
             logger.debug("No usage state file found at %s", self._usage_path)
             return
 
         try:
-            with self._usage_path.open(encoding="utf-8") as f:
-                data = json.load(f)
+            async with aiofiles.open(self._usage_path, encoding="utf-8") as f:
+                content = await f.read()
+                data = json.loads(content)
 
             state = UsageState.model_validate(data)
         except Exception as exc:

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

@@ -8,12 +8,12 @@ DO NOT EDIT MANUALLY.
 SENTRY_DSN = 'https://2818a6d165c64eccc94cfd51ce05d6aa@o4506813296738304.ingest.us.sentry.io/4510045952409600'
 
 # PostHog configuration embedded at build time (empty strings if not provided)
-POSTHOG_API_KEY = ''
+POSTHOG_API_KEY = 'phc_KKnChzZUKeNqZDOTJ6soCBWNQSx3vjiULdwTR9H5Mcr'
 POSTHOG_PROJECT_ID = '191396'
 
 # Logfire configuration embedded at build time (only for dev builds)
-LOGFIRE_ENABLED = 'true'
-LOGFIRE_TOKEN = 'pylf_v1_us_KZ5NM1pP3NwgJkbBJt6Ftdzk8mMhmrXcGJHQQgDJ1LfK'
+LOGFIRE_ENABLED = ''
+LOGFIRE_TOKEN = ''
 
 # Build metadata
 BUILD_TIME_ENV = "production" if SENTRY_DSN else "development"

shotgun/cli/clear.py

@@ -0,0 +1,53 @@
+"""Clear command for shotgun CLI."""
+
+import asyncio
+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)
+        asyncio.run(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