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

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)

shotgun/agents/{history → conversation/history}/compaction.py

@@ -20,6 +20,10 @@ async def apply_persistent_compaction(
     This ensures that compacted history is actually used as the conversation baseline,
     preventing cascading compaction issues across both CLI and TUI usage patterns.
 
+    Compaction happens in two phases:
+    1. Deterministic pre-compaction: Remove file content (no LLM needed)
+    2. LLM-based compaction: Summarize conversation if still over threshold
+
     Args:
         messages: Full message history from agent run
         deps: Agent dependencies containing model config
@@ -28,10 +32,32 @@ async def apply_persistent_compaction(
     Returns:
         Compacted message history that should be stored as conversation state
     """
+    from .file_content_deduplication import deduplicate_file_content
     from .history_processors import token_limit_compactor
 
     try:
-        # Count actual token usage using shared utility
+        # STEP 1: Deterministic pre-compaction (no LLM cost)
+        # Remove file content from tool returns - files are still accessible
+        # via retrieve_code (codebase) or read_file (.shotgun/ folder)
+        messages, tokens_saved = deduplicate_file_content(
+            messages,
+            retention_window=3,  # Keep last 3 messages' file content intact
+        )
+
+        if tokens_saved > 0:
+            logger.info(
+                f"Pre-compaction: removed ~{tokens_saved:,} tokens of file content"
+            )
+            track_event(
+                "file_content_deduplication",
+                {
+                    "tokens_saved_estimate": tokens_saved,
+                    "retention_window": 3,
+                    "model_name": deps.llm_model.name.value,
+                },
+            )
+
+        # STEP 2: Count tokens after pre-compaction
         estimated_tokens = await estimate_tokens_from_messages(messages, deps.llm_model)
 
         # Create minimal usage info for compaction check

shotgun/agents/{history → conversation/history}/constants.py

@@ -10,6 +10,11 @@ INPUT_BUFFER_TOKENS = 500
 MIN_SUMMARY_TOKENS = 100
 TOKEN_LIMIT_RATIO = 0.8
 
+# Chunked compaction constants
+CHUNK_TARGET_RATIO = 0.60  # Target chunk size as % of max_input_tokens
+CHUNK_SAFE_RATIO = 0.70  # Max safe ratio before triggering chunked compaction
+RETENTION_WINDOW_MESSAGES = 5  # Keep last N message groups outside compaction
+
 
 class SummaryType(Enum):
     """Types of summarization requests for logging."""

shotgun/agents/conversation/history/file_content_deduplication.py

@@ -0,0 +1,239 @@
+"""Pre-compaction file content deduplication for conversation history.
+
+This module provides a deterministic pre-pass that removes file content from
+tool returns before LLM-based compaction. Files are still accessible via
+`retrieve_code` (codebase) or `read_file` (.shotgun/ folder).
+"""
+
+from enum import StrEnum
+from typing import Any
+
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ToolReturnPart,
+)
+
+from shotgun.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+class FileReadTool(StrEnum):
+    """Tool names that read file content."""
+
+    CODEBASE = "file_read"  # Reads from indexed codebase (Kuzu graph)
+    SHOTGUN_FOLDER = "read_file"  # Reads from .shotgun/ folder
+
+
+# Minimum content length to bother deduplicating (skip tiny files)
+MIN_CONTENT_LENGTH = 500
+
+# Placeholder templates for each type
+CODEBASE_PLACEHOLDER = (
+    "**File**: `{file_path}`\n"
+    "**Size**: {size_bytes} bytes | **Language**: {language}\n"
+    "**Content**: [Removed for compaction - use `retrieve_code` or `file_read` to access]"
+)
+
+SHOTGUN_PLACEHOLDER = (
+    "**File**: `.shotgun/{filename}`\n"
+    "**Content**: [Removed for compaction - file persisted in .shotgun/ folder]"
+)
+
+# Simple prefix for detecting file_read output format
+# Instead of using regex, we just check for the expected prefix and extract the file path
+CODEBASE_FILE_PREFIX = "**File**: `"
+
+
+def _extract_file_path(content: str) -> str | None:
+    """Extract file path from file_read tool return content.
+
+    Uses simple string operations instead of regex for maximum performance.
+    The file_read tool output format is: **File**: `path`\\n...
+
+    Args:
+        content: The tool return content string
+
+    Returns:
+        The file path or None if format doesn't match
+    """
+    # Fast check: content must start with expected prefix
+    if not content.startswith(CODEBASE_FILE_PREFIX):
+        return None
+
+    # Find the closing backtick after the prefix
+    prefix_len = len(CODEBASE_FILE_PREFIX)
+    backtick_pos = content.find("`", prefix_len)
+
+    if backtick_pos == -1:
+        return None
+
+    return content[prefix_len:backtick_pos]
+
+
+def _get_language_from_path(file_path: str) -> str:
+    """Infer programming language from file extension."""
+    from pathlib import Path
+
+    from shotgun.codebase.core.language_config import get_language_config
+
+    ext = Path(file_path).suffix
+    config = get_language_config(ext)
+    return config.name if config else "unknown"
+
+
+def _create_codebase_placeholder(file_path: str, size_bytes: int, language: str) -> str:
+    """Create placeholder for codebase file content."""
+    return CODEBASE_PLACEHOLDER.format(
+        file_path=file_path,
+        size_bytes=size_bytes,
+        language=language or "unknown",
+    )
+
+
+def _create_shotgun_placeholder(filename: str) -> str:
+    """Create placeholder for .shotgun/ file content."""
+    return SHOTGUN_PLACEHOLDER.format(filename=filename)
+
+
+def _estimate_tokens_saved(original: str, replacement: str) -> int:
+    """Rough estimate of tokens saved (~4 chars per token)."""
+    original_chars = len(original)
+    replacement_chars = len(replacement)
+    # Rough token estimate: ~4 characters per token for code
+    return max(0, (original_chars - replacement_chars) // 4)
+
+
+def deduplicate_file_content(
+    messages: list[ModelMessage],
+    retention_window: int = 3,
+) -> tuple[list[ModelMessage], int]:
+    """Replace file read content with placeholders for indexed/persisted files.
+
+    This is a deterministic pre-compaction pass that reduces tokens without
+    requiring an LLM. Files remain accessible via their respective tools.
+
+    This function uses copy-on-write semantics: only messages that need
+    modification are copied, while unmodified messages are reused by reference.
+    This significantly reduces memory allocation and processing time for large
+    conversations where only a subset of messages contain file content.
+
+    Args:
+        messages: Conversation history
+        retention_window: Keep full content in last N messages (for recent context)
+
+    Returns:
+        Tuple of (modified_messages, estimated_tokens_saved)
+    """
+    if not messages:
+        return messages, 0
+
+    total_tokens_saved = 0
+    files_deduplicated = 0
+
+    # Calculate retention boundary (keep last N messages intact)
+    retention_start = max(0, len(messages) - retention_window)
+
+    # Track which message indices need replacement
+    # We use a dict to store index -> new_message mappings
+    replacements: dict[int, ModelMessage] = {}
+
+    for msg_idx, message in enumerate(messages):
+        # Skip messages in retention window
+        if msg_idx >= retention_start:
+            continue
+
+        # Only process ModelRequest (which contains ToolReturnPart)
+        if not isinstance(message, ModelRequest):
+            continue
+
+        # Build new parts list, replacing file content where appropriate
+        new_parts: list[Any] = []
+        message_modified = False
+
+        for part in message.parts:
+            if not isinstance(part, ToolReturnPart):
+                new_parts.append(part)
+                continue
+
+            tool_name = part.tool_name
+            content = part.content
+
+            # Skip if content is too short to bother
+            if not isinstance(content, str) or len(content) < MIN_CONTENT_LENGTH:
+                new_parts.append(part)
+                continue
+
+            replacement = None
+            original_content = content
+
+            # Handle codebase file reads (file_read)
+            if tool_name == FileReadTool.CODEBASE:
+                file_path = _extract_file_path(content)
+                if file_path:
+                    # Use content length as size estimate (includes formatting overhead
+                    # but close enough for deduplication purposes)
+                    size_bytes = len(content)
+                    language = _get_language_from_path(file_path)
+                    replacement = _create_codebase_placeholder(
+                        file_path, size_bytes, language
+                    )
+                    logger.debug(
+                        f"Deduplicating codebase file: {file_path} ({size_bytes} bytes)"
+                    )
+
+            # Handle .shotgun/ file reads (read_file)
+            elif tool_name == FileReadTool.SHOTGUN_FOLDER:
+                # For read_file, content is raw - we need to figure out filename
+                # from the tool call args (but we only have the return here)
+                # Use a generic placeholder since we don't have the filename
+                if len(content) >= MIN_CONTENT_LENGTH:
+                    # Try to extract filename from content if it looks like markdown
+                    # Otherwise use generic placeholder
+                    replacement = _create_shotgun_placeholder("artifact")
+                    logger.debug(
+                        f"Deduplicating .shotgun/ file read ({len(content)} chars)"
+                    )
+
+            # Apply replacement if we have one
+            if replacement:
+                # Create new ToolReturnPart with replaced content
+                new_part = ToolReturnPart(
+                    tool_name=part.tool_name,
+                    tool_call_id=part.tool_call_id,
+                    content=replacement,
+                    timestamp=part.timestamp,
+                )
+                new_parts.append(new_part)
+                message_modified = True
+
+                tokens_saved = _estimate_tokens_saved(original_content, replacement)
+                total_tokens_saved += tokens_saved
+                files_deduplicated += 1
+            else:
+                new_parts.append(part)
+
+        # Only create a new message if parts were actually modified
+        if message_modified:
+            replacements[msg_idx] = ModelRequest(parts=new_parts)
+
+    # If no modifications were made, return original list (no allocation needed)
+    if not replacements:
+        return messages, 0
+
+    # Build result list with copy-on-write: reuse unmodified messages
+    modified_messages: list[ModelMessage] = []
+    for idx, msg in enumerate(messages):
+        if idx in replacements:
+            modified_messages.append(replacements[idx])
+        else:
+            modified_messages.append(msg)
+
+    if files_deduplicated > 0:
+        logger.info(
+            f"File content deduplication: {files_deduplicated} files, "
+            f"~{total_tokens_saved:,} tokens saved"
+        )
+
+    return modified_messages, total_tokens_saved