shotgun-sh 0.1.14__py3-none-any.whl → 0.2.11__py3-none-any.whl

shotgun/agents/history/context_extraction.py

@@ -1,5 +1,9 @@
 """Context extraction utilities for history processing."""
 
+import json
+import logging
+import traceback
+
 from pydantic_ai.messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
@@ -16,6 +20,46 @@ from pydantic_ai.messages import (
     UserPromptPart,
 )
 
+logger = logging.getLogger(__name__)
+
+
+def _safely_parse_tool_args(args: dict[str, object] | str | None) -> dict[str, object]:
+    """Safely parse tool call arguments, handling incomplete/invalid JSON.
+
+    Args:
+        args: Tool call arguments (dict, JSON string, or None)
+
+    Returns:
+        Parsed args dict, or empty dict if parsing fails
+    """
+    if args is None:
+        return {}
+
+    if isinstance(args, dict):
+        return args
+
+    if not isinstance(args, str):
+        return {}
+
+    try:
+        parsed = json.loads(args)
+        return parsed if isinstance(parsed, dict) else {}
+    except (json.JSONDecodeError, ValueError) as e:
+        # Only log warning if it looks like JSON (starts with { or [) - incomplete JSON
+        # Plain strings are valid args and shouldn't trigger warnings
+        stripped_args = args.strip()
+        if stripped_args.startswith(("{", "[")):
+            args_preview = args[:100] + "..." if len(args) > 100 else args
+            logger.warning(
+                "Detected incomplete/invalid JSON in tool call args during parsing",
+                extra={
+                    "args_preview": args_preview,
+                    "error": str(e),
+                    "args_length": len(args),
+                },
+            )
+        return {}
+
 
 def extract_context_from_messages(messages: list[ModelMessage]) -> str:
     """Extract context from a list of messages for summarization."""
@@ -87,12 +131,55 @@ def extract_context_from_part(
         return f"<ASSISTANT_TEXT>\n{message_part.content}\n</ASSISTANT_TEXT>"
 
     elif isinstance(message_part, ToolCallPart):
-        if isinstance(message_part.args, dict):
-            args_str = ", ".join(f"{k}={repr(v)}" for k, v in message_part.args.items())
-            tool_call_str = f"{message_part.tool_name}({args_str})"
-        else:
-            tool_call_str = f"{message_part.tool_name}({message_part.args})"
-        return f"<TOOL_CALL>\n{tool_call_str}\n</TOOL_CALL>"
+        # Safely parse args to avoid crashes from incomplete JSON during streaming
+        try:
+            parsed_args = _safely_parse_tool_args(message_part.args)
+            if parsed_args:
+                # Successfully parsed as dict - format nicely
+                args_str = ", ".join(f"{k}={repr(v)}" for k, v in parsed_args.items())
+                tool_call_str = f"{message_part.tool_name}({args_str})"
+            elif isinstance(message_part.args, str) and message_part.args:
+                # Non-empty string that didn't parse as JSON
+                # Check if it looks like JSON (starts with { or [) - if so, it's incomplete
+                stripped_args = message_part.args.strip()
+                if stripped_args.startswith(("{", "[")):
+                    # Looks like incomplete JSON - log warning and show empty parens
+                    args_preview = (
+                        stripped_args[:100] + "..."
+                        if len(stripped_args) > 100
+                        else stripped_args
+                    )
+                    stack_trace = "".join(traceback.format_stack())
+                    logger.warning(
+                        "ToolCallPart with unparseable args encountered during context extraction",
+                        extra={
+                            "tool_name": message_part.tool_name,
+                            "tool_call_id": message_part.tool_call_id,
+                            "args_preview": args_preview,
+                            "args_type": type(message_part.args).__name__,
+                            "stack_trace": stack_trace,
+                        },
+                    )
+                    tool_call_str = f"{message_part.tool_name}()"
+                else:
+                    # Plain string arg - display as-is
+                    tool_call_str = f"{message_part.tool_name}({message_part.args})"
+            else:
+                # No args
+                tool_call_str = f"{message_part.tool_name}()"
+            return f"<TOOL_CALL>\n{tool_call_str}\n</TOOL_CALL>"
+        except Exception as e:  # pragma: no cover - defensive catch-all
+            # If anything goes wrong, log full exception with stack trace
+            logger.error(
+                "Unexpected error processing ToolCallPart",
+                exc_info=True,
+                extra={
+                    "tool_name": message_part.tool_name,
+                    "tool_call_id": message_part.tool_call_id,
+                    "error": str(e),
+                },
+            )
+            return f"<TOOL_CALL>\n{message_part.tool_name}()\n</TOOL_CALL>"
 
     elif isinstance(message_part, BuiltinToolCallPart):
         return f"<BUILTIN_TOOL_CALL>\n{message_part.tool_name}\n</BUILTIN_TOOL_CALL>"

shotgun/agents/history/history_processors.py

@@ -1,7 +1,10 @@
 """History processors for managing conversation history in Shotgun agents."""
 
+from collections.abc import Awaitable, Callable
 from typing import TYPE_CHECKING, Any, Protocol
 
+from anthropic import APIStatusError
+from pydantic_ai import ModelSettings
 from pydantic_ai.messages import (
     ModelMessage,
     ModelRequest,
@@ -10,9 +13,10 @@ from pydantic_ai.messages import (
     UserPromptPart,
 )
 
-from shotgun.agents.config.models import shotgun_model_request
+from shotgun.agents.llm import shotgun_model_request
 from shotgun.agents.messages import AgentSystemPrompt, SystemStatusPrompt
 from shotgun.agents.models import AgentDeps
+from shotgun.exceptions import ContextSizeLimitExceeded
 from shotgun.logging_config import get_logger
 from shotgun.posthog_telemetry import track_event
 from shotgun.prompts import PromptLoader
@@ -50,6 +54,86 @@ logger = get_logger(__name__)
 prompt_loader = PromptLoader()
 
 
+async def _safe_token_estimation(
+    estimation_func: Callable[..., Awaitable[int]],
+    model_name: str,
+    max_tokens: int,
+    *args: Any,
+    **kwargs: Any,
+) -> int:
+    """Safely estimate tokens with proper error handling.
+
+    Wraps token estimation functions to handle failures gracefully.
+    Only RuntimeError (from token counters) is wrapped in ContextSizeLimitExceeded.
+    Other errors (network, auth) are allowed to bubble up.
+
+    Args:
+        estimation_func: Async function that estimates tokens
+        model_name: Name of the model for error messages
+        max_tokens: Maximum tokens for the model
+        *args: Arguments to pass to estimation_func
+        **kwargs: Keyword arguments to pass to estimation_func
+
+    Returns:
+        Token count from estimation_func
+
+    Raises:
+        ContextSizeLimitExceeded: If token counting fails with RuntimeError
+        Exception: Any other exceptions from estimation_func
+    """
+    try:
+        return await estimation_func(*args, **kwargs)
+    except Exception as e:
+        # Log the error with full context
+        logger.warning(
+            f"Token counting failed for {model_name}",
+            extra={
+                "error_type": type(e).__name__,
+                "error_message": str(e),
+                "model": model_name,
+            },
+        )
+
+        # Token counting behavior with oversized context (verified via testing):
+        #
+        # 1. OpenAI/tiktoken:
+        #    - Successfully counts any size (tested with 752K tokens, no error)
+        #    - Library errors: ValueError, KeyError, AttributeError, SSLError (file/cache issues)
+        #    - Wrapped as: RuntimeError by our counter
+        #
+        # 2. Gemini/SentencePiece:
+        #    - Successfully counts any size (tested with 752K tokens, no error)
+        #    - Library errors: RuntimeError, IOError, TypeError (file/model loading issues)
+        #    - Wrapped as: RuntimeError by our counter
+        #
+        # 3. Anthropic API:
+        #    - Successfully counts large token counts (tested with 752K tokens, no error)
+        #    - Only enforces 32 MB request size limit (not token count)
+        #    - Raises: APIStatusError(413) with error type 'request_too_large' for 32MB+ requests
+        #    - Other API errors: APIConnectionError, RateLimitError, APIStatusError (4xx/5xx)
+        #    - Wrapped as: RuntimeError by our counter
+        #
+        # IMPORTANT: No provider raises errors for "too many tokens" during counting.
+        # Token count validation happens separately by comparing count to max_input_tokens.
+        #
+        # We wrap RuntimeError (library-level failures from tiktoken/sentencepiece).
+        # We also wrap Anthropic's 413 error (request exceeds 32 MB) as it indicates
+        # context is effectively too large and needs user action to reduce it.
+        if isinstance(e, RuntimeError):
+            raise ContextSizeLimitExceeded(
+                model_name=model_name, max_tokens=max_tokens
+            ) from e
+
+        # Check for Anthropic's 32 MB request size limit (APIStatusError with status 413)
+        if isinstance(e, APIStatusError) and e.status_code == 413:
+            raise ContextSizeLimitExceeded(
+                model_name=model_name, max_tokens=max_tokens
+            ) from e
+
+        # Re-raise other exceptions (network errors, auth failures, etc.)
+        raise
+
+
 def is_summary_part(part: Any) -> bool:
     """Check if a message part is a compacted summary."""
     return isinstance(part, TextPart) and part.content.startswith(SUMMARY_MARKER)
@@ -126,6 +210,7 @@ calculate_max_summarization_tokens = _calculate_max_summarization_tokens
 async def token_limit_compactor(
     ctx: ContextProtocol,
     messages: list[ModelMessage],
+    force: bool = False,
 ) -> list[ModelMessage]:
     """Compact message history based on token limits with incremental processing.
 
@@ -138,6 +223,7 @@ async def token_limit_compactor(
     Args:
         ctx: Run context with usage information and dependencies
         messages: Current conversation history
+        force: If True, force compaction even if below token threshold
 
     Returns:
         Compacted list of messages within token limits
@@ -154,9 +240,15 @@ async def token_limit_compactor(
 
     if last_summary_index is not None:
         # Check if post-summary conversation exceeds threshold for incremental compaction
-        post_summary_tokens = estimate_post_summary_tokens(
-            messages, last_summary_index, deps.llm_model
+        post_summary_tokens = await _safe_token_estimation(
+            estimate_post_summary_tokens,
+            deps.llm_model.name,
+            model_max_tokens,
+            messages,
+            last_summary_index,
+            deps.llm_model,
         )
+
         post_summary_percentage = (
             (post_summary_tokens / max_tokens) * 100 if max_tokens > 0 else 0
         )
@@ -168,7 +260,7 @@ async def token_limit_compactor(
         )
 
         # Only do incremental compaction if post-summary conversation exceeds threshold
-        if post_summary_tokens < max_tokens:
+        if post_summary_tokens < max_tokens and not force:
             logger.debug(
                 f"Post-summary conversation under threshold ({post_summary_tokens} < {max_tokens}), "
                 f"keeping all {len(messages)} messages"
@@ -248,7 +340,7 @@ async def token_limit_compactor(
         ]
 
         # Calculate optimal max_tokens for summarization
-        max_tokens = calculate_max_summarization_tokens(
+        max_tokens = await calculate_max_summarization_tokens(
             deps.llm_model, request_messages
         )
 
@@ -261,7 +353,9 @@ async def token_limit_compactor(
         summary_response = await shotgun_model_request(
             model_config=deps.llm_model,
             messages=request_messages,
-            max_tokens=max_tokens,  # Use calculated optimal tokens for summarization
+            model_settings=ModelSettings(
+                max_tokens=max_tokens  # Use calculated optimal tokens for summarization
+            ),
         )
 
         log_summarization_response(summary_response, "INCREMENTAL")
@@ -328,13 +422,16 @@ async def token_limit_compactor(
 
         # Track compaction completion
         messages_after = len(compacted_messages)
-        tokens_after = estimate_tokens_from_messages(compacted_messages, deps.llm_model)
+        tokens_after = await estimate_tokens_from_messages(
+            compacted_messages, deps.llm_model
+        )
         reduction_percentage = (
             ((messages_before - messages_after) / messages_before * 100)
             if messages_before > 0
             else 0
         )
 
+        # Track incremental compaction with simple metrics (fast, no token counting)
         track_event(
             "context_compaction_triggered",
             {
@@ -347,6 +444,10 @@ async def token_limit_compactor(
                 "agent_mode": deps.agent_mode.value
                 if hasattr(deps, "agent_mode") and deps.agent_mode
                 else "unknown",
+                # Model and provider info (no computation needed)
+                "model_name": deps.llm_model.name.value,
+                "provider": deps.llm_model.provider.value,
+                "key_provider": deps.llm_model.key_provider.value,
             },
         )
 
@@ -354,7 +455,14 @@ async def token_limit_compactor(
 
     else:
         # Check if total conversation exceeds threshold for full compaction
-        total_tokens = estimate_tokens_from_messages(messages, deps.llm_model)
+        total_tokens = await _safe_token_estimation(
+            estimate_tokens_from_messages,
+            deps.llm_model.name,
+            model_max_tokens,
+            messages,
+            deps.llm_model,
+        )
+
         total_percentage = (total_tokens / max_tokens) * 100 if max_tokens > 0 else 0
 
         logger.debug(
@@ -363,7 +471,7 @@ async def token_limit_compactor(
         )
 
         # Only do full compaction if total conversation exceeds threshold
-        if total_tokens < max_tokens:
+        if total_tokens < max_tokens and not force:
             logger.debug(
                 f"Total conversation under threshold ({total_tokens} < {max_tokens}), "
                 f"keeping all {len(messages)} messages"
@@ -392,7 +500,9 @@ async def _full_compaction(
     ]
 
     # Calculate optimal max_tokens for summarization
-    max_tokens = calculate_max_summarization_tokens(deps.llm_model, request_messages)
+    max_tokens = await calculate_max_summarization_tokens(
+        deps.llm_model, request_messages
+    )
 
     # Debug logging using shared utilities
     log_summarization_request(
@@ -403,11 +513,13 @@ async def _full_compaction(
     summary_response = await shotgun_model_request(
         model_config=deps.llm_model,
         messages=request_messages,
-        max_tokens=max_tokens,  # Use calculated optimal tokens for summarization
+        model_settings=ModelSettings(
+            max_tokens=max_tokens  # Use calculated optimal tokens for summarization
+        ),
     )
 
     # Calculate token reduction
-    current_tokens = estimate_tokens_from_messages(messages, deps.llm_model)
+    current_tokens = await estimate_tokens_from_messages(messages, deps.llm_model)
     summary_usage = summary_response.usage
     reduction_percentage = (
         ((current_tokens - summary_usage.output_tokens) / current_tokens) * 100
@@ -459,6 +571,7 @@ async def _full_compaction(
     tokens_before = current_tokens  # Already calculated above
     tokens_after = summary_usage.output_tokens if summary_usage else 0
 
+    # Track full compaction with simple metrics (fast, no token counting)
     track_event(
         "context_compaction_triggered",
         {
@@ -471,6 +584,10 @@ async def _full_compaction(
             "agent_mode": deps.agent_mode.value
             if hasattr(deps, "agent_mode") and deps.agent_mode
             else "unknown",
+            # Model and provider info (no computation needed)
+            "model_name": deps.llm_model.name.value,
+            "provider": deps.llm_model.provider.value,
+            "key_provider": deps.llm_model.key_provider.value,
         },
     )
 

shotgun/agents/history/token_counting/__init__.py

@@ -0,0 +1,31 @@
+"""Real token counting for all supported providers.
+
+This module provides accurate token counting using each provider's official
+APIs and libraries, eliminating the need for rough character-based estimation.
+"""
+
+from .anthropic import AnthropicTokenCounter
+from .base import TokenCounter, extract_text_from_messages
+from .openai import OpenAITokenCounter
+from .sentencepiece_counter import SentencePieceTokenCounter
+from .utils import (
+    count_post_summary_tokens,
+    count_tokens_from_message_parts,
+    count_tokens_from_messages,
+    get_token_counter,
+)
+
+__all__ = [
+    # Base classes
+    "TokenCounter",
+    # Counter implementations
+    "OpenAITokenCounter",
+    "AnthropicTokenCounter",
+    "SentencePieceTokenCounter",
+    # Utility functions
+    "get_token_counter",
+    "count_tokens_from_messages",
+    "count_post_summary_tokens",
+    "count_tokens_from_message_parts",
+    "extract_text_from_messages",
+]

shotgun/agents/history/token_counting/anthropic.py

@@ -0,0 +1,127 @@
+"""Anthropic token counting using official client."""
+
+import logfire
+from pydantic_ai.messages import ModelMessage
+
+from shotgun.agents.config.models import KeyProvider
+from shotgun.llm_proxy import create_anthropic_proxy_provider
+from shotgun.logging_config import get_logger
+
+from .base import TokenCounter, extract_text_from_messages
+
+logger = get_logger(__name__)
+
+
+class AnthropicTokenCounter(TokenCounter):
+    """Token counter for Anthropic models using official client."""
+
+    def __init__(
+        self,
+        model_name: str,
+        api_key: str,
+        key_provider: KeyProvider = KeyProvider.BYOK,
+    ):
+        """Initialize Anthropic token counter.
+
+        Args:
+            model_name: Anthropic model name for token counting
+            api_key: API key (Anthropic for BYOK, Shotgun for proxy)
+            key_provider: Key provider type (BYOK or SHOTGUN)
+
+        Raises:
+            RuntimeError: If client initialization fails
+        """
+        self.model_name = model_name
+        import anthropic
+
+        try:
+            if key_provider == KeyProvider.SHOTGUN:
+                # Use LiteLLM proxy for Shotgun Account
+                # Get async client from AnthropicProvider
+                provider = create_anthropic_proxy_provider(api_key)
+                self.client = provider.client
+                logger.debug(
+                    f"Initialized async Anthropic token counter for {model_name} via LiteLLM proxy"
+                )
+            else:
+                # Direct Anthropic API for BYOK - use async client
+                self.client = anthropic.AsyncAnthropic(api_key=api_key)
+                logger.debug(
+                    f"Initialized async Anthropic token counter for {model_name} via direct API"
+                )
+        except Exception as e:
+            logfire.exception(
+                f"Failed to initialize Anthropic token counter for {model_name}",
+                model_name=model_name,
+                key_provider=key_provider.value,
+                exception_type=type(e).__name__,
+            )
+            raise RuntimeError(
+                f"Failed to initialize Anthropic async client for {model_name}: {type(e).__name__}: {str(e)}"
+            ) from e
+
+    async def count_tokens(self, text: str) -> int:
+        """Count tokens using Anthropic's official API (async).
+
+        Args:
+            text: Text to count tokens for
+
+        Returns:
+            Exact token count from Anthropic API
+
+        Raises:
+            RuntimeError: If API call fails
+        """
+        # Handle empty text to avoid unnecessary API calls
+        # Anthropic API requires non-empty content, so we need a strict check
+        if not text or not text.strip():
+            return 0
+
+        # Additional validation: ensure the text has actual content
+        # Some edge cases might have only whitespace or control characters
+        cleaned_text = text.strip()
+        if not cleaned_text:
+            return 0
+
+        try:
+            # Anthropic API expects messages format and model parameter
+            # Use await with async client
+            result = await self.client.messages.count_tokens(
+                messages=[{"role": "user", "content": cleaned_text}],
+                model=self.model_name,
+            )
+            return result.input_tokens
+        except Exception as e:
+            # Create a preview of the text for logging (truncated to avoid huge logs)
+            text_preview = text[:100] + "..." if len(text) > 100 else text
+
+            logfire.exception(
+                f"Anthropic token counting failed for {self.model_name}",
+                model_name=self.model_name,
+                text_length=len(text),
+                text_preview=text_preview,
+                exception_type=type(e).__name__,
+                exception_message=str(e),
+            )
+            raise RuntimeError(
+                f"Anthropic token counting API failed for {self.model_name}: {type(e).__name__}: {str(e)}"
+            ) from e
+
+    async def count_message_tokens(self, messages: list[ModelMessage]) -> int:
+        """Count tokens across all messages using Anthropic API (async).
+
+        Args:
+            messages: List of PydanticAI messages
+
+        Returns:
+            Total token count for all messages
+
+        Raises:
+            RuntimeError: If token counting fails
+        """
+        # Handle empty message list early
+        if not messages:
+            return 0
+
+        total_text = extract_text_from_messages(messages)
+        return await self.count_tokens(total_text)

shotgun/agents/history/token_counting/base.py

@@ -0,0 +1,78 @@
+"""Base classes and shared utilities for token counting."""
+
+from abc import ABC, abstractmethod
+
+from pydantic_ai.messages import ModelMessage
+
+
+class TokenCounter(ABC):
+    """Abstract base class for provider-specific token counting.
+
+    All methods are async to support non-blocking operations like
+    downloading tokenizer models or making API calls.
+    """
+
+    @abstractmethod
+    async def count_tokens(self, text: str) -> int:
+        """Count tokens in text using provider-specific method (async).
+
+        Args:
+            text: Text to count tokens for
+
+        Returns:
+            Exact token count as determined by the provider
+
+        Raises:
+            RuntimeError: If token counting fails
+        """
+
+    @abstractmethod
+    async def count_message_tokens(self, messages: list[ModelMessage]) -> int:
+        """Count tokens in PydanticAI message structures (async).
+
+        Args:
+            messages: List of messages to count tokens for
+
+        Returns:
+            Total token count across all messages
+
+        Raises:
+            RuntimeError: If token counting fails
+        """
+
+
+def extract_text_from_messages(messages: list[ModelMessage]) -> str:
+    """Extract all text content from messages for token counting.
+
+    Args:
+        messages: List of PydanticAI messages
+
+    Returns:
+        Combined text content from all messages
+    """
+    text_parts = []
+
+    for message in messages:
+        if hasattr(message, "parts"):
+            for part in message.parts:
+                if hasattr(part, "content") and isinstance(part.content, str):
+                    # Only add non-empty content
+                    if part.content.strip():
+                        text_parts.append(part.content)
+                else:
+                    # Handle non-text parts (tool calls, etc.)
+                    part_str = str(part)
+                    if part_str.strip():
+                        text_parts.append(part_str)
+        else:
+            # Handle messages without parts
+            msg_str = str(message)
+            if msg_str.strip():
+                text_parts.append(msg_str)
+
+    # If no valid text parts found, return a minimal placeholder
+    # This ensures we never send completely empty content to APIs
+    if not text_parts:
+        return "."
+
+    return "\n".join(text_parts)