shotgun-sh 0.1.16.dev1__py3-none-any.whl → 0.2.0__py3-none-any.whl

shotgun/agents/tools/web_search/anthropic.py

@@ -1,20 +1,24 @@
 """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.logging_config import get_logger
 
 logger = get_logger(__name__)
 
 
-def anthropic_web_search_tool(query: str) -> str:
-    """Perform a web search using Anthropic's Claude API with streaming.
+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 +31,49 @@ 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
     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)
+    # Build the request messages
+    messages: list[ModelMessage] = [
+        ModelRequest.user_text_prompt(f"Search for: {query}")
+    ]
 
-    # Use the Messages API with web search tool and streaming
+    # 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 +92,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 +113,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 = 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 +143,6 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    import asyncio
+
+    asyncio.run(main())

shotgun/agents/tools/web_search/gemini.py

@@ -1,20 +1,24 @@
 """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.logging_config import get_logger
 
 logger = get_logger(__name__)
 
 
-def gemini_web_search_tool(query: str) -> str:
+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,23 +33,16 @@ 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 = 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]
-
-    # Create a search-optimized prompt that leverages Gemini's knowledge
+    # Create a search-optimized prompt
     search_prompt = f"""Please provide current and accurate information about the following query:
 
 Query: {query}
@@ -56,17 +53,31 @@ Instructions:
 - Focus on current and recent information
 - Be specific and accurate in your response"""
 
-    # Generate response using the model's knowledge
+    # Build the request messages
+    messages: list[ModelMessage] = [ModelRequest.user_text_prompt(search_prompt)]
+
+    # 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(

shotgun/agents/tools/web_search/openai.py

@@ -1,6 +1,6 @@
 """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
@@ -10,7 +10,7 @@ from shotgun.logging_config import get_logger
 logger = get_logger(__name__)
 
 
-def openai_web_search_tool(query: str) -> str:
+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
@@ -54,8 +54,8 @@ Instructions:
 ALWAYS PROVIDE THE SOURCES (urls) TO BACK UP THE INFORMATION YOU PROVIDE.
 """
 
-        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/build_constants.py

@@ -12,8 +12,8 @@ POSTHOG_API_KEY = ''
 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/config.py

@@ -9,6 +9,7 @@ from rich.table import Table
 
 from shotgun.agents.config import ProviderType, get_config_manager
 from shotgun.logging_config import get_logger
+from shotgun.utils.env_utils import is_shotgun_account_enabled
 
 logger = get_logger(__name__)
 console = Console()
@@ -43,11 +44,11 @@ def init(
         console.print()
 
         # Initialize with defaults
-        config = config_manager.initialize()
+        config_manager.initialize()
 
-        # Ask for default provider
+        # Ask for provider
         provider_choices = ["openai", "anthropic", "google"]
-        console.print("Choose your default AI provider:")
+        console.print("Choose your AI provider:")
         for i, provider in enumerate(provider_choices, 1):
             console.print(f"  {i}. {provider}")
 
@@ -55,7 +56,7 @@ def init(
             try:
                 choice = typer.prompt("Enter choice (1-3)", type=int)
                 if 1 <= choice <= 3:
-                    config.default_provider = ProviderType(provider_choices[choice - 1])
+                    provider = ProviderType(provider_choices[choice - 1])
                     break
                 else:
                     console.print(
@@ -65,7 +66,6 @@ def init(
                 console.print("❌ Please enter a valid number.", style="red")
 
         # Ask for API key for the selected provider
-        provider = config.default_provider
         console.print(f"\n🔑 Setting up {provider.upper()} API key...")
 
         api_key = typer.prompt(
@@ -75,9 +75,9 @@ def init(
         )
 
         if api_key:
+            # update_provider will automatically set selected_model for first provider
             config_manager.update_provider(provider, api_key=api_key)
 
-        config_manager.save()
         console.print(
             f"\n✅ [bold green]Configuration saved to {config_manager.config_path}[/bold green]"
         )
@@ -98,16 +98,12 @@ def set(
         str | None,
         typer.Option("--api-key", "-k", help="API key for the provider"),
     ] = None,
-    default: Annotated[
-        bool,
-        typer.Option("--default", "-d", help="Set this provider as default"),
-    ] = False,
 ) -> None:
     """Set configuration for a specific provider."""
     config_manager = get_config_manager()
 
-    # If no API key provided via option and not just setting default, prompt for it
-    if api_key is None and not default:
+    # If no API key provided via option, prompt for it
+    if api_key is None:
         api_key = typer.prompt(
             f"Enter your {provider.upper()} API key",
             hide_input=True,
@@ -118,11 +114,6 @@ def set(
         if api_key:
             config_manager.update_provider(provider, api_key=api_key)
 
-        if default:
-            config = config_manager.load()
-            config.default_provider = provider
-            config_manager.save(config)
-
         console.print(f"✅ Configuration updated for {provider}")
 
     except Exception as e:
@@ -130,41 +121,6 @@ def set(
         raise typer.Exit(1) from e
 
 
-@app.command()
-def set_default(
-    provider: Annotated[
-        ProviderType,
-        typer.Argument(
-            help="AI provider to set as default (openai, anthropic, google)"
-        ),
-    ],
-) -> None:
-    """Set the default AI provider without modifying API keys."""
-    config_manager = get_config_manager()
-
-    try:
-        config = config_manager.load()
-
-        # Check if the provider has an API key configured
-        provider_config = getattr(config, provider.value)
-        if not provider_config.api_key:
-            console.print(
-                f"⚠️  Warning: {provider.upper()} does not have an API key configured.",
-                style="yellow",
-            )
-            console.print(f"Use 'shotgun config set {provider}' to configure it.")
-
-        # Set as default
-        config.default_provider = provider
-        config_manager.save(config)
-
-        console.print(f"✅ Default provider set to: {provider}")
-
-    except Exception as e:
-        console.print(f"❌ Failed to set default provider: {e}", style="red")
-        raise typer.Exit(1) from e
-
-
 @app.command()
 def get(
     provider: Annotated[
@@ -201,16 +157,23 @@ def _show_full_config(config: Any) -> None:
     table.add_column("Setting", style="cyan")
     table.add_column("Value", style="white")
 
-    # Default provider
-    table.add_row("Default Provider", f"[bold]{config.default_provider}[/bold]")
+    # Selected model
+    selected_model = config.selected_model or "None (will auto-detect)"
+    table.add_row("Selected Model", f"[bold]{selected_model}[/bold]")
     table.add_row("", "")  # Separator
 
     # Provider configurations
-    for provider_name, provider_config in [
+    providers_to_show = [
         ("OpenAI", config.openai),
         ("Anthropic", config.anthropic),
         ("Google", config.google),
-    ]:
+    ]
+
+    # Only show Shotgun Account if feature flag is enabled
+    if is_shotgun_account_enabled():
+        providers_to_show.append(("Shotgun Account", config.shotgun))
+
+    for provider_name, provider_config in providers_to_show:
         table.add_row(f"[bold]{provider_name}[/bold]", "")
 
         # API Key
@@ -231,6 +194,8 @@ def _show_provider_config(provider: ProviderType, config: Any) -> None:
         provider_config = config.anthropic
     elif provider_str == "google":
         provider_config = config.google
+    elif provider_str == "shotgun":
+        provider_config = config.shotgun
     else:
         console.print(f"❌ Unknown provider: {provider}", style="red")
         return
@@ -248,7 +213,13 @@ def _show_provider_config(provider: ProviderType, config: Any) -> None:
 
 def _mask_secrets(data: dict[str, Any]) -> None:
     """Mask secrets in configuration data."""
-    for provider in ["openai", "anthropic", "google"]:
+    providers = ["openai", "anthropic", "google"]
+
+    # Only mask shotgun if feature flag is enabled
+    if is_shotgun_account_enabled():
+        providers.append("shotgun")
+
+    for provider in providers:
         if provider in data and isinstance(data[provider], dict):
             if "api_key" in data[provider] and data[provider]["api_key"]:
                 data[provider]["api_key"] = _mask_value(data[provider]["api_key"])

shotgun/cli/models.py

@@ -1,9 +1,9 @@
 """Common models for CLI commands."""
 
-from enum import Enum
+from enum import StrEnum
 
 
-class OutputFormat(str, Enum):
+class OutputFormat(StrEnum):
     """Output format options for CLI commands."""
 
     TEXT = "text"

shotgun/codebase/models.py

@@ -1,13 +1,13 @@
 """Data models for codebase service."""
 
 from collections.abc import Callable
-from enum import Enum
+from enum import StrEnum
 from typing import Any
 
 from pydantic import BaseModel, Field
 
 
-class GraphStatus(str, Enum):
+class GraphStatus(StrEnum):
     """Status of a code knowledge graph."""
 
     READY = "READY"  # Graph is ready for queries
@@ -16,14 +16,14 @@ class GraphStatus(str, Enum):
     ERROR = "ERROR"  # Last operation failed
 
 
-class QueryType(str, Enum):
+class QueryType(StrEnum):
     """Type of query being executed."""
 
     NATURAL_LANGUAGE = "natural_language"
     CYPHER = "cypher"
 
 
-class ProgressPhase(str, Enum):
+class ProgressPhase(StrEnum):
     """Phase of codebase indexing progress."""
 
     STRUCTURE = "structure"  # Identifying packages and folders

shotgun/llm_proxy/__init__.py

@@ -0,0 +1,16 @@
+"""LiteLLM proxy client utilities and configuration."""
+
+from .clients import create_anthropic_proxy_client, create_litellm_provider
+from .constants import (
+    LITELLM_PROXY_ANTHROPIC_BASE,
+    LITELLM_PROXY_BASE_URL,
+    LITELLM_PROXY_OPENAI_BASE,
+)
+
+__all__ = [
+    "LITELLM_PROXY_BASE_URL",
+    "LITELLM_PROXY_ANTHROPIC_BASE",
+    "LITELLM_PROXY_OPENAI_BASE",
+    "create_litellm_provider",
+    "create_anthropic_proxy_client",
+]

shotgun/llm_proxy/clients.py

@@ -0,0 +1,39 @@
+"""Client creation utilities for LiteLLM proxy."""
+
+from anthropic import Anthropic
+from pydantic_ai.providers.litellm import LiteLLMProvider
+
+from .constants import LITELLM_PROXY_ANTHROPIC_BASE, LITELLM_PROXY_BASE_URL
+
+
+def create_litellm_provider(api_key: str) -> LiteLLMProvider:
+    """Create LiteLLM provider for Shotgun Account.
+
+    Args:
+        api_key: Shotgun API key
+
+    Returns:
+        Configured LiteLLM provider pointing to Shotgun's proxy
+    """
+    return LiteLLMProvider(
+        api_base=LITELLM_PROXY_BASE_URL,
+        api_key=api_key,
+    )
+
+
+def create_anthropic_proxy_client(api_key: str) -> Anthropic:
+    """Create Anthropic client configured for LiteLLM proxy.
+
+    This client will proxy token counting requests through the
+    LiteLLM proxy to Anthropic's actual token counting API.
+
+    Args:
+        api_key: Shotgun API key
+
+    Returns:
+        Anthropic client configured to use LiteLLM proxy
+    """
+    return Anthropic(
+        api_key=api_key,
+        base_url=LITELLM_PROXY_ANTHROPIC_BASE,
+    )

shotgun/llm_proxy/constants.py

@@ -0,0 +1,8 @@
+"""LiteLLM proxy constants and configuration."""
+
+# Shotgun's LiteLLM proxy base URL
+LITELLM_PROXY_BASE_URL = "https://litellm-701197220809.us-east1.run.app"
+
+# Provider-specific endpoints
+LITELLM_PROXY_ANTHROPIC_BASE = f"{LITELLM_PROXY_BASE_URL}/anthropic"
+LITELLM_PROXY_OPENAI_BASE = LITELLM_PROXY_BASE_URL

shotgun/main.py

@@ -1,5 +1,11 @@
 """Main CLI application for shotgun."""
 
+# NOTE: These are before we import any Google library to stop the noisy gRPC logs.
+import os  # noqa: I001
+
+os.environ["GRPC_VERBOSITY"] = "ERROR"
+os.environ["GLOG_minloglevel"] = "2"
+
 import logging
 
 # CRITICAL: Add NullHandler to root logger before ANY other imports.

shotgun/posthog_telemetry.py

@@ -1,6 +1,6 @@
 """PostHog analytics setup for Shotgun."""
 
-from enum import Enum
+from enum import StrEnum
 from typing import Any
 
 import posthog
@@ -137,7 +137,7 @@ def shutdown() -> None:
             _posthog_client = None
 
 
-class FeedbackKind(str, Enum):
+class FeedbackKind(StrEnum):
     BUG = "bug"
     FEATURE = "feature"
     OTHER = "other"
@@ -178,7 +178,9 @@ def submit_feedback_survey(feedback: Feedback) -> None:
             ],
             f"$survey_response_{Q_KIND_ID}": feedback.kind,
             f"$survey_response_{Q_DESCRIPTION_ID}": feedback.description,
-            "provider": config.default_provider.value,
+            "selected_model": config.selected_model.value
+            if config.selected_model
+            else None,
             "config_version": config.config_version,
             "last_10_messages": last_10_messages,  # last 10 messages
         },

shotgun/tui/app.py

@@ -14,6 +14,7 @@ from shotgun.utils.update_checker import perform_auto_update_async
 from .screens.chat import ChatScreen
 from .screens.directory_setup import DirectorySetupScreen
 from .screens.feedback import FeedbackScreen
+from .screens.model_picker import ModelPickerScreen
 from .screens.provider_config import ProviderConfigScreen
 
 logger = get_logger(__name__)
@@ -23,6 +24,7 @@ class ShotgunApp(App[None]):
     SCREENS = {
         "chat": ChatScreen,
         "provider_config": ProviderConfigScreen,
+        "model_picker": ModelPickerScreen,
         "directory_setup": DirectorySetupScreen,
         "feedback": FeedbackScreen,
     }
@@ -62,7 +64,8 @@ class ShotgunApp(App[None]):
                 return
 
             self.push_screen(
-                "provider_config", callback=lambda _arg: self.refresh_startup_screen()
+                ProviderConfigScreen(),
+                callback=lambda _arg: self.refresh_startup_screen(),
             )
             return
 
@@ -71,7 +74,8 @@ class ShotgunApp(App[None]):
                 return
 
             self.push_screen(
-                "directory_setup", callback=lambda _arg: self.refresh_startup_screen()
+                DirectorySetupScreen(),
+                callback=lambda _arg: self.refresh_startup_screen(),
             )
             return
 
@@ -108,7 +112,7 @@ class ShotgunApp(App[None]):
                 submit_feedback_survey(feedback)
                 self.notify("Feedback sent. Thank you!")
 
-        self.push_screen("feedback", callback=handle_feedback)
+        self.push_screen(FeedbackScreen(), callback=handle_feedback)
 
 
 def run(no_update_check: bool = False, continue_session: bool = False) -> None: