shotgun-sh 0.2.11.dev7__py3-none-any.whl → 0.2.23.dev1__py3-none-any.whl

shotgun/llm_proxy/client.py

@@ -0,0 +1,215 @@
+"""HTTP client for LiteLLM Proxy API."""
+
+import logging
+from typing import Any
+
+import httpx
+from tenacity import (
+    before_sleep_log,
+    retry,
+    retry_if_exception,
+    stop_after_attempt,
+    wait_exponential_jitter,
+)
+
+from shotgun.api_endpoints import LITELLM_PROXY_BASE_URL
+from shotgun.logging_config import get_logger
+
+from .models import BudgetInfo, KeyInfoResponse, TeamInfoResponse
+
+logger = get_logger(__name__)
+
+
+def _is_retryable_http_error(exception: BaseException) -> bool:
+    """Check if HTTP exception should trigger a retry.
+
+    Args:
+        exception: The exception to check
+
+    Returns:
+        True if the exception is a transient error that should be retried
+    """
+    # Retry on network errors and timeouts
+    if isinstance(exception, (httpx.RequestError, httpx.TimeoutException)):
+        return True
+
+    # Retry on server errors (5xx) and rate limits (429)
+    if isinstance(exception, httpx.HTTPStatusError):
+        status_code = exception.response.status_code
+        return status_code >= 500 or status_code == 429
+
+    # Don't retry on other errors (e.g., 4xx client errors)
+    return False
+
+
+class LiteLLMProxyClient:
+    """HTTP client for LiteLLM Proxy API.
+
+    Provides methods to query budget information and key/team metadata
+    from a LiteLLM proxy server.
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        base_url: str | None = None,
+        timeout: float = 10.0,
+    ):
+        """Initialize LiteLLM Proxy client.
+
+        Args:
+            api_key: LiteLLM API key for authentication
+            base_url: Base URL for LiteLLM proxy. If None, uses LITELLM_PROXY_BASE_URL
+            timeout: Request timeout in seconds
+        """
+        self.api_key = api_key
+        self.base_url = base_url or LITELLM_PROXY_BASE_URL
+        self.timeout = timeout
+
+    @retry(
+        stop=stop_after_attempt(3),
+        wait=wait_exponential_jitter(initial=1, max=8),
+        retry=retry_if_exception(_is_retryable_http_error),
+        before_sleep=before_sleep_log(logger, logging.WARNING),
+        reraise=True,
+    )
+    async def _request_with_retry(
+        self,
+        method: str,
+        url: str,
+        **kwargs: Any,
+    ) -> httpx.Response:
+        """Make async HTTP request with exponential backoff retry and jitter.
+
+        Uses tenacity to retry on transient errors (5xx, 429, network errors)
+        with exponential backoff and jitter. Client errors (4xx except 429)
+        are not retried.
+
+        Args:
+            method: HTTP method (GET, POST, etc.)
+            url: Request URL
+            **kwargs: Additional arguments to pass to httpx request
+
+        Returns:
+            HTTP response
+
+        Raises:
+            httpx.HTTPError: If request fails after all retries
+        """
+        async with httpx.AsyncClient(timeout=self.timeout) as client:
+            response = await client.request(method, url, **kwargs)
+            response.raise_for_status()
+            return response
+
+    async def get_key_info(self) -> KeyInfoResponse:
+        """Get key information from LiteLLM proxy.
+
+        Returns:
+            Key information including spend, budget, and team_id
+
+        Raises:
+            httpx.HTTPError: If request fails
+        """
+        url = f"{self.base_url}/key/info"
+        params = {"key": self.api_key}
+        headers = {"Authorization": f"Bearer {self.api_key}"}
+
+        logger.debug("Fetching key info from %s", url)
+
+        response = await self._request_with_retry(
+            "GET", url, params=params, headers=headers
+        )
+
+        data = response.json()
+        result = KeyInfoResponse.model_validate(data)
+
+        logger.info(
+            "Successfully fetched key info: key_alias=%s, team_id=%s",
+            result.info.key_alias,
+            result.info.team_id,
+        )
+        return result
+
+    async def get_team_info(self, team_id: str) -> TeamInfoResponse:
+        """Get team information from LiteLLM proxy.
+
+        Args:
+            team_id: Team identifier
+
+        Returns:
+            Team information including spend and budget
+
+        Raises:
+            httpx.HTTPError: If request fails
+        """
+        url = f"{self.base_url}/team/info"
+        params = {"team_id": team_id}
+        headers = {"Authorization": f"Bearer {self.api_key}"}
+
+        logger.debug("Fetching team info from %s for team_id=%s", url, team_id)
+
+        response = await self._request_with_retry(
+            "GET", url, params=params, headers=headers
+        )
+
+        data = response.json()
+        result = TeamInfoResponse.model_validate(data)
+
+        logger.info(
+            "Successfully fetched team info: team_alias=%s",
+            result.team_info.team_alias,
+        )
+        return result
+
+    async def get_budget_info(self) -> BudgetInfo:
+        """Get team-level budget information for this key.
+
+        Budget is always configured at the team level, never at the key level.
+        This method fetches the team_id from the key info, then retrieves
+        the team's budget information.
+
+        Returns:
+            Team-level budget information
+
+        Raises:
+            httpx.HTTPError: If request fails
+            ValueError: If team has no budget configured
+        """
+        logger.debug("Fetching budget info")
+
+        # Get key info to retrieve team_id
+        key_response = await self.get_key_info()
+        key_info = key_response.info
+
+        # Fetch team budget (budget is always at team level)
+        logger.debug(
+            "Fetching team budget for team_id=%s",
+            key_info.team_id,
+        )
+        team_response = await self.get_team_info(key_info.team_id)
+        team_info = team_response.team_info
+
+        if team_info.max_budget is None:
+            raise ValueError(
+                f"Team (team_id={key_info.team_id}) has no max_budget configured"
+            )
+
+        logger.debug("Using team-level budget: $%.6f", team_info.max_budget)
+        return BudgetInfo.from_team_info(team_info)
+
+
+# Convenience function for standalone use
+async def get_budget_info(api_key: str, base_url: str | None = None) -> BudgetInfo:
+    """Get budget information for an API key.
+
+    Convenience function that creates a client and calls get_budget_info.
+
+    Args:
+        api_key: LiteLLM API key
+        base_url: Optional base URL for LiteLLM proxy
+
+    Returns:
+        Budget information
+    """
+    client = LiteLLMProxyClient(api_key, base_url=base_url)
+    return await client.get_budget_info()

shotgun/llm_proxy/models.py

@@ -0,0 +1,137 @@
+"""Pydantic models for LiteLLM Proxy API."""
+
+from enum import StrEnum
+
+from pydantic import BaseModel, Field
+
+
+class BudgetSource(StrEnum):
+    """Source of budget information."""
+
+    KEY = "key"
+    TEAM = "team"
+
+
+class KeyInfoData(BaseModel):
+    """Key information data from /key/info endpoint."""
+
+    key_name: str = Field(description="Key name/identifier")
+    key_alias: str | None = Field(default=None, description="Human-readable key alias")
+    spend: float = Field(description="Current spend for this key in USD")
+    max_budget: float | None = Field(
+        default=None, description="Maximum budget for this key in USD"
+    )
+    team_id: str = Field(description="Team ID associated with this key")
+    user_id: str = Field(description="User ID associated with this key")
+    models: list[str] = Field(
+        default_factory=list, description="List of models available to this key"
+    )
+
+
+class KeyInfoResponse(BaseModel):
+    """Response from /key/info endpoint."""
+
+    key: str = Field(description="The API key")
+    info: KeyInfoData = Field(description="Key information data")
+
+
+class TeamInfoData(BaseModel):
+    """Team information data from /team/info endpoint."""
+
+    team_id: str = Field(description="Team identifier")
+    team_alias: str | None = Field(
+        default=None, description="Human-readable team alias"
+    )
+    max_budget: float | None = Field(
+        default=None, description="Maximum budget for this team in USD"
+    )
+    spend: float = Field(description="Current spend for this team in USD")
+    models: list[str] = Field(
+        default_factory=list, description="List of models available to this team"
+    )
+
+
+class TeamInfoResponse(BaseModel):
+    """Response from /team/info endpoint."""
+
+    team_id: str = Field(description="Team identifier")
+    team_info: TeamInfoData = Field(description="Team information data")
+
+
+class BudgetInfo(BaseModel):
+    """Unified budget information.
+
+    Combines key and team budget information to provide a single view
+    of budget status. Budget can come from either key-level or team-level,
+    with key-level taking priority if set.
+    """
+
+    max_budget: float = Field(description="Maximum budget in USD")
+    spend: float = Field(description="Current spend in USD")
+    remaining: float = Field(description="Remaining budget in USD")
+    source: BudgetSource = Field(
+        description="Source of budget information (key or team)"
+    )
+    percentage_used: float = Field(description="Percentage of budget used (0-100)")
+
+    @classmethod
+    def from_key_info(cls, key_info: KeyInfoData) -> "BudgetInfo":
+        """Create BudgetInfo from key-level budget.
+
+        Args:
+            key_info: Key information containing budget data
+
+        Returns:
+            BudgetInfo instance with key-level budget
+
+        Raises:
+            ValueError: If key does not have max_budget set
+        """
+        if key_info.max_budget is None:
+            raise ValueError("Key does not have max_budget set")
+
+        remaining = key_info.max_budget - key_info.spend
+        percentage_used = (
+            (key_info.spend / key_info.max_budget * 100)
+            if key_info.max_budget > 0
+            else 0.0
+        )
+
+        return cls(
+            max_budget=key_info.max_budget,
+            spend=key_info.spend,
+            remaining=remaining,
+            source=BudgetSource.KEY,
+            percentage_used=percentage_used,
+        )
+
+    @classmethod
+    def from_team_info(cls, team_info: TeamInfoData) -> "BudgetInfo":
+        """Create BudgetInfo from team-level budget.
+
+        Args:
+            team_info: Team information containing budget data
+
+        Returns:
+            BudgetInfo instance with team-level budget
+
+        Raises:
+            ValueError: If team does not have max_budget set
+        """
+        if team_info.max_budget is None:
+            raise ValueError("Team does not have max_budget set")
+
+        remaining = team_info.max_budget - team_info.spend
+        percentage_used = (
+            (team_info.spend / team_info.max_budget * 100)
+            if team_info.max_budget > 0
+            else 0.0
+        )
+
+        return cls(
+            max_budget=team_info.max_budget,
+            spend=team_info.spend,
+            remaining=remaining,
+            source=BudgetSource.TEAM,
+            percentage_used=percentage_used,
+        )

shotgun/logging_config.py

@@ -27,6 +27,44 @@ def get_log_directory() -> Path:
     return log_dir
 
 
+def cleanup_old_log_files(log_dir: Path, max_files: int) -> None:
+    """Remove old log files, keeping only the most recent ones.
+
+    Also removes the legacy shotgun.log file if it exists.
+
+    Args:
+        log_dir: Directory containing log files
+        max_files: Maximum number of log files to keep
+    """
+    try:
+        # Remove legacy non-timestamped log file if it exists
+        legacy_log = log_dir / "shotgun.log"
+        if legacy_log.exists():
+            try:
+                legacy_log.unlink()
+            except OSError:
+                pass  # noqa: S110
+
+        # Find all shotgun log files
+        log_files = sorted(
+            log_dir.glob("shotgun-*.log"),
+            key=lambda p: p.stat().st_mtime,
+            reverse=True,  # Newest first
+        )
+
+        # Remove files beyond the limit
+        files_to_delete = log_files[max_files:]
+        for log_file in files_to_delete:
+            try:
+                log_file.unlink()
+            except OSError:
+                # Ignore errors when deleting individual files
+                pass  # noqa: S110
+    except Exception:  # noqa: S110
+        # Silently fail - log cleanup shouldn't break the application
+        pass
+
+
 class ColoredFormatter(logging.Formatter):
     """Custom formatter with colors for different log levels."""
 
@@ -123,6 +161,10 @@ def setup_logger(
         try:
             # Create file handler with ISO8601 timestamp for each run
             log_dir = get_log_directory()
+
+            # Clean up old log files before creating a new one
+            cleanup_old_log_files(log_dir, settings.logging.max_log_files)
+
             log_file = log_dir / f"shotgun-{_RUN_TIMESTAMP}.log"
 
             # Use regular FileHandler - each run gets its own isolated log file

shotgun/main.py

@@ -55,6 +55,8 @@ logger = get_logger(__name__)
 logger.debug("Logfire observability enabled: %s", _logfire_enabled)
 
 # Initialize configuration
+# Note: If config migration fails, ConfigManager will auto-create fresh config
+# and set migration_failed flag for user notification
 try:
     import asyncio
 

shotgun/posthog_telemetry.py

@@ -18,6 +18,9 @@ logger = get_early_logger(__name__)
 # Global PostHog client instance
 _posthog_client = None
 
+# Cache the shotgun instance ID to avoid async calls during event tracking
+_shotgun_instance_id: str | None = None
+
 
 def setup_posthog_observability() -> bool:
     """Set up PostHog analytics for usage tracking.
@@ -25,7 +28,7 @@ def setup_posthog_observability() -> bool:
     Returns:
         True if PostHog was successfully set up, False otherwise
     """
-    global _posthog_client
+    global _posthog_client, _shotgun_instance_id
 
     try:
         # Check if PostHog is already initialized
@@ -57,31 +60,20 @@ def setup_posthog_observability() -> bool:
         # Store the client for later use
         _posthog_client = posthog
 
-        # Set user context with anonymous shotgun instance ID from config
+        # Cache the shotgun instance ID for later use (avoids async issues)
         try:
             import asyncio
 
             config_manager = get_config_manager()
-            shotgun_instance_id = asyncio.run(config_manager.get_shotgun_instance_id())
-
-            # Identify the user in PostHog
-            posthog.identify(  # type: ignore[attr-defined]
-                distinct_id=shotgun_instance_id,
-                properties={
-                    "version": __version__,
-                    "environment": environment,
-                },
-            )
-
-            # Set default properties for all events
-            posthog.disabled = False
-            posthog.personal_api_key = None  # Not needed for event tracking
+            _shotgun_instance_id = asyncio.run(config_manager.get_shotgun_instance_id())
 
             logger.debug(
-                "PostHog user identified with anonymous ID: %s", shotgun_instance_id
+                "PostHog initialized with shotgun instance ID: %s",
+                _shotgun_instance_id,
             )
         except Exception as e:
-            logger.warning("Failed to set user context: %s", e)
+            logger.warning("Failed to load shotgun instance ID: %s", e)
+            # Continue anyway - we'll try to get it during event tracking
 
         logger.debug(
             "PostHog analytics configured successfully (environment: %s, version: %s)",
@@ -102,18 +94,19 @@ def track_event(event_name: str, properties: dict[str, Any] | None = None) -> No
         event_name: Name of the event to track
         properties: Optional properties to include with the event
     """
-    global _posthog_client
+    global _posthog_client, _shotgun_instance_id
 
     if _posthog_client is None:
         logger.debug("PostHog not initialized, skipping event: %s", event_name)
         return
 
     try:
-        import asyncio
-
-        # Get shotgun instance ID for tracking
-        config_manager = get_config_manager()
-        shotgun_instance_id = asyncio.run(config_manager.get_shotgun_instance_id())
+        # Use cached instance ID (loaded during setup)
+        if _shotgun_instance_id is None:
+            logger.warning(
+                "Shotgun instance ID not available, skipping event: %s", event_name
+            )
+            return
 
         # Add version and environment to properties
         if properties is None:
@@ -128,7 +121,7 @@ def track_event(event_name: str, properties: dict[str, Any] | None = None) -> No
 
         # Track the event using PostHog's capture method
         _posthog_client.capture(
-            distinct_id=shotgun_instance_id, event=event_name, properties=properties
+            distinct_id=_shotgun_instance_id, event=event_name, properties=properties
         )
         logger.debug("Tracked PostHog event: %s", event_name)
     except Exception as e:

shotgun/prompts/agents/partials/common_agent_system_prompt.j2

@@ -7,10 +7,11 @@ Your extensive expertise spans, among other things:
 ## KEY RULES
 
 {% if interactive_mode %}
-0. Always ask CLARIFYING QUESTIONS using structured output if the user's request is ambiguous or lacks sufficient detail.
+0. Always ask CLARIFYING QUESTIONS using structured output before doing work.
    - Return your response with the clarifying_questions field populated
-   - Do not make assumptions about what the user wants
+   - Do not make assumptions about what the user wants, get a clear understanding first.
    - Questions should be clear, specific, and answerable
+   - Do not ask too many questions that might overwhelm the user; prioritize the most important ones.
 {% endif %}
 1. Above all, prefer using tools to do the work and NEVER respond with text.
 2. IMPORTANT: Always ask for review and go ahead to move forward after using write_file().

shotgun/sdk/codebase.py

@@ -93,6 +93,19 @@ class CodebaseSDK:
         if indexed_from_cwd is None:
             indexed_from_cwd = str(Path.cwd().resolve())
 
+        # Track codebase indexing started event
+        source = detect_source()
+        logger.debug(
+            "Tracking codebase_index_started event: source=%s",
+            source,
+        )
+        track_event(
+            "codebase_index_started",
+            {
+                "source": source,
+            },
+        )
+
         graph = await self.service.create_graph(
             resolved_path,
             name,
@@ -101,9 +114,7 @@ class CodebaseSDK:
         )
         file_count = sum(graph.language_stats.values()) if graph.language_stats else 0
 
-        # Track codebase indexing event
-        # Detect if called from TUI by checking the call stack
-        source = detect_source()
+        # Track codebase indexing completion event (reuse source from start event)
 
         logger.debug(
             "Tracking codebase_indexed event: file_count=%d, node_count=%d, relationship_count=%d, source=%s",