shotgun-sh 0.2.8.dev2__py3-none-any.whl → 0.3.3.dev1__py3-none-any.whl

shotgun/agents/context_analyzer/models.py

@@ -0,0 +1,212 @@
+"""Pydantic models for context analysis."""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class TokenAllocation(BaseModel):
+    """Token counts allocated from API usage data by message/tool type.
+
+    Used internally by ContextAnalyzer to track token distribution across
+    different message types and tool categories.
+    """
+
+    user: int = Field(ge=0, default=0, description="Tokens from user prompts")
+    agent_responses: int = Field(
+        ge=0, default=0, description="Tokens from agent text responses"
+    )
+    system_prompts: int = Field(
+        ge=0, default=0, description="Tokens from system prompts"
+    )
+    system_status: int = Field(
+        ge=0, default=0, description="Tokens from system status messages"
+    )
+    codebase_understanding: int = Field(
+        ge=0, default=0, description="Tokens from codebase understanding tools"
+    )
+    artifact_management: int = Field(
+        ge=0, default=0, description="Tokens from artifact management tools"
+    )
+    web_research: int = Field(
+        ge=0, default=0, description="Tokens from web research tools"
+    )
+    unknown: int = Field(ge=0, default=0, description="Tokens from uncategorized tools")
+
+
+class MessageTypeStats(BaseModel):
+    """Statistics for a specific message type."""
+
+    count: int = Field(ge=0, description="Number of messages of this type")
+    tokens: int = Field(ge=0, description="Total tokens consumed by this type")
+
+    @property
+    def avg_tokens(self) -> float:
+        """Calculate average tokens per message."""
+        return self.tokens / self.count if self.count > 0 else 0.0
+
+
+class ContextAnalysis(BaseModel):
+    """Complete analysis of conversation context composition."""
+
+    user_messages: MessageTypeStats
+    agent_responses: MessageTypeStats
+    system_prompts: MessageTypeStats
+    system_status: MessageTypeStats
+    codebase_understanding: MessageTypeStats
+    artifact_management: MessageTypeStats
+    web_research: MessageTypeStats
+    unknown: MessageTypeStats
+    hint_messages: MessageTypeStats
+    total_tokens: int = Field(ge=0, description="Total tokens including hints")
+    total_messages: int = Field(ge=0, description="Total message count including hints")
+    context_window: int = Field(ge=0, description="Model's maximum input tokens")
+    agent_context_tokens: int = Field(
+        ge=0,
+        description="Tokens that actually consume agent context (excluding UI-only)",
+    )
+    model_name: str = Field(description="Name of the model being used")
+    max_usable_tokens: int = Field(
+        ge=0, description="80% of max_input_tokens (usable limit)"
+    )
+    free_space_tokens: int = Field(
+        description="Remaining tokens available (negative if over capacity)"
+    )
+
+    def get_percentage(self, stats: MessageTypeStats) -> float:
+        """Calculate percentage of agent context tokens for a message type.
+
+        Args:
+            stats: Message type statistics to calculate percentage for
+
+        Returns:
+            Percentage of total agent context tokens (0-100)
+        """
+        return (
+            (stats.tokens / self.agent_context_tokens * 100)
+            if self.agent_context_tokens > 0
+            else 0.0
+        )
+
+
+class ContextCompositionTelemetry(BaseModel):
+    """Telemetry data for context composition tracking to PostHog."""
+
+    # Context usage
+    total_messages: int = Field(ge=0)
+    agent_context_tokens: int = Field(ge=0)
+    context_window: int = Field(ge=0)
+    max_usable_tokens: int = Field(ge=0)
+    free_space_tokens: int = Field(ge=0)
+    usage_percentage: float = Field(ge=0, le=100)
+
+    # Message type counts
+    user_messages_count: int = Field(ge=0)
+    agent_responses_count: int = Field(ge=0)
+    system_prompts_count: int = Field(ge=0)
+    system_status_count: int = Field(ge=0)
+    codebase_understanding_count: int = Field(ge=0)
+    artifact_management_count: int = Field(ge=0)
+    web_research_count: int = Field(ge=0)
+    unknown_tools_count: int = Field(ge=0)
+
+    # Token distribution percentages
+    user_messages_pct: float = Field(ge=0, le=100)
+    agent_responses_pct: float = Field(ge=0, le=100)
+    system_prompts_pct: float = Field(ge=0, le=100)
+    system_status_pct: float = Field(ge=0, le=100)
+    codebase_understanding_pct: float = Field(ge=0, le=100)
+    artifact_management_pct: float = Field(ge=0, le=100)
+    web_research_pct: float = Field(ge=0, le=100)
+    unknown_tools_pct: float = Field(ge=0, le=100)
+
+    # Compaction info
+    compaction_occurred: bool
+    messages_before_compaction: int | None = None
+    messages_after_compaction: int | None = None
+    compaction_reduction_pct: float | None = None
+
+    @classmethod
+    def from_analysis(
+        cls,
+        analysis: "ContextAnalysis",
+        compaction_occurred: bool = False,
+        messages_before_compaction: int | None = None,
+    ) -> "ContextCompositionTelemetry":
+        """Create telemetry from context analysis.
+
+        Args:
+            analysis: The context analysis to convert
+            compaction_occurred: Whether message compaction occurred
+            messages_before_compaction: Number of messages before compaction
+
+        Returns:
+            ContextCompositionTelemetry instance
+        """
+        total_messages = analysis.total_messages - analysis.hint_messages.count
+        usage_pct = (
+            round((analysis.agent_context_tokens / analysis.max_usable_tokens * 100), 1)
+            if analysis.max_usable_tokens > 0
+            else 0
+        )
+
+        # Calculate compaction metrics
+        messages_after: int | None = None
+        compaction_reduction_pct: float | None = None
+
+        if compaction_occurred and messages_before_compaction is not None:
+            messages_after = total_messages
+            if messages_before_compaction > 0:
+                compaction_reduction_pct = round(
+                    (1 - (total_messages / messages_before_compaction)) * 100, 1
+                )
+
+        return cls(
+            # Context usage
+            total_messages=total_messages,
+            agent_context_tokens=analysis.agent_context_tokens,
+            context_window=analysis.context_window,
+            max_usable_tokens=analysis.max_usable_tokens,
+            free_space_tokens=analysis.free_space_tokens,
+            usage_percentage=usage_pct,
+            # Message type counts
+            user_messages_count=analysis.user_messages.count,
+            agent_responses_count=analysis.agent_responses.count,
+            system_prompts_count=analysis.system_prompts.count,
+            system_status_count=analysis.system_status.count,
+            codebase_understanding_count=analysis.codebase_understanding.count,
+            artifact_management_count=analysis.artifact_management.count,
+            web_research_count=analysis.web_research.count,
+            unknown_tools_count=analysis.unknown.count,
+            # Token distribution percentages
+            user_messages_pct=round(analysis.get_percentage(analysis.user_messages), 1),
+            agent_responses_pct=round(
+                analysis.get_percentage(analysis.agent_responses), 1
+            ),
+            system_prompts_pct=round(
+                analysis.get_percentage(analysis.system_prompts), 1
+            ),
+            system_status_pct=round(analysis.get_percentage(analysis.system_status), 1),
+            codebase_understanding_pct=round(
+                analysis.get_percentage(analysis.codebase_understanding), 1
+            ),
+            artifact_management_pct=round(
+                analysis.get_percentage(analysis.artifact_management), 1
+            ),
+            web_research_pct=round(analysis.get_percentage(analysis.web_research), 1),
+            unknown_tools_pct=round(analysis.get_percentage(analysis.unknown), 1),
+            # Compaction info
+            compaction_occurred=compaction_occurred,
+            messages_before_compaction=messages_before_compaction,
+            messages_after_compaction=messages_after,
+            compaction_reduction_pct=compaction_reduction_pct,
+        )
+
+
+class ContextAnalysisOutput(BaseModel):
+    """Output format for context analysis with multiple representations."""
+
+    markdown: str = Field(description="Markdown-formatted analysis for display")
+    json_data: dict[str, Any] = Field(
+        description="JSON representation of analysis data"
+    )

shotgun/agents/conversation/__init__.py

@@ -0,0 +1,18 @@
+"""Conversation module for managing conversation history and persistence."""
+
+from .filters import (
+    filter_incomplete_messages,
+    filter_orphaned_tool_responses,
+    is_tool_call_complete,
+)
+from .manager import ConversationManager
+from .models import ConversationHistory, ConversationState
+
+__all__ = [
+    "ConversationHistory",
+    "ConversationManager",
+    "ConversationState",
+    "filter_incomplete_messages",
+    "filter_orphaned_tool_responses",
+    "is_tool_call_complete",
+]

shotgun/agents/conversation/filters.py

@@ -0,0 +1,164 @@
+"""Filter functions for conversation message validation."""
+
+import json
+import logging
+
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelRequestPart,
+    ModelResponse,
+    ToolCallPart,
+    ToolReturnPart,
+)
+
+logger = logging.getLogger(__name__)
+
+
+def is_tool_call_complete(tool_call: ToolCallPart) -> bool:
+    """Check if a tool call has valid, complete JSON arguments.
+
+    Args:
+        tool_call: The tool call part to validate
+
+    Returns:
+        True if the tool call args are valid JSON, False otherwise
+    """
+    if tool_call.args is None:
+        return True  # No args is valid
+
+    if isinstance(tool_call.args, dict):
+        return True  # Already parsed dict is valid
+
+    if not isinstance(tool_call.args, str):
+        return False
+
+    # Try to parse the JSON string
+    try:
+        json.loads(tool_call.args)
+        return True
+    except (json.JSONDecodeError, ValueError) as e:
+        # Log incomplete tool call detection
+        args_preview = (
+            tool_call.args[:100] + "..."
+            if len(tool_call.args) > 100
+            else tool_call.args
+        )
+        logger.info(
+            "Detected incomplete tool call in validation",
+            extra={
+                "tool_name": tool_call.tool_name,
+                "tool_call_id": tool_call.tool_call_id,
+                "args_preview": args_preview,
+                "error": str(e),
+            },
+        )
+        return False
+
+
+def filter_incomplete_messages(messages: list[ModelMessage]) -> list[ModelMessage]:
+    """Filter out messages with incomplete tool calls.
+
+    Args:
+        messages: List of messages to filter
+
+    Returns:
+        List of messages with only complete tool calls
+    """
+    filtered: list[ModelMessage] = []
+    filtered_count = 0
+    filtered_tool_names: list[str] = []
+
+    for message in messages:
+        # Only check ModelResponse messages for tool calls
+        if not isinstance(message, ModelResponse):
+            filtered.append(message)
+            continue
+
+        # Check if any tool calls are incomplete
+        has_incomplete_tool_call = False
+        for part in message.parts:
+            if isinstance(part, ToolCallPart) and not is_tool_call_complete(part):
+                has_incomplete_tool_call = True
+                filtered_tool_names.append(part.tool_name)
+                break
+
+        # Only include messages without incomplete tool calls
+        if not has_incomplete_tool_call:
+            filtered.append(message)
+        else:
+            filtered_count += 1
+
+    # Log if any messages were filtered
+    if filtered_count > 0:
+        logger.info(
+            "Filtered incomplete messages before saving",
+            extra={
+                "filtered_count": filtered_count,
+                "total_messages": len(messages),
+                "filtered_tool_names": filtered_tool_names,
+            },
+        )
+
+    return filtered
+
+
+def filter_orphaned_tool_responses(messages: list[ModelMessage]) -> list[ModelMessage]:
+    """Filter out tool responses without corresponding tool calls.
+
+    This ensures message history is valid for OpenAI API which requires
+    tool responses to follow their corresponding tool calls.
+
+    Args:
+        messages: List of messages to filter
+
+    Returns:
+        List of messages with orphaned tool responses removed
+    """
+    # Collect all tool_call_ids from ToolCallPart in ModelResponse
+    valid_tool_call_ids: set[str] = set()
+    for msg in messages:
+        if isinstance(msg, ModelResponse):
+            for part in msg.parts:
+                if isinstance(part, ToolCallPart) and part.tool_call_id:
+                    valid_tool_call_ids.add(part.tool_call_id)
+
+    # Filter out orphaned ToolReturnPart from ModelRequest
+    filtered: list[ModelMessage] = []
+    orphaned_count = 0
+    orphaned_tool_names: list[str] = []
+
+    for msg in messages:
+        if isinstance(msg, ModelRequest):
+            # Filter parts, removing orphaned ToolReturnPart
+            filtered_parts: list[ModelRequestPart] = []
+            request_part: ModelRequestPart
+            for request_part in msg.parts:
+                if isinstance(request_part, ToolReturnPart):
+                    if request_part.tool_call_id in valid_tool_call_ids:
+                        filtered_parts.append(request_part)
+                    else:
+                        # Skip orphaned tool response
+                        orphaned_count += 1
+                        orphaned_tool_names.append(request_part.tool_name or "unknown")
+                else:
+                    filtered_parts.append(request_part)
+
+            # Only add if there are remaining parts
+            if filtered_parts:
+                filtered.append(ModelRequest(parts=filtered_parts))
+        else:
+            filtered.append(msg)
+
+    # Log if any tool responses were filtered
+    if orphaned_count > 0:
+        logger.info(
+            "Filtered orphaned tool responses",
+            extra={
+                "orphaned_count": orphaned_count,
+                "total_messages": len(messages),
+                "orphaned_tool_names": orphaned_tool_names,
+            },
+        )
+
+    return filtered

shotgun/agents/conversation/history/chunking.py

@@ -0,0 +1,278 @@
+"""Pattern-based chunking for oversized conversation compaction.
+
+This module provides functions to break oversized conversations into logical
+chunks for summarization, preserving semantic units like tool call sequences.
+"""
+
+import logging
+from dataclasses import dataclass, field
+
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+
+from shotgun.agents.config.models import ModelConfig
+
+from .constants import CHUNK_TARGET_RATIO, RETENTION_WINDOW_MESSAGES
+from .token_estimation import estimate_tokens_from_messages
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MessageGroup:
+    """A logical group of messages that must stay together.
+
+    Examples:
+    - A single user message
+    - A tool call sequence: ModelResponse(ToolCallPart) -> ModelRequest(ToolReturnPart)
+    - A standalone assistant response
+    """
+
+    messages: list[ModelMessage]
+    is_tool_sequence: bool = False
+    start_index: int = 0
+    end_index: int = 0
+    _token_count: int | None = field(default=None, repr=False)
+
+    async def get_token_count(self, model_config: ModelConfig) -> int:
+        """Lazily compute and cache token count for this group."""
+        if self._token_count is None:
+            self._token_count = await estimate_tokens_from_messages(
+                self.messages, model_config
+            )
+        return self._token_count
+
+
+@dataclass
+class Chunk:
+    """A chunk of message groups ready for summarization."""
+
+    groups: list[MessageGroup]
+    chunk_index: int
+    total_token_estimate: int = 0
+
+    def get_all_messages(self) -> list[ModelMessage]:
+        """Flatten all messages in this chunk."""
+        messages: list[ModelMessage] = []
+        for group in self.groups:
+            messages.extend(group.messages)
+        return messages
+
+
+def identify_message_groups(messages: list[ModelMessage]) -> list[MessageGroup]:
+    """Identify logical message groups that must stay together.
+
+    Rules:
+    1. Tool calls must include their responses (matched by tool_call_id)
+    2. User messages are individual groups
+    3. Standalone assistant responses are individual groups
+
+    Args:
+        messages: The full message history
+
+    Returns:
+        List of MessageGroup objects
+    """
+    groups: list[MessageGroup] = []
+
+    # Track pending tool calls that need their returns
+    # Maps tool_call_id -> group index
+    pending_tool_calls: dict[str, int] = {}
+
+    for i, msg in enumerate(messages):
+        if isinstance(msg, ModelResponse):
+            # Check for tool calls in response
+            tool_calls = [p for p in msg.parts if isinstance(p, ToolCallPart)]
+
+            if tool_calls:
+                # Start a tool sequence group
+                group = MessageGroup(
+                    messages=[msg],
+                    is_tool_sequence=True,
+                    start_index=i,
+                    end_index=i,
+                )
+                group_idx = len(groups)
+                groups.append(group)
+
+                # Track all tool call IDs in this response
+                for tc in tool_calls:
+                    if tc.tool_call_id:
+                        pending_tool_calls[tc.tool_call_id] = group_idx
+            else:
+                # Standalone assistant response (text only)
+                groups.append(
+                    MessageGroup(
+                        messages=[msg],
+                        is_tool_sequence=False,
+                        start_index=i,
+                        end_index=i,
+                    )
+                )
+
+        elif isinstance(msg, ModelRequest):
+            # Check for tool returns in request
+            tool_returns = [p for p in msg.parts if isinstance(p, ToolReturnPart)]
+            user_prompts = [p for p in msg.parts if isinstance(p, UserPromptPart)]
+
+            if tool_returns:
+                # Add to corresponding tool call groups
+                for tr in tool_returns:
+                    if tr.tool_call_id and tr.tool_call_id in pending_tool_calls:
+                        group_idx = pending_tool_calls.pop(tr.tool_call_id)
+                        groups[group_idx].messages.append(msg)
+                        groups[group_idx].end_index = i
+                    # Note: orphaned tool returns are handled by filter_orphaned_tool_responses
+
+            elif user_prompts:
+                # User message - standalone group
+                groups.append(
+                    MessageGroup(
+                        messages=[msg],
+                        is_tool_sequence=False,
+                        start_index=i,
+                        end_index=i,
+                    )
+                )
+            # Note: System prompts are handled separately by compaction
+
+    logger.debug(
+        f"Identified {len(groups)} message groups "
+        f"({sum(1 for g in groups if g.is_tool_sequence)} tool sequences)"
+    )
+
+    return groups
+
+
+async def create_chunks(
+    groups: list[MessageGroup],
+    model_config: ModelConfig,
+    retention_window: int = RETENTION_WINDOW_MESSAGES,
+) -> tuple[list[Chunk], list[ModelMessage]]:
+    """Create chunks from message groups, respecting token limits.
+
+    Args:
+        groups: List of message groups from identify_message_groups()
+        model_config: Model configuration for token limits
+        retention_window: Number of recent groups to keep outside compaction
+
+    Returns:
+        Tuple of (chunks_to_summarize, retained_recent_messages)
+    """
+    max_chunk_tokens = int(model_config.max_input_tokens * CHUNK_TARGET_RATIO)
+
+    # Handle edge case: too few groups
+    if len(groups) <= retention_window:
+        all_messages: list[ModelMessage] = []
+        for g in groups:
+            all_messages.extend(g.messages)
+        return [], all_messages
+
+    # Separate retention window from groups to chunk
+    groups_to_chunk = groups[:-retention_window]
+    retained_groups = groups[-retention_window:]
+
+    # Build chunks
+    chunks: list[Chunk] = []
+    current_groups: list[MessageGroup] = []
+    current_tokens = 0
+
+    for group in groups_to_chunk:
+        group_tokens = await group.get_token_count(model_config)
+
+        # Handle oversized single group - becomes its own chunk
+        if group_tokens > max_chunk_tokens:
+            # Finish current chunk if any
+            if current_groups:
+                chunks.append(
+                    Chunk(
+                        groups=current_groups,
+                        chunk_index=len(chunks),
+                        total_token_estimate=current_tokens,
+                    )
+                )
+                current_groups = []
+                current_tokens = 0
+
+            # Add oversized as its own chunk
+            chunks.append(
+                Chunk(
+                    groups=[group],
+                    chunk_index=len(chunks),
+                    total_token_estimate=group_tokens,
+                )
+            )
+            logger.warning(
+                f"Oversized message group ({group_tokens:,} tokens) "
+                f"added as single chunk - may need special handling"
+            )
+            continue
+
+        # Would adding this group exceed limit?
+        if current_tokens + group_tokens > max_chunk_tokens:
+            # Finish current chunk
+            if current_groups:
+                chunks.append(
+                    Chunk(
+                        groups=current_groups,
+                        chunk_index=len(chunks),
+                        total_token_estimate=current_tokens,
+                    )
+                )
+            current_groups = [group]
+            current_tokens = group_tokens
+        else:
+            current_groups.append(group)
+            current_tokens += group_tokens
+
+    # Don't forget last chunk
+    if current_groups:
+        chunks.append(
+            Chunk(
+                groups=current_groups,
+                chunk_index=len(chunks),
+                total_token_estimate=current_tokens,
+            )
+        )
+
+    # Extract retained messages
+    retained_messages: list[ModelMessage] = []
+    for g in retained_groups:
+        retained_messages.extend(g.messages)
+
+    # Update chunk indices (in case any were out of order)
+    for i, chunk in enumerate(chunks):
+        chunk.chunk_index = i
+
+    logger.info(
+        f"Created {len(chunks)} chunks for compaction, "
+        f"retaining {len(retained_messages)} recent messages"
+    )
+
+    return chunks, retained_messages
+
+
+async def chunk_messages_for_compaction(
+    messages: list[ModelMessage],
+    model_config: ModelConfig,
+) -> tuple[list[Chunk], list[ModelMessage]]:
+    """Main entry point: chunk oversized conversation for summarization.
+
+    This function identifies logical message groups (preserving tool call sequences),
+    then packs them into chunks that fit within model token limits.
+
+    Args:
+        messages: Full conversation message history
+        model_config: Model configuration for token limits
+
+    Returns:
+        Tuple of (chunks_to_summarize, retention_window_messages)
+    """
+    groups = identify_message_groups(messages)
+    return await create_chunks(groups, model_config)