shotgun-sh 0.1.0.dev22__py3-none-any.whl → 0.1.0.dev24__py3-none-any.whl

shotgun/agents/agent_manager.py

@@ -1,10 +1,12 @@
 """Agent manager for coordinating multiple AI agents with shared message history."""
 
 import logging
-from collections.abc import AsyncIterable
-from dataclasses import dataclass, field
-from enum import Enum
-from typing import Any, cast
+from collections.abc import AsyncIterable, Sequence
+from dataclasses import dataclass, field, is_dataclass, replace
+from typing import TYPE_CHECKING, Any, cast
+
+if TYPE_CHECKING:
+    from shotgun.agents.conversation_history import ConversationState
 
 from pydantic_ai import (
     Agent,
@@ -21,6 +23,7 @@ from pydantic_ai.messages import (
     FunctionToolResultEvent,
     ModelMessage,
     ModelRequest,
+    ModelRequestPart,
     ModelResponse,
     ModelResponsePart,
     PartDeltaEvent,
@@ -33,10 +36,11 @@ from textual.message import Message
 from textual.widget import Widget
 
 from shotgun.agents.common import add_system_prompt_message, add_system_status_message
+from shotgun.agents.models import AgentType, FileOperation
 
 from .export import create_export_agent
 from .history.compaction import apply_persistent_compaction
-from .models import AgentDeps, AgentRuntimeOptions, FileOperation
+from .models import AgentDeps, AgentRuntimeOptions
 from .plan import create_plan_agent
 from .research import create_research_agent
 from .specify import create_specify_agent
@@ -45,16 +49,6 @@ from .tasks import create_tasks_agent
 logger = logging.getLogger(__name__)
 
 
-class AgentType(Enum):
-    """Enumeration for available agent types (for Python < 3.11)."""
-
-    RESEARCH = "research"
-    PLAN = "plan"
-    TASKS = "tasks"
-    SPECIFY = "specify"
-    EXPORT = "export"
-
-
 class MessageHistoryUpdated(Message):
     """Event posted when the message history is updated."""
 
@@ -281,6 +275,8 @@ class AgentManager(Widget):
         # Start with persistent message history
         message_history = self.message_history
 
+        deps.agent_mode = self._current_agent_type
+
         # Add a system status message so the agent knows whats going on
         message_history = await add_system_status_message(deps, message_history)
 
@@ -463,3 +459,87 @@ class AgentManager(Widget):
                 file_operations=file_operations,
             )
         )
+
+    def _filter_system_prompts(
+        self, messages: list[ModelMessage]
+    ) -> list[ModelMessage]:
+        """Filter out system prompts from messages for UI display.
+
+        Args:
+            messages: List of messages that may contain system prompts
+
+        Returns:
+            List of messages without system prompt parts
+        """
+        from pydantic_ai.messages import SystemPromptPart
+
+        filtered_messages: list[ModelMessage] = []
+        for msg in messages:
+            parts: Sequence[ModelRequestPart] | Sequence[ModelResponsePart] | None = (
+                msg.parts if hasattr(msg, "parts") else None
+            )
+            if not parts:
+                filtered_messages.append(msg)
+                continue
+
+            non_system_parts = [
+                part for part in parts if not isinstance(part, SystemPromptPart)
+            ]
+
+            if not non_system_parts:
+                # Skip messages made up entirely of system prompt parts (e.g. system message)
+                continue
+
+            if len(non_system_parts) == len(parts):
+                # Nothing was filtered – keep original message
+                filtered_messages.append(msg)
+                continue
+
+            if is_dataclass(msg):
+                filtered_messages.append(
+                    # ignore types because of the convoluted Request | Response types
+                    replace(msg, parts=cast(Any, non_system_parts))
+                )
+            else:
+                filtered_messages.append(msg)
+        return filtered_messages
+
+    def get_conversation_state(self) -> "ConversationState":
+        """Get the current conversation state.
+
+        Returns:
+            ConversationState object containing UI and agent messages and current type
+        """
+        from shotgun.agents.conversation_history import ConversationState
+
+        return ConversationState(
+            agent_messages=self.message_history.copy(),
+            agent_type=self._current_agent_type.value,
+        )
+
+    def restore_conversation_state(self, state: "ConversationState") -> None:
+        """Restore conversation state from a saved state.
+
+        Args:
+            state: ConversationState object to restore
+        """
+        # Restore message history for agents (includes system prompts)
+        self.message_history = state.agent_messages.copy()
+
+        # Filter out system prompts for UI display
+        self.ui_message_history = self._filter_system_prompts(state.agent_messages)
+
+        # Restore agent type
+        self._current_agent_type = AgentType(state.agent_type)
+
+        # Notify listeners about the restored messages
+        self._post_messages_updated()
+
+
+# Re-export AgentType for backward compatibility
+__all__ = [
+    "AgentManager",
+    "AgentType",
+    "MessageHistoryUpdated",
+    "PartialResponseMessage",
+]

shotgun/agents/common.py

@@ -19,10 +19,12 @@ from pydantic_ai.messages import (
 )
 
 from shotgun.agents.config import ProviderType, get_config_manager, get_provider_model
+from shotgun.agents.models import AgentType
 from shotgun.logging_config import get_logger
 from shotgun.prompts import PromptLoader
-from shotgun.sdk.services import get_artifact_service, get_codebase_service
+from shotgun.sdk.services import get_codebase_service
 from shotgun.utils import ensure_shotgun_directory_exists
+from shotgun.utils.file_system_utils import get_shotgun_base_path
 
 from .history import token_limit_compactor
 from .history.compaction import apply_persistent_compaction
@@ -38,14 +40,7 @@ from .tools import (
     retrieve_code,
     write_file,
 )
-from .tools.artifact_management import (
-    create_artifact,
-    list_artifact_templates,
-    list_artifacts,
-    read_artifact,
-    read_artifact_section,
-    write_artifact_section,
-)
+from .tools.file_management import AGENT_DIRECTORIES
 
 logger = get_logger(__name__)
 
@@ -67,18 +62,23 @@ async def add_system_status_message(
         Updated message history with system status message prepended
     """
     message_history = message_history or []
-    codebase_understanding_graphs = await deps.codebase_service.list_graphs()
+    # Only show graphs for the current working directory
+    codebase_understanding_graphs = (
+        await deps.codebase_service.list_graphs_for_directory()
+    )
 
-    # Collect artifact state information
-    from .artifact_state import collect_artifact_state
+    # Get existing files for the agent
+    existing_files = get_agent_existing_files(deps.agent_mode)
 
-    artifact_state = collect_artifact_state()
+    # Extract table of contents from the agent's markdown file
+    markdown_toc = extract_markdown_toc(deps.agent_mode)
 
     system_state = prompt_loader.render(
         "agents/state/system_state.j2",
         codebase_understanding_graphs=codebase_understanding_graphs,
         is_tui_context=deps.is_tui_context,
-        **artifact_state,
+        existing_files=existing_files,
+        markdown_toc=markdown_toc,
     )
 
     message_history.append(
@@ -97,14 +97,17 @@ def create_base_agent(
     load_codebase_understanding_tools: bool = True,
     additional_tools: list[Any] | None = None,
     provider: ProviderType | None = None,
+    agent_mode: AgentType | None = None,
 ) -> tuple[Agent[AgentDeps, str | DeferredToolRequests], AgentDeps]:
     """Create a base agent with common configuration.
 
     Args:
         system_prompt_fn: Function that will be decorated as system_prompt
         agent_runtime_options: Agent runtime options for the agent
+        load_codebase_understanding_tools: Whether to load codebase understanding tools
         additional_tools: Optional list of additional tools
         provider: Optional provider override. If None, uses configured default
+        agent_mode: The mode of the agent (research, plan, tasks, specify, export)
 
     Returns:
         Tuple of (Configured Pydantic AI agent, Agent dependencies)
@@ -126,13 +129,12 @@ def create_base_agent(
 
         # Create deps with model config and services
         codebase_service = get_codebase_service()
-        artifact_service = get_artifact_service()
         deps = AgentDeps(
             **agent_runtime_options.model_dump(),
             llm_model=model_config,
             codebase_service=codebase_service,
-            artifact_service=artifact_service,
             system_prompt_fn=system_prompt_fn,
+            agent_mode=agent_mode,
         )
 
     except Exception as e:
@@ -180,14 +182,6 @@ def create_base_agent(
     agent.tool(append_file)
     agent.tool(read_file)
 
-    # Register artifact management tools (always available)
-    agent.tool(create_artifact)
-    agent.tool(list_artifacts)
-    agent.tool(list_artifact_templates)
-    agent.tool(read_artifact)
-    agent.tool(read_artifact_section)
-    agent.tool(write_artifact_section)
-
     # Register codebase understanding tools (conditional)
     if load_codebase_understanding_tools:
         agent.tool(query_graph)
@@ -199,10 +193,134 @@ def create_base_agent(
     else:
         logger.debug("🚫🧠 Codebase understanding tools not registered")
 
-    logger.debug("✅ Agent creation complete with artifact and codebase tools")
+    logger.debug("✅ Agent creation complete with codebase tools")
     return agent, deps
 
 
+def extract_markdown_toc(agent_mode: AgentType | None) -> str | None:
+    """Extract table of contents from agent's markdown file.
+
+    Args:
+        agent_mode: The agent mode to extract TOC for
+
+    Returns:
+        Formatted TOC string (up to 2000 chars) or None if not applicable
+    """
+    # Skip for EXPORT mode or no mode
+    if (
+        not agent_mode
+        or agent_mode == AgentType.EXPORT
+        or agent_mode not in AGENT_DIRECTORIES
+    ):
+        return None
+
+    base_path = get_shotgun_base_path()
+    md_file = AGENT_DIRECTORIES[agent_mode]
+    md_path = base_path / md_file
+
+    # Check if the markdown file exists
+    if not md_path.exists():
+        return None
+
+    try:
+        content = md_path.read_text(encoding="utf-8")
+        lines = content.split("\n")
+
+        # Extract headings
+        toc_lines = []
+        for line in lines:
+            stripped = line.strip()
+            if stripped.startswith("#"):
+                # Count the heading level
+                level = 0
+                for char in stripped:
+                    if char == "#":
+                        level += 1
+                    else:
+                        break
+
+                # Get the heading text (remove the # symbols and clean up)
+                heading_text = stripped[level:].strip()
+                if heading_text:
+                    # Add indentation based on level
+                    indent = "  " * (level - 1)
+                    toc_lines.append(f"{indent}{'#' * level} {heading_text}")
+
+        if not toc_lines:
+            return None
+
+        # Join and truncate to 2000 characters
+        toc = "\n".join(toc_lines)
+        if len(toc) > 2000:
+            toc = toc[:1997] + "..."
+
+        return toc
+
+    except Exception as e:
+        logger.debug(f"Failed to extract TOC from {md_file}: {e}")
+        return None
+
+
+def get_agent_existing_files(agent_mode: AgentType | None = None) -> list[str]:
+    """Get list of existing files for the given agent mode.
+
+    Args:
+        agent_mode: The agent mode to check files for. If None, lists all files.
+
+    Returns:
+        List of existing file paths relative to .shotgun directory
+    """
+    base_path = get_shotgun_base_path()
+    existing_files = []
+
+    # If no agent mode, list all files in base path and first level subdirectories
+    if agent_mode is None:
+        # List files in the root .shotgun directory
+        for item in base_path.iterdir():
+            if item.is_file():
+                existing_files.append(item.name)
+            elif item.is_dir():
+                # List files in first-level subdirectories
+                for subitem in item.iterdir():
+                    if subitem.is_file():
+                        relative_path = subitem.relative_to(base_path)
+                        existing_files.append(str(relative_path))
+        return existing_files
+
+    # Handle specific agent modes
+    if agent_mode not in AGENT_DIRECTORIES:
+        return []
+
+    if agent_mode == AgentType.EXPORT:
+        # For export agent, list all files in exports directory
+        exports_dir = base_path / "exports"
+        if exports_dir.exists():
+            for file_path in exports_dir.rglob("*"):
+                if file_path.is_file():
+                    relative_path = file_path.relative_to(base_path)
+                    existing_files.append(str(relative_path))
+    else:
+        # For other agents, check both .md file and directory with same name
+        allowed_file = AGENT_DIRECTORIES[agent_mode]
+
+        # Check for the .md file
+        md_file_path = base_path / allowed_file
+        if md_file_path.exists():
+            existing_files.append(allowed_file)
+
+        # Check for directory with same base name (e.g., research/ for research.md)
+        base_name = allowed_file.replace(".md", "")
+        dir_path = base_path / base_name
+        if dir_path.exists() and dir_path.is_dir():
+            # List all files in the directory
+            for file_path in dir_path.rglob("*"):
+                if file_path.is_file():
+                    relative_path = file_path.relative_to(base_path)
+                    existing_files.append(str(relative_path))
+
+    return existing_files
+
+
 def build_agent_system_prompt(
     agent_type: str,
     ctx: RunContext[AgentDeps],

shotgun/agents/conversation_history.py

@@ -0,0 +1,56 @@
+"""Models and utilities for persisting TUI conversation history."""
+
+from datetime import datetime
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, Field
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelMessagesTypeAdapter,
+)
+from pydantic_core import to_jsonable_python
+
+
+class ConversationState(BaseModel):
+    """Represents the complete state of a conversation in memory."""
+
+    agent_messages: list[ModelMessage]
+    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[dict[str, Any]] = Field(
+        default_factory=list
+    )  # Will store serialized ModelMessage 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 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)

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

@@ -19,7 +19,7 @@ from .common import (
     create_usage_limits,
     run_agent,
 )
-from .models import AgentDeps, AgentRuntimeOptions
+from .models import AgentDeps, AgentRuntimeOptions, AgentType
 
 logger = get_logger(__name__)
 
@@ -41,7 +41,10 @@ def create_export_agent(
     system_prompt_fn = partial(build_agent_system_prompt, "export")
 
     agent, deps = create_base_agent(
-        system_prompt_fn, agent_runtime_options, provider=provider
+        system_prompt_fn,
+        agent_runtime_options,
+        provider=provider,
+        agent_mode=AgentType.EXPORT,
     )
     return agent, deps
 

shotgun/agents/models.py

@@ -4,7 +4,7 @@ import os
 from asyncio import Future, Queue
 from collections.abc import Callable
 from datetime import datetime
-from enum import Enum
+from enum import Enum, StrEnum
 from pathlib import Path
 from typing import TYPE_CHECKING
 
@@ -14,10 +14,19 @@ from pydantic_ai import RunContext
 from .config.models import ModelConfig
 
 if TYPE_CHECKING:
-    from shotgun.artifacts.service import ArtifactService
     from shotgun.codebase.service import CodebaseService
 
 
+class AgentType(StrEnum):
+    """Enumeration for available agent types."""
+
+    RESEARCH = "research"
+    SPECIFY = "specify"
+    PLAN = "plan"
+    TASKS = "tasks"
+    EXPORT = "export"
+
+
 class UserAnswer(BaseModel):
     """A answer from the user."""
 
@@ -224,10 +233,6 @@ class AgentDeps(AgentRuntimeOptions):
         description="Codebase service for code analysis tools",
     )
 
-    artifact_service: "ArtifactService" = Field(
-        description="Artifact service for managing structured artifacts",
-    )
-
     system_prompt_fn: Callable[[RunContext["AgentDeps"]], str] = Field(
         description="Function that generates the system prompt for this agent",
     )
@@ -237,10 +242,14 @@ class AgentDeps(AgentRuntimeOptions):
         description="Tracker for file operations during agent run",
     )
 
+    agent_mode: AgentType | None = Field(
+        default=None,
+        description="Current agent mode for file scoping",
+    )
+
 
 # Rebuild model to resolve forward references after imports are available
 try:
-    from shotgun.artifacts.service import ArtifactService
     from shotgun.codebase.service import CodebaseService
 
     AgentDeps.model_rebuild()

shotgun/agents/plan.py

@@ -19,7 +19,7 @@ from .common import (
     create_usage_limits,
     run_agent,
 )
-from .models import AgentDeps, AgentRuntimeOptions
+from .models import AgentDeps, AgentRuntimeOptions, AgentType
 
 logger = get_logger(__name__)
 
@@ -46,6 +46,7 @@ def create_plan_agent(
         load_codebase_understanding_tools=True,
         additional_tools=None,
         provider=provider,
+        agent_mode=AgentType.PLAN,
     )
     return agent, deps
 

shotgun/agents/research.py

@@ -21,7 +21,7 @@ from .common import (
     create_usage_limits,
     run_agent,
 )
-from .models import AgentDeps, AgentRuntimeOptions
+from .models import AgentDeps, AgentRuntimeOptions, AgentType
 from .tools import get_available_web_search_tools
 
 logger = get_logger(__name__)
@@ -60,6 +60,7 @@ def create_research_agent(
         load_codebase_understanding_tools=True,
         additional_tools=web_search_tools,
         provider=provider,
+        agent_mode=AgentType.RESEARCH,
     )
     return agent, deps
 

shotgun/agents/specify.py

@@ -19,7 +19,7 @@ from .common import (
     create_usage_limits,
     run_agent,
 )
-from .models import AgentDeps, AgentRuntimeOptions
+from .models import AgentDeps, AgentRuntimeOptions, AgentType
 
 logger = get_logger(__name__)
 
@@ -46,6 +46,7 @@ def create_specify_agent(
         load_codebase_understanding_tools=True,
         additional_tools=None,
         provider=provider,
+        agent_mode=AgentType.SPECIFY,
     )
     return agent, deps