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

shotgun/codebase/core/ingestor.py

@@ -6,6 +6,7 @@ import os
 import time
 import uuid
 from collections import defaultdict
+from collections.abc import Callable
 from pathlib import Path
 from typing import Any
 
@@ -198,11 +199,21 @@ class Ingestor:
                     return True
         return False
 
-    def flush_nodes(self) -> None:
-        """Flush pending node insertions to the database."""
+    def flush_nodes(
+        self,
+        progress_callback: Callable[[int, int], None] | None = None,
+    ) -> None:
+        """Flush pending node insertions to the database.
+
+        Args:
+            progress_callback: Optional callback(current, total) for progress reporting
+        """
         if not self.node_buffer:
             return
 
+        total_nodes = len(self.node_buffer)
+        processed = 0
+
         # Group nodes by label
         nodes_by_label: dict[str, list[dict[str, Any]]] = defaultdict(list)
         for label, properties in self.node_buffer:
@@ -239,9 +250,18 @@ class Ingestor:
                     params = dict(zip(prop_names, prop_values, strict=False))
                     self.conn.execute(query, params)
 
+                    # Report progress
+                    processed += 1
+                    if progress_callback and processed % 10 == 0:
+                        progress_callback(processed, total_nodes)
+
             except Exception as e:
                 logger.error(f"Failed to insert {label} nodes: {e}")
 
+        # Final progress report
+        if progress_callback:
+            progress_callback(total_nodes, total_nodes)
+
         # Log node counts by type
         node_type_counts: dict[str, int] = {}
         for label, _ in self.node_buffer:
@@ -280,11 +300,21 @@ class Ingestor:
 
         # Don't auto-flush relationships - wait for explicit flush_all() to ensure nodes exist first
 
-    def flush_relationships(self) -> None:
-        """Flush pending relationship insertions to the database."""
+    def flush_relationships(
+        self,
+        progress_callback: Callable[[int, int], None] | None = None,
+    ) -> None:
+        """Flush pending relationship insertions to the database.
+
+        Args:
+            progress_callback: Optional callback(current, total) for progress reporting
+        """
         if not self.relationship_buffer:
             return
 
+        total_rels = len(self.relationship_buffer)
+        processed = 0
+
         # Group relationships by type
         rels_by_type: dict[
             str, list[tuple[str, str, Any, str, str, str, Any, dict[str, Any] | None]]
@@ -299,7 +329,7 @@ class Ingestor:
                 to_label,
                 to_key,
                 to_value,
-                properties,
+                _properties,
             ) = rel_data
 
             # Determine actual table name
@@ -323,7 +353,7 @@ class Ingestor:
                         to_label,
                         to_key,
                         to_value,
-                        properties,
+                        _properties,
                     ) = rel_data
 
                     # Build MATCH and MERGE query (use MERGE to avoid duplicate relationships)
@@ -337,6 +367,11 @@ class Ingestor:
                     try:
                         self.conn.execute(query, params)
                         success_count += 1
+
+                        # Report progress
+                        processed += 1
+                        if progress_callback and processed % 10 == 0:
+                            progress_callback(processed, total_rels)
                     except Exception as e:
                         logger.error(
                             f"Failed to create single relationship {table_name}: {from_label}({from_value}) -> {to_label}({to_value})"
@@ -360,6 +395,10 @@ class Ingestor:
                 # Don't swallow the exception - let it propagate
                 raise
 
+        # Final progress report
+        if progress_callback:
+            progress_callback(total_rels, total_rels)
+
         # Log summary of flushed relationships
         logger.info(
             f"Flushed {len(self.relationship_buffer)} relationships: {relationship_counts}"
@@ -586,6 +625,9 @@ class SimpleGraphBuilder:
             self.ignore_dirs = self.ignore_dirs.union(set(exclude_patterns))
         self.progress_callback = progress_callback
 
+        # Generate unique session ID for correlating timing events in PostHog
+        self._index_session_id = str(uuid.uuid4())[:8]
+
         # Caches
         self.structural_elements: dict[Path, str | None] = {}
         self.ast_cache: dict[Path, tuple[Node, str]] = {}
@@ -621,25 +663,129 @@ class SimpleGraphBuilder:
             # Don't let progress callback errors crash the build
             logger.debug(f"Progress callback error: {e}")
 
+    def _log_timing(
+        self,
+        phase: str,
+        duration: float,
+        items: int,
+        extra_props: dict[str, Any] | None = None,
+    ) -> None:
+        """Log timing data to PostHog for analysis."""
+        from shotgun.posthog_telemetry import track_event
+
+        properties: dict[str, Any] = {
+            "session_id": self._index_session_id,
+            "phase": phase,
+            "duration_seconds": round(duration, 3),
+            "item_count": items,
+        }
+        if extra_props:
+            properties.update(extra_props)
+
+        track_event("codebase_index_phase_completed", properties)
+
+    def _log_summary(
+        self,
+        total_duration: float,
+        total_files: int,
+        total_nodes: int,
+        total_relationships: int,
+    ) -> None:
+        """Log indexing summary event to PostHog."""
+        from shotgun.posthog_telemetry import track_event
+
+        track_event(
+            "codebase_index_completed",
+            {
+                "session_id": self._index_session_id,
+                "total_duration_seconds": round(total_duration, 3),
+                "total_files": total_files,
+                "total_nodes": total_nodes,
+                "total_relationships": total_relationships,
+            },
+        )
+
     async def run(self) -> None:
         """Run the three-pass graph building process."""
         logger.info(f"Building graph for project: {self.project_name}")
 
         # Pass 1: Structure
         logger.info("Pass 1: Identifying packages and folders...")
+        t0 = time.time()
         self._identify_structure()
+        t1 = time.time()
+        self._log_timing("structure", t1 - t0, len(self.structural_elements))
 
         # Pass 2: Definitions
         logger.info("Pass 2: Processing files and extracting definitions...")
+        t2 = time.time()
         await self._process_files()
+        t3 = time.time()
+        self._log_timing(
+            "definitions",
+            t3 - t2,
+            len(self.ast_cache),
+            {"file_count": len(self.ast_cache)},
+        )
 
         # Pass 3: Relationships
         logger.info("Pass 3: Processing relationships (calls, imports)...")
+        t4 = time.time()
         self._process_relationships()
+        t5 = time.time()
+        self._log_timing("relationships", t5 - t4, len(self.ast_cache))
 
         # Flush all pending operations
         logger.info("Flushing all data to database...")
-        self.ingestor.flush_all()
+        t6 = time.time()
+        node_count = len(self.ingestor.node_buffer)
+
+        # Create progress callback for flush_nodes
+        def node_progress(current: int, total: int) -> None:
+            self._report_progress(
+                "flush_nodes", "Flushing nodes to database", current, total
+            )
+
+        self.ingestor.flush_nodes(progress_callback=node_progress)
+        self._report_progress(
+            "flush_nodes", "Flushing nodes to database", node_count, node_count, True
+        )
+        t7 = time.time()
+        self._log_timing("flush_nodes", t7 - t6, node_count, {"node_count": node_count})
+
+        rel_count = len(self.ingestor.relationship_buffer)
+
+        # Create progress callback for flush_relationships
+        def rel_progress(current: int, total: int) -> None:
+            self._report_progress(
+                "flush_relationships",
+                "Flushing relationships to database",
+                current,
+                total,
+            )
+
+        self.ingestor.flush_relationships(progress_callback=rel_progress)
+        self._report_progress(
+            "flush_relationships",
+            "Flushing relationships to database",
+            rel_count,
+            rel_count,
+            True,
+        )
+        t8 = time.time()
+        self._log_timing(
+            "flush_relationships", t8 - t7, rel_count, {"relationship_count": rel_count}
+        )
+
+        # Track summary event with totals (no PII - only numeric metadata)
+        total_duration = t8 - t0
+        self._log_summary(
+            total_duration=total_duration,
+            total_files=len(self.ast_cache),
+            total_nodes=node_count,
+            total_relationships=rel_count,
+        )
+
         logger.info("Graph building complete!")
 
     def _identify_structure(self) -> None:

shotgun/codebase/models.py

@@ -29,6 +29,8 @@ class ProgressPhase(StrEnum):
     STRUCTURE = "structure"  # Identifying packages and folders
     DEFINITIONS = "definitions"  # Processing files and extracting definitions
     RELATIONSHIPS = "relationships"  # Processing relationships (calls, imports)
+    FLUSH_NODES = "flush_nodes"  # Flushing nodes to database
+    FLUSH_RELATIONSHIPS = "flush_relationships"  # Flushing relationships to database
 
 
 class IndexProgress(BaseModel):

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",
 ]