shotgun-sh 0.1.15.dev2__tar.gz → 0.1.16.dev2__tar.gz

{shotgun_sh-0.1.15.dev2 → shotgun_sh-0.1.16.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: shotgun-sh
-Version: 0.1.15.dev2
+Version: 0.1.16.dev2
 Summary: AI-powered research, planning, and task management CLI tool
 Project-URL: Homepage, https://shotgun.sh/
 Project-URL: Repository, https://github.com/shotgun-sh/shotgun
@@ -22,6 +22,7 @@ Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Topic :: Utilities
 Requires-Python: >=3.11
 Requires-Dist: anthropic>=0.39.0
+Requires-Dist: genai-prices>=0.0.27
 Requires-Dist: google-generativeai>=0.8.5
 Requires-Dist: httpx>=0.27.0
 Requires-Dist: jinja2>=3.1.0

{shotgun_sh-0.1.15.dev2 → shotgun_sh-0.1.16.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "shotgun-sh"
-version = "0.1.15.dev2"
+version = "0.1.16.dev2"
 description = "AI-powered research, planning, and task management CLI tool"
 readme = "README.md"
 license = { text = "MIT" }
@@ -45,6 +45,7 @@ dependencies = [
     "google-generativeai>=0.8.5",
     "tiktoken>=0.7.0",
     "packaging>=23.0",
+    "genai-prices>=0.0.27",
 ]
 
 [project.urls]

{shotgun_sh-0.1.15.dev2 → shotgun_sh-0.1.16.dev2}/src/shotgun/agents/agent_manager.py

@@ -115,12 +115,12 @@ class AgentManager(Widget):
         super().__init__()
         self.display = False
 
+        if deps is None:
+            raise ValueError("AgentDeps must be provided to AgentManager")
+
         # Use provided deps or create default with interactive mode
         self.deps = deps
 
-        if self.deps is None:
-            raise ValueError("AgentDeps must be provided to AgentManager")
-
         # Create AgentRuntimeOptions from deps for agent creation
         agent_runtime_options = AgentRuntimeOptions(
             interactive_mode=self.deps.interactive_mode,
@@ -269,6 +269,7 @@ class AgentManager(Widget):
         Returns:
             The agent run result.
         """
+        logger.info(f"Running agent {self._current_agent_type.value}")
         # Use merged deps (shared state + agent-specific system prompt) if not provided
         if deps is None:
             deps = self._create_merged_deps(self._current_agent_type)
@@ -395,6 +396,10 @@ class AgentManager(Widget):
         # Apply compaction to persistent message history to prevent cascading growth
         all_messages = result.all_messages()
         self.message_history = await apply_persistent_compaction(all_messages, deps)
+        usage = result.usage()
+        deps.usage_manager.add_usage(
+            usage, model_name=deps.llm_model.name, provider=deps.llm_model.provider
+        )
 
         # Log file operations summary if any files were modified
         file_operations = deps.file_tracker.operations.copy()
@@ -641,6 +646,9 @@ class AgentManager(Widget):
                 filtered_messages.append(msg)
         return filtered_messages
 
+    def get_usage_hint(self) -> str | None:
+        return self.deps.usage_manager.build_usage_hint()
+
     def get_conversation_state(self) -> "ConversationState":
         """Get the current conversation state.
 

{shotgun_sh-0.1.15.dev2 → shotgun_sh-0.1.16.dev2}/src/shotgun/agents/models.py

@@ -11,6 +11,8 @@ from typing import TYPE_CHECKING
 from pydantic import BaseModel, ConfigDict, Field
 from pydantic_ai import RunContext
 
+from shotgun.agents.usage_manager import SessionUsageManager, get_session_usage_manager
+
 from .config.models import ModelConfig
 
 if TYPE_CHECKING:
@@ -108,6 +110,11 @@ class AgentRuntimeOptions(BaseModel):
         description="Tasks for storing deferred tool results",
     )
 
+    usage_manager: SessionUsageManager = Field(
+        default_factory=get_session_usage_manager,
+        description="Usage manager for tracking usage",
+    )
+
 
 class FileOperationType(str, Enum):
     """Types of file operations that can be tracked."""

shotgun_sh-0.1.16.dev2/src/shotgun/agents/usage_manager.py

@@ -0,0 +1,159 @@
+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
+
+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"
+        self.restore_usage_state()
+
+    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)
+        )
+        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())
+
+    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:
+            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)
+            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():
+            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)
+
+            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_sh-0.1.15.dev2 → shotgun_sh-0.1.16.dev2}/src/shotgun/tui/screens/chat.py

@@ -58,6 +58,7 @@ from .chat_screen.command_providers import (
     CodebaseCommandProvider,
     DeleteCodebasePaletteProvider,
     ProviderSetupProvider,
+    UsageProvider,
 )
 
 logger = logging.getLogger(__name__)
@@ -228,9 +229,15 @@ class ChatScreen(Screen[None]):
     BINDINGS = [
         ("ctrl+p", "command_palette", "Command Palette"),
         ("shift+tab", "toggle_mode", "Toggle mode"),
+        ("ctrl+u", "show_usage", "Show usage"),
     ]
 
-    COMMANDS = {AgentModeProvider, ProviderSetupProvider, CodebaseCommandProvider}
+    COMMANDS = {
+        AgentModeProvider,
+        ProviderSetupProvider,
+        CodebaseCommandProvider,
+        UsageProvider,
+    }
 
     value = reactive("")
     mode = reactive(AgentType.RESEARCH)
@@ -401,6 +408,14 @@ class ChatScreen(Screen[None]):
         # whoops it actually changes focus. Let's be brutal for now
         self.call_later(lambda: self.query_one(PromptInput).focus())
 
+    def action_show_usage(self) -> None:
+        usage_hint = self.agent_manager.get_usage_hint()
+        logger.info(f"Usage hint: {usage_hint}")
+        if usage_hint:
+            self.mount_hint(usage_hint)
+        else:
+            self.notify("No usage hint available", severity="error")
+
     @work
     async def add_question_listener(self) -> None:
         while True:
@@ -589,7 +604,7 @@ class ChatScreen(Screen[None]):
     async def index_codebase(self, selection: CodebaseIndexSelection) -> None:
         label = self.query_one("#indexing-job-display", Static)
         label.update(
-            f"[$foreground-muted]Indexing [bold $text-accent]{selection.name}[/]...[/]"
+            f"[$foreground-muted]Indexing codebase: [bold $text-accent]{selection.name}[/][/]"
         )
         label.refresh()
 
@@ -599,33 +614,33 @@ class ChatScreen(Screen[None]):
             empty = width - filled
             return "▓" * filled + "░" * empty
 
-        # Spinner animation state (shared between timer and progress callback)
+        # Spinner animation frames
         spinner_frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
-        spinner_state: dict[str, int | str | float] = {
+
+        # Progress state (shared between timer and progress callback)
+        progress_state: dict[str, int | float] = {
             "frame_index": 0,
-            "phase_name": "Starting...",
             "percentage": 0.0,
         }
 
-        def update_spinner_display() -> None:
-            """Update spinner frame on timer - runs every 100ms."""
+        def update_progress_display() -> None:
+            """Update progress bar on timer - runs every 100ms."""
             # Advance spinner frame
-            frame_idx = int(spinner_state["frame_index"])
-            spinner_state["frame_index"] = (frame_idx + 1) % len(spinner_frames)
+            frame_idx = int(progress_state["frame_index"])
+            progress_state["frame_index"] = (frame_idx + 1) % len(spinner_frames)
             spinner = spinner_frames[frame_idx]
 
             # Get current state
-            phase = str(spinner_state["phase_name"])
-            pct = float(spinner_state["percentage"])
+            pct = float(progress_state["percentage"])
             bar = create_progress_bar(pct)
 
             # Update label
             label.update(
-                f"[$foreground-muted]Indexing codebase: {spinner} {phase}... {bar} {pct:.0f}%[/]"
+                f"[$foreground-muted]Indexing codebase: {spinner} {bar} {pct:.0f}%[/]"
             )
 
         def progress_callback(progress_info: IndexProgress) -> None:
-            """Update progress state (spinner animates independently on timer)."""
+            """Update progress state (timer renders it independently)."""
             # Calculate overall percentage (0-95%, reserve 95-100% for finalization)
             if progress_info.phase == ProgressPhase.STRUCTURE:
                 # Phase 1: 0-10%, always show 5% while running, 10% when complete
@@ -648,11 +663,10 @@ class ChatScreen(Screen[None]):
                 overall_pct = 0.0
 
             # Update shared state (timer will render it)
-            spinner_state["phase_name"] = progress_info.phase_name
-            spinner_state["percentage"] = overall_pct
+            progress_state["percentage"] = overall_pct
 
-        # Start spinner animation timer (10 fps = 100ms interval)
-        spinner_timer = self.set_interval(0.1, update_spinner_display)
+        # Start progress animation timer (10 fps = 100ms interval)
+        progress_timer = self.set_interval(0.1, update_progress_display)
 
         try:
             # Pass the current working directory as the indexed_from_cwd
@@ -667,14 +681,12 @@ class ChatScreen(Screen[None]):
                 progress_callback=progress_callback,
             )
 
-            # Stop spinner animation
-            spinner_timer.stop()
+            # Stop progress animation
+            progress_timer.stop()
 
             # Show 100% completion after indexing finishes
             final_bar = create_progress_bar(100.0)
-            label.update(
-                f"[$foreground-muted]Indexing codebase: ✓ Complete {final_bar} 100%[/]"
-            )
+            label.update(f"[$foreground-muted]Indexing codebase: {final_bar} 100%[/]")
             label.refresh()
 
             logger.info(
@@ -687,12 +699,12 @@ class ChatScreen(Screen[None]):
             )
 
         except CodebaseAlreadyIndexedError as exc:
-            spinner_timer.stop()
+            progress_timer.stop()
             logger.warning(f"Codebase already indexed: {exc}")
             self.notify(str(exc), severity="warning")
             return
         except InvalidPathError as exc:
-            spinner_timer.stop()
+            progress_timer.stop()
             logger.error(f"Invalid path error: {exc}")
             self.notify(str(exc), severity="error")
 
@@ -704,8 +716,8 @@ class ChatScreen(Screen[None]):
             )
             self.notify(f"Failed to index codebase: {exc}", severity="error")
         finally:
-            # Always stop the spinner timer
-            spinner_timer.stop()
+            # Always stop the progress timer
+            progress_timer.stop()
             label.update("")
             label.refresh()
 
@@ -774,6 +786,7 @@ class ChatScreen(Screen[None]):
 
         # Update the current mode
         self.mode = AgentType(conversation.last_agent_model)
+        self.deps.usage_manager.restore_usage_state()
 
 
 def help_text_with_codebase(already_indexed: bool = False) -> str:

{shotgun_sh-0.1.15.dev2 → shotgun_sh-0.1.16.dev2}/src/shotgun/tui/screens/chat_screen/command_providers.py

@@ -96,6 +96,38 @@ class AgentModeProvider(Provider):
                 yield Hit(score, matcher.highlight(title), callback, help=help_text)
 
 
+class UsageProvider(Provider):
+    """Command provider for agent mode switching."""
+
+    @property
+    def chat_screen(self) -> "ChatScreen":
+        from shotgun.tui.screens.chat import ChatScreen
+
+        return cast(ChatScreen, self.screen)
+
+    async def discover(self) -> AsyncGenerator[DiscoveryHit, None]:
+        """Provide default mode switching commands when palette opens."""
+        yield DiscoveryHit(
+            "Show usage",
+            self.chat_screen.action_show_usage,
+            help="Display usage information for the current session",
+        )
+
+    async def search(self, query: str) -> AsyncGenerator[Hit, None]:
+        """Search for mode commands."""
+        matcher = self.matcher(query)
+
+        async for discovery_hit in self.discover():
+            score = matcher.match(discovery_hit.text or "")
+            if score > 0:
+                yield Hit(
+                    score,
+                    matcher.highlight(discovery_hit.text or ""),
+                    discovery_hit.command,
+                    help=discovery_hit.help,
+                )
+
+
 class ProviderSetupProvider(Provider):
     """Command palette entries for provider configuration."""
 

{shotgun_sh-0.1.15.dev2 → shotgun_sh-0.1.16.dev2}/src/shotgun/tui/screens/provider_config.py

@@ -9,7 +9,7 @@ from textual.app import ComposeResult
 from textual.containers import Horizontal, Vertical
 from textual.reactive import reactive
 from textual.screen import Screen
-from textual.widgets import Button, Input, Label, ListItem, ListView, Static
+from textual.widgets import Button, Input, Label, ListItem, ListView, Markdown, Static
 
 from shotgun.agents.config import ConfigManager, ProviderType
 
@@ -47,6 +47,10 @@ class ProviderConfigScreen(Screen[None]):
             color: $text-accent;
         }
 
+        #provider-links {
+            padding: 1 0;
+        }
+
         #provider-list {
             margin: 2 0;
             height: auto;
@@ -78,6 +82,10 @@ class ProviderConfigScreen(Screen[None]):
                 "Select a provider and enter the API key needed to activate it.",
                 id="provider-config-summary",
             )
+            yield Markdown(
+                "Don't have an API Key? Use these links to get one: [OpenAI](https://platform.openai.com/api-keys) | [Anthropic](https://console.anthropic.com) | [Google Gemini](https://aistudio.google.com)",
+                id="provider-links",
+            )
         yield ListView(*self._build_provider_items(), id="provider-list")
         yield Input(
             placeholder=self._input_placeholder(self.selected_provider),