shotgun-sh 0.1.0__py3-none-any.whl

shotgun/agents/conversation_history.py

@@ -0,0 +1,106 @@
+"""Models and utilities for persisting TUI conversation history."""
+
+from datetime import datetime
+from typing import Any, cast
+
+from pydantic import BaseModel, ConfigDict, Field
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelMessagesTypeAdapter,
+)
+from pydantic_core import to_jsonable_python
+
+from shotgun.tui.screens.chat_screen.hint_message import HintMessage
+
+SerializedMessage = dict[str, Any]
+
+
+class ConversationState(BaseModel):
+    """Represents the complete state of a conversation in memory."""
+
+    agent_messages: list[ModelMessage]
+    ui_messages: list[ModelMessage | HintMessage] = Field(default_factory=list)
+    agent_type: str  # Will store AgentType.value
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+
+class ConversationHistory(BaseModel):
+    """Persistent conversation history for TUI sessions."""
+
+    version: int = 1
+    agent_history: list[SerializedMessage] = Field(
+        default_factory=list
+    )  # Stores serialized ModelMessage objects
+    ui_history: list[SerializedMessage] = Field(
+        default_factory=list
+    )  # Stores serialized ModelMessage and HintMessage objects
+    last_agent_model: str = "research"
+    updated_at: datetime = Field(default_factory=datetime.now)
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    def set_agent_messages(self, messages: list[ModelMessage]) -> None:
+        """Set agent_history from a list of ModelMessage objects.
+
+        Args:
+            messages: List of ModelMessage objects to serialize and store
+        """
+        # Serialize ModelMessage list to JSON-serializable format
+        self.agent_history = to_jsonable_python(
+            messages, fallback=lambda x: str(x), exclude_none=True
+        )
+
+    def set_ui_messages(self, messages: list[ModelMessage | HintMessage]) -> None:
+        """Set ui_history from a list of UI messages."""
+
+        def _serialize_message(
+            message: ModelMessage | HintMessage,
+        ) -> Any:
+            if isinstance(message, HintMessage):
+                data = message.model_dump()
+                data["message_type"] = "hint"
+                return data
+            payload = to_jsonable_python(
+                message, fallback=lambda x: str(x), exclude_none=True
+            )
+            if isinstance(payload, dict):
+                payload.setdefault("message_type", "model")
+            return payload
+
+        self.ui_history = [_serialize_message(msg) for msg in messages]
+
+    def get_agent_messages(self) -> list[ModelMessage]:
+        """Get agent_history as a list of ModelMessage objects.
+
+        Returns:
+            List of deserialized ModelMessage objects
+        """
+        if not self.agent_history:
+            return []
+
+        # Deserialize from JSON format back to ModelMessage objects
+        return ModelMessagesTypeAdapter.validate_python(self.agent_history)
+
+    def get_ui_messages(self) -> list[ModelMessage | HintMessage]:
+        """Get ui_history as a list of Model or hint messages."""
+
+        if not self.ui_history:
+            # Fallback for older conversation files without UI history
+            return cast(list[ModelMessage | HintMessage], self.get_agent_messages())
+
+        messages: list[ModelMessage | HintMessage] = []
+        for item in self.ui_history:
+            message_type = item.get("message_type") if isinstance(item, dict) else None
+            if message_type == "hint":
+                messages.append(HintMessage.model_validate(item))
+                continue
+
+            # Backwards compatibility: data may not include the type marker
+            payload = item
+            if isinstance(payload, dict):
+                payload = {k: v for k, v in payload.items() if k != "message_type"}
+            deserialized = ModelMessagesTypeAdapter.validate_python([payload])
+            messages.append(deserialized[0])
+
+        return messages

shotgun/agents/conversation_manager.py

@@ -0,0 +1,105 @@
+"""Manager for handling conversation persistence operations."""
+
+import json
+from pathlib import Path
+
+from shotgun.logging_config import get_logger
+from shotgun.utils import get_shotgun_home
+
+from .conversation_history import ConversationHistory
+
+logger = get_logger(__name__)
+
+
+class ConversationManager:
+    """Handles saving and loading conversation history."""
+
+    def __init__(self, conversation_path: Path | None = None):
+        """Initialize ConversationManager.
+
+        Args:
+            conversation_path: Path to conversation file.
+                              If None, uses default ~/.shotgun-sh/conversation.json
+        """
+        if conversation_path is None:
+            self.conversation_path = get_shotgun_home() / "conversation.json"
+        else:
+            self.conversation_path = conversation_path
+
+    def save(self, conversation: ConversationHistory) -> None:
+        """Save conversation history to file.
+
+        Args:
+            conversation: ConversationHistory to save
+        """
+        # Ensure directory exists
+        self.conversation_path.parent.mkdir(parents=True, exist_ok=True)
+
+        try:
+            # Update timestamp
+            from datetime import datetime
+
+            conversation.updated_at = datetime.now()
+
+            # Serialize to JSON using Pydantic's model_dump
+            data = conversation.model_dump(mode="json")
+
+            with open(self.conversation_path, "w", encoding="utf-8") as f:
+                json.dump(data, f, indent=2, ensure_ascii=False)
+
+            logger.debug("Conversation saved to %s", self.conversation_path)
+
+        except Exception as e:
+            logger.error(
+                "Failed to save conversation to %s: %s", self.conversation_path, e
+            )
+            # Don't raise - we don't want to interrupt the user's session
+
+    def load(self) -> ConversationHistory | None:
+        """Load conversation history from file.
+
+        Returns:
+            ConversationHistory if file exists and is valid, None otherwise
+        """
+        if not self.conversation_path.exists():
+            logger.debug("No conversation history found at %s", self.conversation_path)
+            return None
+
+        try:
+            with open(self.conversation_path, encoding="utf-8") as f:
+                data = json.load(f)
+
+            conversation = ConversationHistory.model_validate(data)
+            logger.debug(
+                "Conversation loaded from %s with %d agent messages",
+                self.conversation_path,
+                len(conversation.agent_history),
+            )
+            return conversation
+
+        except Exception as e:
+            logger.error(
+                "Failed to load conversation from %s: %s", self.conversation_path, e
+            )
+            return None
+
+    def clear(self) -> None:
+        """Delete the conversation history file."""
+        if self.conversation_path.exists():
+            try:
+                self.conversation_path.unlink()
+                logger.debug(
+                    "Conversation history cleared at %s", self.conversation_path
+                )
+            except Exception as e:
+                logger.error(
+                    "Failed to clear conversation at %s: %s", self.conversation_path, e
+                )
+
+    def exists(self) -> bool:
+        """Check if a conversation history file exists.
+
+        Returns:
+            True if conversation file exists, False otherwise
+        """
+        return self.conversation_path.exists()

shotgun/agents/export.py

@@ -0,0 +1,96 @@
+"""Export agent factory and functions using Pydantic AI with file-based memory."""
+
+from functools import partial
+
+from pydantic_ai import (
+    Agent,
+    DeferredToolRequests,
+)
+from pydantic_ai.agent import AgentRunResult
+from pydantic_ai.messages import ModelMessage
+
+from shotgun.agents.config import ProviderType
+from shotgun.logging_config import get_logger
+
+from .common import (
+    add_system_status_message,
+    build_agent_system_prompt,
+    create_base_agent,
+    create_usage_limits,
+    run_agent,
+)
+from .models import AgentDeps, AgentRuntimeOptions, AgentType
+
+logger = get_logger(__name__)
+
+
+def create_export_agent(
+    agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
+) -> tuple[Agent[AgentDeps, str | DeferredToolRequests], AgentDeps]:
+    """Create an export agent with file management capabilities.
+
+    Args:
+        agent_runtime_options: Agent runtime options for the agent
+        provider: Optional provider override. If None, uses configured default
+
+    Returns:
+        Tuple of (Configured Pydantic AI agent for export management, Agent dependencies)
+    """
+    logger.debug("Initializing export agent")
+    # Use partial to create system prompt function for export agent
+    system_prompt_fn = partial(build_agent_system_prompt, "export")
+
+    agent, deps = create_base_agent(
+        system_prompt_fn,
+        agent_runtime_options,
+        provider=provider,
+        agent_mode=AgentType.EXPORT,
+    )
+    return agent, deps
+
+
+async def run_export_agent(
+    agent: Agent[AgentDeps, str | DeferredToolRequests],
+    instruction: str,
+    deps: AgentDeps,
+    message_history: list[ModelMessage] | None = None,
+) -> AgentRunResult[str | DeferredToolRequests]:
+    """Export artifacts based on the given instruction.
+
+    Args:
+        agent: The configured export agent
+        instruction: The export instruction
+        deps: Agent dependencies
+        message_history: Optional message history for conversation continuity
+
+    Returns:
+        AgentRunResult containing the export process output
+    """
+    logger.debug("📤 Starting export for instruction: %s", instruction)
+
+    message_history = await add_system_status_message(deps, message_history)
+
+    # Let the agent use its tools to read existing artifacts and export them
+    full_prompt = f"Export artifacts or findings based on: {instruction}"
+
+    try:
+        # Create usage limits for responsible API usage
+        usage_limits = create_usage_limits()
+
+        result = await run_agent(
+            agent=agent,
+            prompt=full_prompt,
+            deps=deps,
+            message_history=message_history,
+            usage_limits=usage_limits,
+        )
+
+        logger.debug("✅ Export completed successfully")
+        return result
+
+    except Exception as e:
+        import traceback
+
+        logger.error("Full traceback:\n%s", traceback.format_exc())
+        logger.error("❌ Export failed: %s", str(e))
+        raise

shotgun/agents/history/__init__.py

@@ -0,0 +1,5 @@
+"""History management utilities for Shotgun agents."""
+
+from .history_processors import token_limit_compactor
+
+__all__ = ["token_limit_compactor"]

shotgun/agents/history/compaction.py

@@ -0,0 +1,85 @@
+"""Conversation compaction utilities."""
+
+from pydantic_ai.messages import ModelMessage
+from pydantic_ai.usage import RequestUsage
+
+from shotgun.agents.models import AgentDeps
+from shotgun.logging_config import get_logger
+
+from .token_estimation import estimate_tokens_from_messages
+
+logger = get_logger(__name__)
+
+
+async def apply_persistent_compaction(
+    messages: list[ModelMessage], deps: AgentDeps
+) -> list[ModelMessage]:
+    """Apply compaction to message history for persistent storage.
+
+    This ensures that compacted history is actually used as the conversation baseline,
+    preventing cascading compaction issues across both CLI and TUI usage patterns.
+
+    Args:
+        messages: Full message history from agent run
+        deps: Agent dependencies containing model config
+
+    Returns:
+        Compacted message history that should be stored as conversation state
+    """
+    from .history_processors import token_limit_compactor
+
+    try:
+        # Count actual token usage using shared utility
+        estimated_tokens = estimate_tokens_from_messages(messages, deps.llm_model)
+
+        # Create minimal usage info for compaction check
+        usage = RequestUsage(
+            input_tokens=estimated_tokens,
+            output_tokens=0,
+        )
+
+        # Create a minimal context object for compaction
+        class MockContext:
+            def __init__(self, deps: AgentDeps, usage: RequestUsage | None):
+                self.deps = deps
+                self.usage = usage
+
+        ctx = MockContext(deps, usage)
+        compacted_messages = await token_limit_compactor(ctx, messages)
+
+        # Log the result for monitoring
+        original_size = len(messages)
+        compacted_size = len(compacted_messages)
+
+        if compacted_size < original_size:
+            reduction_pct = ((original_size - compacted_size) / original_size) * 100
+            logger.debug(
+                f"Persistent compaction applied: {original_size} → {compacted_size} messages "
+                f"({reduction_pct:.1f}% reduction)"
+            )
+        else:
+            logger.debug(
+                f"No persistent compaction needed: {original_size} messages unchanged"
+            )
+
+        return compacted_messages
+
+    except Exception as e:
+        # If compaction fails, return original messages
+        # This ensures the system remains functional even if compaction has issues
+        logger.warning(f"Persistent compaction failed, using original history: {e}")
+        return messages
+
+
+def should_apply_persistent_compaction(deps: AgentDeps) -> bool:
+    """Check if persistent compaction should be applied.
+
+    Args:
+        deps: Agent dependencies
+
+    Returns:
+        True if persistent compaction should be applied
+    """
+    # For now, always apply persistent compaction
+    # Future: Add configuration option in deps or environment variable
+    return True

shotgun/agents/history/constants.py

@@ -0,0 +1,19 @@
+"""Constants for history processing and compaction."""
+
+from enum import Enum
+
+# Summary marker for compacted history
+SUMMARY_MARKER = "📌 COMPACTED_HISTORY:"
+
+# Token calculation constants
+INPUT_BUFFER_TOKENS = 500
+MIN_SUMMARY_TOKENS = 100
+TOKEN_LIMIT_RATIO = 0.8
+
+
+class SummaryType(Enum):
+    """Types of summarization requests for logging."""
+
+    INCREMENTAL = "INCREMENTAL"
+    FULL = "FULL"
+    CONTEXT_EXTRACTION = "CONTEXT_EXTRACTION"

shotgun/agents/history/context_extraction.py

@@ -0,0 +1,108 @@
+"""Context extraction utilities for history processing."""
+
+from pydantic_ai.messages import (
+    BuiltinToolCallPart,
+    BuiltinToolReturnPart,
+    ModelMessage,
+    ModelRequest,
+    ModelResponse,
+    ModelResponsePart,
+    RetryPromptPart,
+    SystemPromptPart,
+    TextPart,
+    ThinkingPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
+
+
+def extract_context_from_messages(messages: list[ModelMessage]) -> str:
+    """Extract context from a list of messages for summarization."""
+    context = ""
+    for msg in messages:
+        if isinstance(msg, ModelResponse | ModelRequest):
+            for part in msg.parts:
+                message_content = extract_context_from_part(part)
+                if message_content:
+                    context += message_content + "\n"
+    return context
+
+
+def extract_context_from_message_range(
+    messages: list[ModelMessage],
+    start_index: int,
+    end_index: int | None = None,
+) -> str:
+    """Extract context from a specific range of messages."""
+    if end_index is None:
+        end_index = len(messages)
+
+    message_slice = messages[start_index:end_index]
+    return extract_context_from_messages(message_slice)
+
+
+def has_meaningful_content(messages: list[ModelMessage]) -> bool:
+    """Check if messages contain meaningful content worth summarizing.
+
+    Only ModelResponse messages are considered meaningful for summarization.
+    User requests alone don't need summarization.
+    """
+    for msg in messages:
+        if isinstance(msg, ModelResponse):
+            for part in msg.parts:
+                if extract_context_from_part(part):
+                    return True
+    return False
+
+
+def extract_context_from_part(
+    message_part: (
+        SystemPromptPart
+        | UserPromptPart
+        | ToolReturnPart
+        | RetryPromptPart
+        | ModelResponsePart
+    ),
+) -> str:
+    """Extract context from a single message part."""
+    if isinstance(message_part, SystemPromptPart):
+        return ""  # Exclude system prompts from summary
+
+    elif isinstance(message_part, UserPromptPart):
+        if isinstance(message_part.content, str):
+            return f"<USER_PROMPT>\n{message_part.content}\n</USER_PROMPT>"
+        return ""
+
+    elif isinstance(message_part, ToolReturnPart):
+        return f"<TOOL_RETURN>\n{str(message_part.content)}\n</TOOL_RETURN>"
+
+    elif isinstance(message_part, RetryPromptPart):
+        if isinstance(message_part.content, str):
+            return f"<RETRY_PROMPT>\n{message_part.content}\n</RETRY_PROMPT>"
+        return ""
+
+    # Handle ModelResponsePart types
+    elif isinstance(message_part, TextPart):
+        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>"
+
+    elif isinstance(message_part, BuiltinToolCallPart):
+        return f"<BUILTIN_TOOL_CALL>\n{message_part.tool_name}\n</BUILTIN_TOOL_CALL>"
+
+    elif isinstance(message_part, BuiltinToolReturnPart):
+        return (
+            f"<BUILTIN_TOOL_RETURN>\n{message_part.tool_name}\n</BUILTIN_TOOL_RETURN>"
+        )
+
+    elif isinstance(message_part, ThinkingPart):
+        return f"<THINKING>\n{message_part.content}\n</THINKING>"
+
+    return ""

shotgun/agents/history/history_building.py

@@ -0,0 +1,104 @@
+"""Functions for building compacted message history."""
+
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelRequest,
+    ModelRequestPart,
+    ModelResponse,
+    SystemPromptPart,
+    TextPart,
+    UserPromptPart,
+)
+
+from .message_utils import (
+    get_first_user_request,
+    get_last_user_request,
+    get_system_prompt,
+    get_user_content_from_request,
+)
+
+
+def build_clean_compacted_history(
+    summary_part: TextPart,
+    all_messages: list[ModelMessage],
+    last_summary_index: int | None = None,
+) -> list[ModelMessage]:
+    """Build a clean compacted history without preserving old verbose content.
+
+    Args:
+        summary_part: The marked summary part to include
+        all_messages: Original message history
+        last_summary_index: Index of the last summary (if any)
+
+    Returns:
+        Clean compacted message history
+    """
+    # Extract essential context from pre-summary messages (if any)
+    system_prompt = ""
+    first_user_prompt = ""
+
+    if last_summary_index is not None and last_summary_index > 0:
+        # Get system and first user from original conversation
+        pre_summary_messages = all_messages[:last_summary_index]
+        system_prompt = get_system_prompt(pre_summary_messages) or ""
+        first_user_prompt = get_first_user_request(pre_summary_messages) or ""
+
+    # Build the base structure
+    compacted_messages: list[ModelMessage] = []
+
+    # Add system/user context if it exists and is meaningful
+    if system_prompt or first_user_prompt:
+        compacted_messages.append(
+            _create_base_request(system_prompt, first_user_prompt)
+        )
+
+    # Add the summary
+    summary_message = ModelResponse(parts=[summary_part])
+    compacted_messages.append(summary_message)
+
+    # Ensure proper ending
+    return ensure_ends_with_model_request(compacted_messages, all_messages)
+
+
+def ensure_ends_with_model_request(
+    compacted_messages: list[ModelMessage],
+    original_messages: list[ModelMessage],
+) -> list[ModelMessage]:
+    """Ensure the message history ends with ModelRequest for PydanticAI compatibility."""
+    last_user_request = get_last_user_request(original_messages)
+
+    if not last_user_request:
+        return compacted_messages
+
+    # Check if we need to add the last request or restructure
+    if compacted_messages and isinstance(compacted_messages[0], ModelRequest):
+        first_request = compacted_messages[0]
+        last_user_content = get_user_content_from_request(last_user_request)
+        first_user_content = get_user_content_from_request(first_request)
+
+        if last_user_content != first_user_content:
+            # Different messages - append the last request
+            compacted_messages.append(last_user_request)
+        else:
+            # Same message - restructure to end with ModelRequest
+            if len(compacted_messages) >= 2:
+                summary_message = compacted_messages[1]  # The summary
+                compacted_messages = [summary_message, first_request]
+    else:
+        # No first request, just add the last one
+        compacted_messages.append(last_user_request)
+
+    return compacted_messages
+
+
+def _create_base_request(system_prompt: str, user_prompt: str) -> ModelRequest:
+    """Create the base ModelRequest with system and user prompts."""
+    parts: list[ModelRequestPart] = []
+
+    if system_prompt:
+        parts.append(SystemPromptPart(content=system_prompt))
+
+    if user_prompt:
+        parts.append(UserPromptPart(content=user_prompt))
+
+    return ModelRequest(parts=parts)