shotgun-sh 0.1.0.dev1__py3-none-any.whl

shotgun/agents/models.py

@@ -0,0 +1,94 @@
+"""Pydantic models for agent dependencies and configuration."""
+
+from asyncio import Future, Queue
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from .config.models import ModelConfig
+
+if TYPE_CHECKING:
+    from shotgun.codebase.service import CodebaseService
+
+
+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",
+    )
+
+    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 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",
+    )
+
+
+# Rebuild model to resolve forward references after imports are available
+try:
+    from shotgun.codebase.service import CodebaseService
+
+    AgentDeps.model_rebuild()
+except ImportError:
+    # CodebaseService may not be available in all contexts
+    pass

shotgun/agents/plan.py

@@ -0,0 +1,119 @@
+"""Plan agent factory and functions using Pydantic AI with file-based memory."""
+
+from pydantic_ai import (
+    Agent,
+    DeferredToolRequests,
+    RunContext,
+)
+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 shotgun.prompts import PromptLoader
+
+from .common import (
+    add_system_status_message,
+    create_base_agent,
+    create_usage_limits,
+    ensure_file_exists,
+    get_file_history,
+    run_agent,
+)
+from .models import AgentDeps, AgentRuntimeOptions
+
+logger = get_logger(__name__)
+
+# Global prompt loader instance
+prompt_loader = PromptLoader()
+
+
+def _build_plan_agent_system_prompt(ctx: RunContext[AgentDeps]) -> str:
+    """Build the system prompt for the plan agent.
+
+    Args:
+        ctx: RunContext containing AgentDeps with interactive_mode and other settings
+
+    Returns:
+        The complete system prompt string for the plan agent
+    """
+    return prompt_loader.render(
+        "agents/plan.j2", interactive_mode=ctx.deps.interactive_mode, context="plans"
+    )
+
+
+def create_plan_agent(
+    agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
+) -> tuple[Agent[AgentDeps, str | DeferredToolRequests], AgentDeps]:
+    """Create a plan 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 planning tasks, Agent dependencies)
+    """
+    logger.debug("Initializing plan agent")
+    agent, deps = create_base_agent(
+        _build_plan_agent_system_prompt, agent_runtime_options, provider=provider
+    )
+    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.
+
+    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)
+
+    # Ensure plan.md exists
+    ensure_file_exists("plan.md", "# Plan")
+
+    # Let the agent use its tools to read existing plan and research
+    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
+
+
+def get_plan_history() -> str:
+    """Get the full plan history from the file.
+
+    Returns:
+        Plan history content or fallback message
+    """
+    return get_file_history("plan.md")

shotgun/agents/research.py

@@ -0,0 +1,131 @@
+"""Research agent factory and functions using Pydantic AI with file-based memory."""
+
+from pydantic_ai import (
+    Agent,
+    DeferredToolRequests,
+    RunContext,
+)
+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 shotgun.prompts import PromptLoader
+
+from .common import (
+    add_system_status_message,
+    create_base_agent,
+    create_usage_limits,
+    ensure_file_exists,
+    get_file_history,
+    run_agent,
+)
+from .models import AgentDeps, AgentRuntimeOptions
+from .tools import web_search_tool
+
+logger = get_logger(__name__)
+
+# Global prompt loader instance
+prompt_loader = PromptLoader()
+
+
+def _build_research_agent_system_prompt(ctx: RunContext[AgentDeps]) -> str:
+    """Build the system prompt for the research agent.
+
+    Args:
+        ctx: RunContext containing AgentDeps with interactive_mode and other settings
+
+    Returns:
+        The complete system prompt string for the research agent
+    """
+    return prompt_loader.render(
+        "agents/research.j2",
+        interactive_mode=ctx.deps.interactive_mode,
+        context="research output",
+    )
+
+
+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 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")
+    agent, deps = create_base_agent(
+        _build_research_agent_system_prompt,
+        agent_runtime_options,
+        load_codebase_understanding_tools=True,
+        additional_tools=[web_search_tool],
+        provider=provider,
+    )
+    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 the research file.
+
+    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)
+
+    # Ensure research.md exists
+    ensure_file_exists("research.md", "# Research")
+
+    message_history = await add_system_status_message(deps, message_history)
+
+    user_prompt = prompt_loader.render(
+        "user/research.j2",
+        user_query=query,
+        context="research output",
+    )
+
+    try:
+        # Create usage limits for responsible API usage
+        usage_limits = create_usage_limits()
+
+        result = await run_agent(
+            agent=agent,
+            prompt=user_prompt,
+            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
+
+
+def get_research_history() -> str:
+    """Get the full research history from the file.
+
+    Returns:
+        Research history content or fallback message
+    """
+    return get_file_history("research.md")

shotgun/agents/tasks.py

@@ -0,0 +1,122 @@
+"""Tasks agent factory and functions using Pydantic AI with file-based memory."""
+
+from pydantic_ai import (
+    Agent,
+    DeferredToolRequests,
+    RunContext,
+)
+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 shotgun.prompts import PromptLoader
+
+from .common import (
+    add_system_status_message,
+    create_base_agent,
+    create_usage_limits,
+    ensure_file_exists,
+    get_file_history,
+    run_agent,
+)
+from .models import AgentDeps, AgentRuntimeOptions
+
+logger = get_logger(__name__)
+
+# Global prompt loader instance
+prompt_loader = PromptLoader()
+
+
+def _build_tasks_agent_system_prompt(ctx: RunContext[AgentDeps]) -> str:
+    """Build the system prompt for the tasks agent.
+
+    Args:
+        ctx: RunContext containing AgentDeps with interactive_mode and other settings
+
+    Returns:
+        The complete system prompt string for the tasks agent
+    """
+    return prompt_loader.render(
+        "agents/tasks.j2",
+        interactive_mode=ctx.deps.interactive_mode,
+        context="task lists",
+    )
+
+
+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")
+    agent, deps = create_base_agent(
+        _build_tasks_agent_system_prompt, agent_runtime_options, provider=provider
+    )
+    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)
+
+    # Ensure tasks.md exists
+    ensure_file_exists("tasks.md", "# Tasks")
+
+    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
+
+
+def get_tasks_history() -> str:
+    """Get the full tasks history from the file.
+
+    Returns:
+        Tasks history content or fallback message
+    """
+    return get_file_history("tasks.md")

shotgun/agents/tools/__init__.py

@@ -0,0 +1,26 @@
+"""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 web_search_tool
+
+__all__ = [
+    "web_search_tool",
+    "ask_user",
+    "read_file",
+    "write_file",
+    "append_file",
+    # Codebase understanding tools
+    "query_graph",
+    "retrieve_code",
+    "file_read",
+    "directory_lister",
+    "codebase_shell",
+]

shotgun/agents/tools/codebase/__init__.py

@@ -0,0 +1,28 @@
+"""Codebase understanding tools for Pydantic AI agents."""
+
+from .codebase_shell import codebase_shell
+from .directory_lister import directory_lister
+from .file_read import file_read
+from .models import (
+    CodeSnippetResult,
+    DirectoryListResult,
+    FileReadResult,
+    QueryGraphResult,
+    ShellCommandResult,
+)
+from .query_graph import query_graph
+from .retrieve_code import retrieve_code
+
+__all__ = [
+    "query_graph",
+    "retrieve_code",
+    "file_read",
+    "directory_lister",
+    "codebase_shell",
+    # Result models
+    "QueryGraphResult",
+    "CodeSnippetResult",
+    "FileReadResult",
+    "DirectoryListResult",
+    "ShellCommandResult",
+]