shotgun-sh 0.1.0__py3-none-any.whl

shotgun/agents/models.py

@@ -0,0 +1,275 @@
+"""Pydantic models for agent dependencies and configuration."""
+
+import os
+from asyncio import Future, Queue
+from collections.abc import Callable
+from datetime import datetime
+from enum import Enum, StrEnum
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, ConfigDict, Field
+from pydantic_ai import RunContext
+
+from .config.models import ModelConfig
+
+if TYPE_CHECKING:
+    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 PipelineConfigEntry(BaseModel):
+    """Configuration for each agent in the pipeline.
+
+    This model defines what files an agent can write to and what
+    files from prior agents it should read for context.
+    """
+
+    own_file: str | None = Field(
+        default=None,
+        description="The file this agent writes to (None for export agent)",
+    )
+    prior_files: list[str] = Field(
+        default_factory=list,
+        description="Files from prior agents in pipeline to read for context",
+    )
+
+
+class UserAnswer(BaseModel):
+    """A answer from the user."""
+
+    answer: str = Field(
+        description="The answer from the user",
+    )
+    tool_call_id: str = Field(
+        description="Tool call id",
+    )
+
+
+class UserQuestion(BaseModel):
+    """A question asked by the user."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    question: str = Field(
+        description="The question asked by the user",
+    )
+    tool_call_id: str = Field(
+        description="Tool call id",
+    )
+    result: Future[UserAnswer] = Field(
+        description="Future that will contain the user's answer"
+    )
+
+
+class AgentRuntimeOptions(BaseModel):
+    """User interface options for agents."""
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    interactive_mode: bool = Field(
+        default=True,
+        description="Whether agents can interact with users (ask questions, etc.)",
+    )
+
+    working_directory: Path = Field(
+        default_factory=lambda: Path.cwd(),
+        description="Working directory for agent operations",
+    )
+
+    is_tui_context: bool = Field(
+        default=False,
+        description="Whether the agent is running in TUI context",
+    )
+
+    max_iterations: int = Field(
+        default=10,
+        ge=1,
+        le=100,
+        description="Maximum number of iterations for agent loops",
+    )
+
+    queue: Queue[UserQuestion] = Field(
+        default_factory=Queue,
+        description="Queue for storing user responses",
+    )
+
+    tasks: list[Future[UserAnswer]] = Field(
+        default_factory=list,
+        description="Tasks for storing deferred tool results",
+    )
+
+
+class FileOperationType(str, Enum):
+    """Types of file operations that can be tracked."""
+
+    CREATED = "created"
+    UPDATED = "updated"
+    DELETED = "deleted"
+
+
+class FileOperation(BaseModel):
+    """Single file operation record."""
+
+    file_path: str = Field(
+        description="Full absolute path to the file",
+    )
+    operation: FileOperationType = Field(
+        description="Type of operation performed",
+    )
+    timestamp: datetime = Field(
+        default_factory=datetime.now,
+        description="When the operation occurred",
+    )
+
+
+class FileOperationTracker(BaseModel):
+    """Tracks file operations during a single agent run."""
+
+    operations: list[FileOperation] = Field(
+        default_factory=list,
+        description="List of file operations performed",
+    )
+
+    def add_operation(
+        self, file_path: Path | str, operation: FileOperationType
+    ) -> None:
+        """Record a file operation.
+
+        Args:
+            file_path: Path to the file (will be converted to absolute)
+            operation: Type of operation performed
+        """
+        # Convert to absolute path string
+        if isinstance(file_path, Path):
+            absolute_path = str(file_path.resolve())
+        else:
+            absolute_path = str(Path(file_path).resolve())
+
+        self.operations.append(
+            FileOperation(file_path=absolute_path, operation=operation)
+        )
+
+    def clear(self) -> None:
+        """Clear all tracked operations for a new run."""
+        self.operations = []
+
+    def get_summary(self) -> dict[FileOperationType, list[str]]:
+        """Get operations grouped by type.
+
+        Returns:
+            Dictionary mapping operation types to lists of file paths
+        """
+        summary: dict[FileOperationType, list[str]] = {
+            FileOperationType.CREATED: [],
+            FileOperationType.UPDATED: [],
+            FileOperationType.DELETED: [],
+        }
+
+        for op in self.operations:
+            summary[op.operation].append(op.file_path)
+
+        # Remove duplicates while preserving order
+        for op_type in summary:
+            seen = set()
+            unique_paths = []
+            for path in summary[op_type]:
+                if path not in seen:
+                    seen.add(path)
+                    unique_paths.append(path)
+            summary[op_type] = unique_paths
+
+        return summary
+
+    def format_summary(self) -> str:
+        """Generate human-readable summary for the user.
+
+        Returns:
+            Formatted string showing files modified during the run
+        """
+        if not self.operations:
+            return "No files were modified during this run."
+
+        summary = self.get_summary()
+        lines = ["Files modified during this run:"]
+
+        if summary[FileOperationType.CREATED]:
+            lines.append("\nCreated:")
+            for path in summary[FileOperationType.CREATED]:
+                lines.append(f"  - {path}")
+
+        if summary[FileOperationType.UPDATED]:
+            lines.append("\nUpdated:")
+            for path in summary[FileOperationType.UPDATED]:
+                lines.append(f"  - {path}")
+
+        if summary[FileOperationType.DELETED]:
+            lines.append("\nDeleted:")
+            for path in summary[FileOperationType.DELETED]:
+                lines.append(f"  - {path}")
+
+        return "\n".join(lines)
+
+    def get_display_path(self) -> str | None:
+        """Get a single file path or common parent directory for display.
+
+        Returns:
+            Path string to display, or None if no files were modified
+        """
+        if not self.operations:
+            return None
+
+        unique_paths = list({op.file_path for op in self.operations})
+
+        if len(unique_paths) == 1:
+            # Single file - return its path
+            return unique_paths[0]
+
+        # Multiple files - find common parent directory
+        common_path = os.path.commonpath(unique_paths)
+        return common_path
+
+
+class AgentDeps(AgentRuntimeOptions):
+    """Dependencies passed to all agents for configuration and runtime behavior."""
+
+    llm_model: ModelConfig = Field(
+        description="Model configuration with token limits and provider info",
+    )
+
+    codebase_service: "CodebaseService" = Field(
+        description="Codebase service for code analysis tools",
+    )
+
+    system_prompt_fn: Callable[[RunContext["AgentDeps"]], str] = Field(
+        description="Function that generates the system prompt for this agent",
+    )
+
+    file_tracker: FileOperationTracker = Field(
+        default_factory=FileOperationTracker,
+        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.codebase.service import CodebaseService
+
+    AgentDeps.model_rebuild()
+except ImportError:
+    # Services may not be available in all contexts
+    pass

shotgun/agents/plan.py

@@ -0,0 +1,98 @@
+"""Plan 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_plan_agent(
+    agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
+) -> tuple[Agent[AgentDeps, str | DeferredToolRequests], AgentDeps]:
+    """Create a plan agent with artifact 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 planning tasks, Agent dependencies)
+    """
+    logger.debug("Initializing plan agent")
+    # Use partial to create system prompt function for plan agent
+    system_prompt_fn = partial(build_agent_system_prompt, "plan")
+
+    agent, deps = create_base_agent(
+        system_prompt_fn,
+        agent_runtime_options,
+        load_codebase_understanding_tools=True,
+        additional_tools=None,
+        provider=provider,
+        agent_mode=AgentType.PLAN,
+    )
+    return agent, deps
+
+
+async def run_plan_agent(
+    agent: Agent[AgentDeps, str | DeferredToolRequests],
+    goal: str,
+    deps: AgentDeps,
+    message_history: list[ModelMessage] | None = None,
+) -> AgentRunResult[str | DeferredToolRequests]:
+    """Create or update a plan based on the given goal using artifacts.
+
+    Args:
+        agent: The configured plan agent
+        goal: The planning goal or instruction
+        deps: Agent dependencies
+        message_history: Optional message history for conversation continuity
+
+    Returns:
+        AgentRunResult containing the planning process output
+    """
+    logger.debug("📋 Starting planning for goal: %s", goal)
+
+    # Simple prompt - the agent system prompt has all the artifact instructions
+    full_prompt = f"Create a comprehensive plan for: {goal}"
+
+    try:
+        # Create usage limits for responsible API usage
+        usage_limits = create_usage_limits()
+
+        message_history = await add_system_status_message(deps, message_history)
+
+        result = await run_agent(
+            agent=agent,
+            prompt=full_prompt,
+            deps=deps,
+            message_history=message_history,
+            usage_limits=usage_limits,
+        )
+
+        logger.debug("✅ Planning completed successfully")
+        return result
+
+    except Exception as e:
+        import traceback
+
+        logger.error("Full traceback:\n%s", traceback.format_exc())
+        logger.error("❌ Planning failed: %s", str(e))
+        raise

shotgun/agents/research.py

@@ -0,0 +1,108 @@
+"""Research 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
+from .tools import get_available_web_search_tools
+
+logger = get_logger(__name__)
+
+
+def create_research_agent(
+    agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
+) -> tuple[Agent[AgentDeps, str | DeferredToolRequests], AgentDeps]:
+    """Create a research agent with web search and artifact 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 research tasks, Agent dependencies)
+    """
+    logger.debug("Initializing research agent")
+
+    # Get available web search tools based on configured API keys
+    web_search_tools = get_available_web_search_tools()
+    if web_search_tools:
+        logger.info(
+            "Research agent configured with %d web search tool(s)",
+            len(web_search_tools),
+        )
+    else:
+        logger.warning("Research agent configured without web search tools")
+
+    # Use partial to create system prompt function for research agent
+    system_prompt_fn = partial(build_agent_system_prompt, "research")
+
+    agent, deps = create_base_agent(
+        system_prompt_fn,
+        agent_runtime_options,
+        load_codebase_understanding_tools=True,
+        additional_tools=web_search_tools,
+        provider=provider,
+        agent_mode=AgentType.RESEARCH,
+    )
+    return agent, deps
+
+
+async def run_research_agent(
+    agent: Agent[AgentDeps, str | DeferredToolRequests],
+    query: str,
+    deps: AgentDeps,
+    message_history: list[ModelMessage] | None = None,
+) -> AgentRunResult[str | DeferredToolRequests]:
+    """Perform research on the given query and update research artifacts.
+
+    Args:
+        agent: The configured research agent
+        query: The research query to investigate
+        deps: Agent dependencies
+
+    Returns:
+        Summary of research findings
+    """
+    logger.debug("🔬 Starting research for query: %s", query)
+
+    message_history = await add_system_status_message(deps, message_history)
+
+    try:
+        # Create usage limits for responsible API usage
+        usage_limits = create_usage_limits()
+
+        result = await run_agent(
+            agent=agent,
+            prompt=query,
+            deps=deps,
+            message_history=message_history,
+            usage_limits=usage_limits,
+        )
+
+        logger.debug("✅ Research completed successfully")
+        return result
+
+    except Exception as e:
+        import traceback
+
+        logger.error("Full traceback:\n%s", traceback.format_exc())
+        logger.error("❌ Research failed: %s", str(e))
+        raise

shotgun/agents/specify.py

@@ -0,0 +1,98 @@
+"""Specify 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_specify_agent(
+    agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
+) -> tuple[Agent[AgentDeps, str | DeferredToolRequests], AgentDeps]:
+    """Create a specify agent with artifact 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 specification tasks, Agent dependencies)
+    """
+    logger.debug("Initializing specify agent")
+    # Use partial to create system prompt function for specify agent
+    system_prompt_fn = partial(build_agent_system_prompt, "specify")
+
+    agent, deps = create_base_agent(
+        system_prompt_fn,
+        agent_runtime_options,
+        load_codebase_understanding_tools=True,
+        additional_tools=None,
+        provider=provider,
+        agent_mode=AgentType.SPECIFY,
+    )
+    return agent, deps
+
+
+async def run_specify_agent(
+    agent: Agent[AgentDeps, str | DeferredToolRequests],
+    requirement: str,
+    deps: AgentDeps,
+    message_history: list[ModelMessage] | None = None,
+) -> AgentRunResult[str | DeferredToolRequests]:
+    """Create or update specifications based on the given requirement.
+
+    Args:
+        agent: The configured specify agent
+        requirement: The specification requirement or instruction
+        deps: Agent dependencies
+        message_history: Optional message history for conversation continuity
+
+    Returns:
+        AgentRunResult containing the specification process output
+    """
+    logger.debug("📋 Starting specification for requirement: %s", requirement)
+
+    # Simple prompt - the agent system prompt has all the artifact instructions
+    full_prompt = f"Create a comprehensive specification for: {requirement}"
+
+    try:
+        # Create usage limits for responsible API usage
+        usage_limits = create_usage_limits()
+
+        message_history = await add_system_status_message(deps, message_history)
+
+        result = await run_agent(
+            agent=agent,
+            prompt=full_prompt,
+            deps=deps,
+            message_history=message_history,
+            usage_limits=usage_limits,
+        )
+
+        logger.debug("✅ Specification completed successfully")
+        return result
+
+    except Exception as e:
+        import traceback
+
+        logger.error("Full traceback:\n%s", traceback.format_exc())
+        logger.error("❌ Specification failed: %s", str(e))
+        raise

shotgun/agents/tasks.py

@@ -0,0 +1,96 @@
+"""Tasks 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_tasks_agent(
+    agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
+) -> tuple[Agent[AgentDeps, str | DeferredToolRequests], AgentDeps]:
+    """Create a tasks 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 task management, Agent dependencies)
+    """
+    logger.debug("Initializing tasks agent")
+    # Use partial to create system prompt function for tasks agent
+    system_prompt_fn = partial(build_agent_system_prompt, "tasks")
+
+    agent, deps = create_base_agent(
+        system_prompt_fn,
+        agent_runtime_options,
+        provider=provider,
+        agent_mode=AgentType.TASKS,
+    )
+    return agent, deps
+
+
+async def run_tasks_agent(
+    agent: Agent[AgentDeps, str | DeferredToolRequests],
+    instruction: str,
+    deps: AgentDeps,
+    message_history: list[ModelMessage] | None = None,
+) -> AgentRunResult[str | DeferredToolRequests]:
+    """Create or update tasks based on the given instruction.
+
+    Args:
+        agent: The configured tasks agent
+        instruction: The task creation/update instruction
+        deps: Agent dependencies
+        message_history: Optional message history for conversation continuity
+
+    Returns:
+        AgentRunResult containing the task creation process output
+    """
+    logger.debug("📋 Starting task creation for instruction: %s", instruction)
+
+    message_history = await add_system_status_message(deps, message_history)
+
+    # Let the agent use its tools to read existing tasks, plan, and research
+    full_prompt = f"Create or update tasks 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("✅ Task creation completed successfully")
+        return result
+
+    except Exception as e:
+        import traceback
+
+        logger.error("Full traceback:\n%s", traceback.format_exc())
+        logger.error("❌ Task creation failed: %s", str(e))
+        raise

shotgun/agents/tools/__init__.py

@@ -0,0 +1,34 @@
+"""Tools package for Pydantic AI agents."""
+
+from .codebase import (
+    codebase_shell,
+    directory_lister,
+    file_read,
+    query_graph,
+    retrieve_code,
+)
+from .file_management import append_file, read_file, write_file
+from .user_interaction import ask_user
+from .web_search import (
+    anthropic_web_search_tool,
+    gemini_web_search_tool,
+    get_available_web_search_tools,
+    openai_web_search_tool,
+)
+
+__all__ = [
+    "openai_web_search_tool",
+    "anthropic_web_search_tool",
+    "gemini_web_search_tool",
+    "get_available_web_search_tools",
+    "ask_user",
+    "read_file",
+    "write_file",
+    "append_file",
+    # Codebase understanding tools
+    "query_graph",
+    "retrieve_code",
+    "file_read",
+    "directory_lister",
+    "codebase_shell",
+]