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

shotgun/agents/tools/web_search/openai.py

@@ -1,16 +1,27 @@
 """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
 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()
 
-def openai_web_search_tool(query: str) -> str:
+
+@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.
 
     This tool uses OpenAI's web search capabilities to find current information
@@ -32,7 +43,7 @@ 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)}"
@@ -40,22 +51,20 @@ 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 = 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}]}

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

@@ -0,0 +1,164 @@
+import json
+from collections import defaultdict
+from dataclasses import dataclass
+from datetime import datetime
+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
+
+from shotgun.agents.config.models import ProviderType
+from shotgun.utils import get_shotgun_home
+
+logger = getLogger(__name__)
+ModelName: TypeAlias = str
+
+
+@dataclass(frozen=True)
+class UsageSummaryEntry:
+    model_name: ModelName
+    provider: ProviderType
+    usage: RunUsage
+
+
+class UsageLogEntry(BaseModel):
+    timestamp: datetime = Field(default_factory=datetime.now)
+    model_name: ModelName
+    usage: RunUsage
+    provider: ProviderType
+
+
+class SessionUsage(BaseModel):
+    usage: RunUsage
+    log: list[UsageLogEntry]
+
+
+class UsageState(BaseModel):
+    usage: dict[ModelName, RunUsage] = Field(default_factory=dict)
+    model_providers: dict[ModelName, ProviderType] = Field(default_factory=dict)
+    usage_log: list[UsageLogEntry] = Field(default_factory=list)
+
+
+class SessionUsageManager:
+    def __init__(self) -> None:
+        self.usage: defaultdict[ModelName, RunUsage] = defaultdict(RunUsage)
+        self._model_providers: dict[ModelName, ProviderType] = {}
+        self._usage_log: list[UsageLogEntry] = []
+        self._usage_path: Path = get_shotgun_home() / "usage.json"
+        # Note: restore_usage_state needs to be called asynchronously after init
+        # Caller should use: manager = SessionUsageManager(); await manager.restore_usage_state()
+
+    async def add_usage(
+        self, usage: RunUsage, *, model_name: ModelName, provider: ProviderType
+    ) -> None:
+        self.usage[model_name] += usage
+        self._model_providers[model_name] = provider
+        self._usage_log.append(
+            UsageLogEntry(model_name=model_name, usage=usage, provider=provider)
+        )
+        await self.persist_usage_state()
+
+    def get_usage_report(self) -> dict[ModelName, RunUsage]:
+        return self.usage.copy()
+
+    def get_usage_breakdown(self) -> list[UsageSummaryEntry]:
+        breakdown: list[UsageSummaryEntry] = []
+        for model_name, usage in self.usage.items():
+            provider = self._model_providers.get(model_name)
+            if provider is None:
+                continue
+            breakdown.append(
+                UsageSummaryEntry(model_name=model_name, provider=provider, usage=usage)
+            )
+        breakdown.sort(key=lambda entry: entry.model_name.lower())
+        return breakdown
+
+    def build_usage_hint(self) -> str | None:
+        return format_usage_hint(self.get_usage_breakdown())
+
+    async def persist_usage_state(self) -> None:
+        state = UsageState(
+            usage=dict(self.usage.items()),
+            model_providers=self._model_providers.copy(),
+            usage_log=self._usage_log.copy(),
+        )
+
+        try:
+            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
+            )
+
+    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:
+            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:
+            logger.error(
+                "Failed to restore usage state from %s: %s", self._usage_path, exc
+            )
+            return
+
+        self.usage = defaultdict(RunUsage)
+        for model_name, usage in state.usage.items():
+            self.usage[model_name] = usage
+
+        self._model_providers = state.model_providers.copy()
+        self._usage_log = state.usage_log.copy()
+
+
+def format_usage_hint(breakdown: list[UsageSummaryEntry]) -> str | None:
+    if not breakdown:
+        return None
+
+    lines = ["# Token usage by model"]
+
+    for entry in breakdown:
+        usage = entry.usage
+        input_tokens = usage.input_tokens
+        output_tokens = usage.output_tokens
+        cached_tokens = usage.cache_read_tokens
+
+        cost = calc_price(usage=usage, model_ref=entry.model_name)
+        input_line = f"* Input: {input_tokens:,}"
+        if cached_tokens > 0:
+            input_line += f" (+ {cached_tokens:,} cached)"
+        input_line += " tokens"
+        section = f"""
+### {entry.model_name}
+
+{input_line}
+* Output: {output_tokens:,} tokens
+* Total: {input_tokens + output_tokens:,} tokens
+* Cost: ${cost.total_price:,.2f}
+""".strip()
+        lines.append(section)
+
+    return "\n\n".join(lines)
+
+
+_usage_manager = None
+
+
+def get_session_usage_manager() -> SessionUsageManager:
+    global _usage_manager
+    if _usage_manager is None:
+        _usage_manager = SessionUsageManager()
+    return _usage_manager

shotgun/api_endpoints.py

@@ -0,0 +1,15 @@
+"""Shotgun backend service API endpoints and URLs."""
+
+from shotgun.settings import settings
+
+# Shotgun Web API base URL (for authentication/subscription)
+# 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)
+# 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"
+LITELLM_PROXY_OPENAI_BASE = LITELLM_PROXY_BASE_URL

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

shotgun/cli/compact.py

@@ -0,0 +1,186 @@
+"""Compact command for shotgun CLI."""
+
+import asyncio
+import json
+from pathlib import Path
+from typing import Annotated, Any
+
+import typer
+from pydantic_ai.usage import RequestUsage
+from rich.console import Console
+
+from shotgun.agents.config import get_provider_model
+from shotgun.agents.conversation_manager import ConversationManager
+from shotgun.agents.history.history_processors import token_limit_compactor
+from shotgun.agents.history.token_estimation import estimate_tokens_from_messages
+from shotgun.cli.models import OutputFormat
+from shotgun.logging_config import get_logger
+
+app = typer.Typer(
+    name="compact", help="Compact the conversation history", no_args_is_help=False
+)
+logger = get_logger(__name__)
+console = Console()
+
+
+@app.callback(invoke_without_command=True)
+def compact(
+    format: Annotated[
+        OutputFormat,
+        typer.Option(
+            "--format",
+            "-f",
+            help="Output format: markdown or json",
+        ),
+    ] = OutputFormat.MARKDOWN,
+) -> None:
+    """Compact the current conversation history to reduce size.
+
+    This command compacts the conversation in ~/.shotgun-sh/conversation.json
+    by summarizing older messages while preserving recent context. The compacted
+    conversation is automatically saved back to the file.
+    """
+    try:
+        result = asyncio.run(compact_conversation())
+
+        if format == OutputFormat.JSON:
+            # Output as JSON
+            console.print_json(json.dumps(result, indent=2))
+        else:
+            # Output as markdown
+            console.print(format_markdown(result))
+
+    except FileNotFoundError as e:
+        console.print(
+            f"[red]Error:[/red] {e}\n\n"
+            "No conversation found. Start a TUI session first with: [cyan]shotgun[/cyan]",
+            style="bold",
+        )
+        raise typer.Exit(code=1) from e
+    except Exception as e:
+        console.print(
+            f"[red]Error:[/red] Failed to compact conversation: {e}", style="bold"
+        )
+        logger.debug("Full traceback:", exc_info=True)
+        raise typer.Exit(code=1) from e
+
+
+async def compact_conversation() -> dict[str, Any]:
+    """Compact the conversation and return statistics.
+
+    Returns:
+        Dictionary with compaction statistics including before/after metrics
+    """
+    # Get conversation file path
+    conversation_file = Path.home() / ".shotgun-sh" / "conversation.json"
+
+    if not conversation_file.exists():
+        raise FileNotFoundError(f"Conversation file not found at {conversation_file}")
+
+    # Load conversation
+    manager = ConversationManager(conversation_file)
+    conversation = await manager.load()
+
+    if not conversation:
+        raise ValueError("Conversation file is empty or corrupted")
+
+    # Get agent messages only (not UI messages)
+    agent_messages = conversation.get_agent_messages()
+
+    if not agent_messages:
+        raise ValueError("No agent messages found in conversation")
+
+    # Get model config
+    model_config = await get_provider_model()
+
+    # Calculate before metrics
+    original_message_count = len(agent_messages)
+    original_tokens = await estimate_tokens_from_messages(agent_messages, model_config)
+
+    # For CLI, we can call token_limit_compactor directly without full AgentDeps
+    # since we only need the model config and message history
+    # Create a minimal context object for compaction
+    class CompactContext:
+        def __init__(self, model_config: Any, usage: RequestUsage) -> None:
+            self.deps = type("Deps", (), {"llm_model": model_config})()
+            self.usage = usage
+
+    # Create minimal usage info for compaction check
+    usage = RequestUsage(input_tokens=original_tokens, output_tokens=0)
+    ctx = CompactContext(model_config, usage)
+
+    # Apply compaction with force=True to bypass threshold checks
+    compacted_messages = await token_limit_compactor(ctx, agent_messages, force=True)
+
+    # Calculate after metrics
+    compacted_message_count = len(compacted_messages)
+    compacted_tokens = await estimate_tokens_from_messages(
+        compacted_messages, model_config
+    )
+
+    # Calculate reduction percentages
+    message_reduction = (
+        ((original_message_count - compacted_message_count) / original_message_count)
+        * 100
+        if original_message_count > 0
+        else 0
+    )
+    token_reduction = (
+        ((original_tokens - compacted_tokens) / original_tokens) * 100
+        if original_tokens > 0
+        else 0
+    )
+
+    # Save compacted conversation
+    conversation.set_agent_messages(compacted_messages)
+    await manager.save(conversation)
+
+    logger.info(
+        f"Compacted conversation: {original_message_count} → {compacted_message_count} messages "
+        f"({message_reduction:.1f}% reduction)"
+    )
+
+    return {
+        "success": True,
+        "before": {
+            "messages": original_message_count,
+            "estimated_tokens": original_tokens,
+        },
+        "after": {
+            "messages": compacted_message_count,
+            "estimated_tokens": compacted_tokens,
+        },
+        "reduction": {
+            "messages_percent": round(message_reduction, 1),
+            "tokens_percent": round(token_reduction, 1),
+        },
+    }
+
+
+def format_markdown(result: dict[str, Any]) -> str:
+    """Format compaction result as markdown.
+
+    Args:
+        result: Dictionary with compaction statistics
+
+    Returns:
+        Formatted markdown string
+    """
+    before = result["before"]
+    after = result["after"]
+    reduction = result["reduction"]
+
+    return f"""# Conversation Compacted ✓
+
+## Before
+- **Messages:** {before["messages"]:,}
+- **Estimated Tokens:** {before["estimated_tokens"]:,}
+
+## After
+- **Messages:** {after["messages"]:,}
+- **Estimated Tokens:** {after["estimated_tokens"]:,}
+
+## Reduction
+- **Messages:** {reduction["messages_percent"]}%
+- **Tokens:** {reduction["tokens_percent"]}%
+"""