shotgun-sh 0.2.17__py3-none-any.whl → 0.4.0.dev1__py3-none-any.whl

shotgun/exceptions.py

@@ -1,13 +1,57 @@
 """General exceptions for Shotgun application."""
 
+from shotgun.utils import get_shotgun_home
+
+# Shotgun Account signup URL for BYOK users
+SHOTGUN_SIGNUP_URL = "https://shotgun.sh"
+SHOTGUN_CONTACT_EMAIL = "contact@shotgun.sh"
+
 
 class ErrorNotPickedUpBySentry(Exception):  # noqa: N818
     """Base for user-actionable errors that shouldn't be sent to Sentry.
 
     These errors represent expected user conditions requiring action
     rather than bugs that need tracking.
+
+    All subclasses should implement to_markdown() and to_plain_text() methods
+    for consistent error message formatting.
     """
 
+    def to_markdown(self) -> str:
+        """Generate markdown-formatted error message for TUI.
+
+        Subclasses should override this method.
+        """
+        return f"⚠️ {str(self)}"
+
+    def to_plain_text(self) -> str:
+        """Generate plain text error message for CLI.
+
+        Subclasses should override this method.
+        """
+        return f"⚠️  {str(self)}"
+
+
+# ============================================================================
+# User Action Required Errors
+# ============================================================================
+
+
+class AgentCancelledException(ErrorNotPickedUpBySentry):
+    """Raised when user cancels an agent operation."""
+
+    def __init__(self) -> None:
+        """Initialize the exception."""
+        super().__init__("Operation cancelled by user")
+
+    def to_markdown(self) -> str:
+        """Generate markdown-formatted error message for TUI."""
+        return "⚠️ Operation cancelled by user"
+
+    def to_plain_text(self) -> str:
+        """Generate plain text error message for CLI."""
+        return "⚠️  Operation cancelled by user"
+
 
 class ContextSizeLimitExceeded(ErrorNotPickedUpBySentry):
     """Raised when conversation context exceeds the model's limits.
@@ -30,3 +74,284 @@ class ContextSizeLimitExceeded(ErrorNotPickedUpBySentry):
         super().__init__(
             f"Context too large for {model_name} (limit: {max_tokens:,} tokens)"
         )
+
+    def to_markdown(self) -> str:
+        """Generate markdown-formatted error message for TUI."""
+        return (
+            f"⚠️ **Context too large for {self.model_name}**\n\n"
+            f"Your conversation history exceeds this model's limit ({self.max_tokens:,} tokens).\n\n"
+            f"**Choose an action:**\n\n"
+            f"1. Switch to a larger model (`Ctrl+P` → Change Model)\n"
+            f"2. Switch to a larger model, compact (`/compact`), then switch back to {self.model_name}\n"
+            f"3. Clear conversation (`/clear`)\n"
+        )
+
+    def to_plain_text(self) -> str:
+        """Generate plain text error message for CLI."""
+        return (
+            f"⚠️  Context too large for {self.model_name}\n\n"
+            f"Your conversation history exceeds this model's limit ({self.max_tokens:,} tokens).\n\n"
+            f"Choose an action:\n"
+            f"1. Switch to a larger model\n"
+            f"2. Switch to a larger model, compact, then switch back to {self.model_name}\n"
+            f"3. Clear conversation\n"
+        )
+
+
+# ============================================================================
+# Shotgun Account Errors (show contact email in TUI)
+# ============================================================================
+
+
+class ShotgunAccountException(ErrorNotPickedUpBySentry):
+    """Base class for Shotgun Account service errors.
+
+    TUI will check isinstance() of this class to show contact email UI.
+    """
+
+
+class BudgetExceededException(ShotgunAccountException):
+    """Raised when Shotgun Account budget has been exceeded.
+
+    This is a user-actionable error - they need to contact support
+    to increase their budget limit. This is a temporary exception
+    until self-service budget increases are implemented.
+    """
+
+    def __init__(
+        self,
+        current_cost: float | None = None,
+        max_budget: float | None = None,
+        message: str | None = None,
+    ):
+        """Initialize the exception.
+
+        Args:
+            current_cost: Current total spend/cost (optional)
+            max_budget: Maximum budget limit (optional)
+            message: Optional custom error message from API
+        """
+        self.current_cost = current_cost
+        self.max_budget = max_budget
+        self.api_message = message
+
+        if message:
+            error_msg = message
+        elif current_cost is not None and max_budget is not None:
+            error_msg = f"Budget exceeded: ${current_cost:.2f} / ${max_budget:.2f}"
+        else:
+            error_msg = "Budget exceeded"
+
+        super().__init__(error_msg)
+
+    def to_markdown(self) -> str:
+        """Generate markdown-formatted error message for TUI.
+
+        Note: TUI will detect ShotgunAccountException and automatically
+        show email contact UI component.
+        """
+        return (
+            "⚠️ **Your Shotgun Account budget has been exceeded!**\n\n"
+            "Your account has reached its spending limit and cannot process more requests.\n\n"
+            "**Action Required:** Top up your account to continue using Shotgun.\n\n"
+            "👉 **[Top Up Now at https://app.shotgun.sh/dashboard](https://app.shotgun.sh/dashboard)**\n\n"
+            "**Need help?** Contact us if you have questions about your budget.\n\n"
+            f"_Error details: {str(self)}_"
+        )
+
+    def to_plain_text(self) -> str:
+        """Generate plain text error message for CLI."""
+        return (
+            "⚠️  Your Shotgun Account budget has been exceeded!\n\n"
+            "Your account has reached its spending limit and cannot process more requests.\n\n"
+            "Action Required: Top up your account to continue using Shotgun.\n\n"
+            "→ Top Up Now: https://app.shotgun.sh/dashboard\n\n"
+            f"Need help? Contact: {SHOTGUN_CONTACT_EMAIL}\n\n"
+            f"Error details: {str(self)}"
+        )
+
+
+class ShotgunServiceOverloadException(ShotgunAccountException):
+    """Raised when Shotgun Account AI service is overloaded."""
+
+    def __init__(self, message: str | None = None):
+        """Initialize the exception.
+
+        Args:
+            message: Optional custom error message from API
+        """
+        super().__init__(message or "Service temporarily overloaded")
+
+    def to_markdown(self) -> str:
+        """Generate markdown-formatted error message for TUI."""
+        return "⚠️ The AI service is temporarily overloaded. Please wait a moment and try again."
+
+    def to_plain_text(self) -> str:
+        """Generate plain text error message for CLI."""
+        return "⚠️  The AI service is temporarily overloaded. Please wait a moment and try again."
+
+
+class ShotgunRateLimitException(ShotgunAccountException):
+    """Raised when Shotgun Account rate limit is reached."""
+
+    def __init__(self, message: str | None = None):
+        """Initialize the exception.
+
+        Args:
+            message: Optional custom error message from API
+        """
+        super().__init__(message or "Rate limit reached")
+
+    def to_markdown(self) -> str:
+        """Generate markdown-formatted error message for TUI."""
+        return "⚠️ Rate limit reached. Please wait before trying again."
+
+    def to_plain_text(self) -> str:
+        """Generate plain text error message for CLI."""
+        return "⚠️  Rate limit reached. Please wait before trying again."
+
+
+# ============================================================================
+# BYOK (Bring Your Own Key) API Errors
+# ============================================================================
+
+
+class BYOKAPIException(ErrorNotPickedUpBySentry):
+    """Base class for BYOK API errors.
+
+    All BYOK errors suggest using Shotgun Account to avoid the issue.
+    """
+
+    def __init__(self, message: str, specific_error: str = "API error"):
+        """Initialize the exception.
+
+        Args:
+            message: The error message from the API
+            specific_error: Human-readable error type label
+        """
+        self.api_message = message
+        self.specific_error = specific_error
+        super().__init__(message)
+
+    def to_markdown(self) -> str:
+        """Generate markdown-formatted error message for TUI."""
+        return (
+            f"⚠️ **{self.specific_error}**: {self.api_message}\n\n"
+            f"_This could be avoided with a [Shotgun Account]({SHOTGUN_SIGNUP_URL})._"
+        )
+
+    def to_plain_text(self) -> str:
+        """Generate plain text error message for CLI."""
+        return (
+            f"⚠️  {self.specific_error}: {self.api_message}\n\n"
+            f"This could be avoided with a Shotgun Account: {SHOTGUN_SIGNUP_URL}"
+        )
+
+
+class BYOKRateLimitException(BYOKAPIException):
+    """Raised when BYOK user hits rate limit."""
+
+    def __init__(self, message: str):
+        """Initialize the exception.
+
+        Args:
+            message: The error message from the API
+        """
+        super().__init__(message, specific_error="Rate limit reached")
+
+
+class BYOKQuotaBillingException(BYOKAPIException):
+    """Raised when BYOK user has quota or billing issues."""
+
+    def __init__(self, message: str):
+        """Initialize the exception.
+
+        Args:
+            message: The error message from the API
+        """
+        super().__init__(message, specific_error="Quota or billing issue")
+
+
+class BYOKAuthenticationException(BYOKAPIException):
+    """Raised when BYOK authentication fails."""
+
+    def __init__(self, message: str):
+        """Initialize the exception.
+
+        Args:
+            message: The error message from the API
+        """
+        super().__init__(message, specific_error="Authentication error")
+
+
+class BYOKServiceOverloadException(BYOKAPIException):
+    """Raised when BYOK service is overloaded."""
+
+    def __init__(self, message: str):
+        """Initialize the exception.
+
+        Args:
+            message: The error message from the API
+        """
+        super().__init__(message, specific_error="Service overloaded")
+
+
+class BYOKGenericAPIException(BYOKAPIException):
+    """Raised for generic BYOK API errors."""
+
+    def __init__(self, message: str):
+        """Initialize the exception.
+
+        Args:
+            message: The error message from the API
+        """
+        super().__init__(message, specific_error="API error")
+
+
+# ============================================================================
+# Generic Errors
+# ============================================================================
+
+
+class GenericAPIStatusException(ErrorNotPickedUpBySentry):
+    """Raised for generic API status errors that don't fit other categories."""
+
+    def __init__(self, message: str):
+        """Initialize the exception.
+
+        Args:
+            message: The error message from the API
+        """
+        self.api_message = message
+        super().__init__(message)
+
+    def to_markdown(self) -> str:
+        """Generate markdown-formatted error message for TUI."""
+        return f"⚠️ AI service error: {self.api_message}"
+
+    def to_plain_text(self) -> str:
+        """Generate plain text error message for CLI."""
+        return f"⚠️  AI service error: {self.api_message}"
+
+
+class UnknownAgentException(ErrorNotPickedUpBySentry):
+    """Raised for unknown/unclassified agent errors."""
+
+    def __init__(self, original_exception: Exception):
+        """Initialize the exception.
+
+        Args:
+            original_exception: The original exception that was caught
+        """
+        self.original_exception = original_exception
+        super().__init__(str(original_exception))
+
+    def to_markdown(self) -> str:
+        """Generate markdown-formatted error message for TUI."""
+        log_path = get_shotgun_home() / "logs" / "shotgun.log"
+        return f"⚠️ An error occurred: {str(self.original_exception)}\n\nCheck logs at {log_path}"
+
+    def to_plain_text(self) -> str:
+        """Generate plain text error message for CLI."""
+        log_path = get_shotgun_home() / "logs" / "shotgun.log"
+        return f"⚠️  An error occurred: {str(self.original_exception)}\n\nCheck logs at {log_path}"

shotgun/llm_proxy/__init__.py

@@ -1,5 +1,6 @@
 """LiteLLM proxy client utilities and configuration."""
 
+from .client import LiteLLMProxyClient, get_budget_info
 from .clients import (
     create_anthropic_proxy_provider,
     create_litellm_provider,
@@ -9,6 +10,14 @@ from .constants import (
     LITELLM_PROXY_BASE_URL,
     LITELLM_PROXY_OPENAI_BASE,
 )
+from .models import (
+    BudgetInfo,
+    BudgetSource,
+    KeyInfoData,
+    KeyInfoResponse,
+    TeamInfoData,
+    TeamInfoResponse,
+)
 
 __all__ = [
     "LITELLM_PROXY_BASE_URL",
@@ -16,4 +25,12 @@ __all__ = [
     "LITELLM_PROXY_OPENAI_BASE",
     "create_litellm_provider",
     "create_anthropic_proxy_provider",
+    "LiteLLMProxyClient",
+    "get_budget_info",
+    "BudgetInfo",
+    "BudgetSource",
+    "KeyInfoData",
+    "KeyInfoResponse",
+    "TeamInfoData",
+    "TeamInfoResponse",
 ]

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()