shotgun-sh 0.3.3.dev1__py3-none-any.whl → 0.6.2__py3-none-any.whl

shotgun/agents/messages.py

@@ -1,12 +1,13 @@
 """Custom message types for Shotgun agents.
 
-This module defines specialized SystemPromptPart subclasses to distinguish
-between different types of system prompts in the agent pipeline.
+This module defines specialized message part subclasses to distinguish
+between different types of prompts in the agent pipeline.
 """
 
 from dataclasses import dataclass, field
+from typing import Literal
 
-from pydantic_ai.messages import SystemPromptPart
+from pydantic_ai.messages import SystemPromptPart, UserPromptPart
 
 from shotgun.agents.models import AgentType
 
@@ -33,3 +34,14 @@ class SystemStatusPrompt(SystemPromptPart):
     """
 
     prompt_type: str = "status"
+
+
+@dataclass
+class InternalPromptPart(UserPromptPart):
+    """User prompt that is system-generated rather than actual user input.
+
+    Used for internal continuation prompts like file resume messages.
+    These should be hidden from the UI but preserved in agent history for context.
+    """
+
+    part_kind: Literal["internal-prompt"] = "internal-prompt"  # type: ignore[assignment]

shotgun/agents/models.py

@@ -1,12 +1,12 @@
 """Pydantic models for agent dependencies and configuration."""
 
 import os
-from asyncio import Future, Queue
+from asyncio import Event, Future, Queue
 from collections.abc import Callable
 from datetime import datetime
 from enum import StrEnum
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeAlias
 
 from pydantic import BaseModel, ConfigDict, Field
 from pydantic_ai import RunContext
@@ -16,9 +16,57 @@ from shotgun.agents.usage_manager import SessionUsageManager, get_session_usage_
 from .config.models import ModelConfig
 
 if TYPE_CHECKING:
+    from pydantic_ai import Agent
+
+    from shotgun.agents.router.models import RouterDeps
     from shotgun.codebase.service import CodebaseService
 
 
+class SubAgentContext(BaseModel):
+    """
+    Context passed to sub-agents so they know they're being orchestrated.
+
+    When sub-agents receive this context, they should:
+    - Be more concise (router handles user communication)
+    - Focus on their specific task
+    - Return structured results
+    """
+
+    is_router_delegated: bool = Field(
+        default=True, description="Always True when passed to sub-agent"
+    )
+    plan_goal: str = Field(
+        default="", description="High-level goal from execution plan"
+    )
+    current_step_id: str | None = Field(
+        default=None, description="ID of the current execution step"
+    )
+    current_step_title: str | None = Field(
+        default=None, description="Title of the current execution step"
+    )
+
+
+class AgentSystemPromptContext(BaseModel):
+    """Context passed to agent system prompt templates.
+
+    This model standardizes the context variables passed to Jinja2 templates
+    when rendering agent system prompts. Using a model makes it easier to
+    test template context construction and ensures type safety.
+    """
+
+    interactive_mode: bool = Field(
+        description="Whether the agent is running in interactive mode"
+    )
+    mode: str = Field(description="The agent type (research, specify, plan, etc.)")
+    sub_agent_context: SubAgentContext | None = Field(
+        default=None, description="Context when running as a sub-agent of the router"
+    )
+    router_mode: str | None = Field(
+        default=None,
+        description="Router mode value (planning/drafting) if router agent",
+    )
+
+
 class AgentResponse(BaseModel):
     """Structured response from an agent with optional clarifying questions.
 
@@ -39,6 +87,23 @@ class AgentResponse(BaseModel):
 Optional list of clarifying questions to ask the user.
 - Single question: Shown as a non-blocking suggestion (user can answer or continue with other prompts)
 - Multiple questions (2+): Asked sequentially in Q&A mode (blocks input until all answered or cancelled)
+""",
+    )
+    files_found: list[str] | None = Field(
+        default=None,
+        description="""
+Optional list of absolute file paths found by the agent.
+Used by FileReadAgent to return paths of files it searched and found.
+The delegation tool can then load these files as multimodal content.
+""",
+    )
+    file_requests: list[str] | None = Field(
+        default=None,
+        description="""
+Optional list of file paths the agent wants to read.
+When set, the agent loop exits, files are loaded as BinaryContent,
+and the loop resumes with file content in the next prompt.
+Use this for PDFs, images, or other binary files you need to analyze.
 """,
     )
 
@@ -51,6 +116,8 @@ class AgentType(StrEnum):
     PLAN = "plan"
     TASKS = "tasks"
     EXPORT = "export"
+    ROUTER = "router"
+    FILE_READ = "file_read"
 
 
 class PipelineConfigEntry(BaseModel):
@@ -319,6 +386,16 @@ class AgentDeps(AgentRuntimeOptions):
         description="Current agent mode for file scoping",
     )
 
+    sub_agent_context: SubAgentContext | None = Field(
+        default=None,
+        description="Context when agent is delegated to by router",
+    )
+
+    cancellation_event: Event | None = Field(
+        default=None,
+        description="Event set when the operation should be cancelled",
+    )
+
 
 # Rebuild model to resolve forward references after imports are available
 try:
@@ -328,3 +405,14 @@ try:
 except ImportError:
     # Services may not be available in all contexts
     pass
+
+
+# Type alias for the standard agent type used throughout the codebase
+ShotgunAgent: TypeAlias = "Agent[AgentDeps, AgentResponse]"
+
+# Type alias for router agent (uses RouterDeps which extends AgentDeps)
+# Note: Agent is contravariant in deps, so RouterAgent is NOT a subtype of ShotgunAgent
+RouterAgent: TypeAlias = "Agent[RouterDeps, AgentResponse]"
+
+# Union type for any agent type (used in AgentManager)
+AnyAgent: TypeAlias = "ShotgunAgent | RouterAgent"

shotgun/agents/plan.py

@@ -2,16 +2,15 @@
 
 from functools import partial
 
-from pydantic_ai import (
-    Agent,
-)
 from pydantic_ai.agent import AgentRunResult
 from pydantic_ai.messages import ModelMessage
 
 from shotgun.agents.config import ProviderType
+from shotgun.agents.models import ShotgunAgent
 from shotgun.logging_config import get_logger
 
 from .common import (
+    EventStreamHandler,
     add_system_status_message,
     build_agent_system_prompt,
     create_base_agent,
@@ -25,7 +24,7 @@ logger = get_logger(__name__)
 
 async def create_plan_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
-) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
+) -> tuple[ShotgunAgent, AgentDeps]:
     """Create a plan agent with artifact management capabilities.
 
     Args:
@@ -51,26 +50,25 @@ async def create_plan_agent(
 
 
 async def run_plan_agent(
-    agent: Agent[AgentDeps, AgentResponse],
-    goal: str,
+    agent: ShotgunAgent,
+    prompt: str,
     deps: AgentDeps,
     message_history: list[ModelMessage] | None = None,
+    event_stream_handler: EventStreamHandler | None = None,
 ) -> AgentRunResult[AgentResponse]:
-    """Create or update a plan based on the given goal using artifacts.
+    """Create or update a plan based on the given prompt using artifacts.
 
     Args:
         agent: The configured plan agent
-        goal: The planning goal or instruction
+        prompt: The planning prompt or instruction
         deps: Agent dependencies
         message_history: Optional message history for conversation continuity
+        event_stream_handler: Optional callback for streaming events
 
     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}"
+    logger.debug("📋 Starting planning for prompt: %s", prompt)
 
     try:
         # Create usage limits for responsible API usage
@@ -80,10 +78,11 @@ async def run_plan_agent(
 
         result = await run_agent(
             agent=agent,
-            prompt=full_prompt,
+            prompt=prompt,
             deps=deps,
             message_history=message_history,
             usage_limits=usage_limits,
+            event_stream_handler=event_stream_handler,
         )
 
         logger.debug("✅ Planning completed successfully")

shotgun/agents/research.py

@@ -2,18 +2,17 @@
 
 from functools import partial
 
-from pydantic_ai import (
-    Agent,
-)
 from pydantic_ai.agent import AgentRunResult
 from pydantic_ai.messages import (
     ModelMessage,
 )
 
 from shotgun.agents.config import ProviderType
+from shotgun.agents.models import ShotgunAgent
 from shotgun.logging_config import get_logger
 
 from .common import (
+    EventStreamHandler,
     add_system_status_message,
     build_agent_system_prompt,
     create_base_agent,
@@ -28,7 +27,7 @@ logger = get_logger(__name__)
 
 async def create_research_agent(
     agent_runtime_options: AgentRuntimeOptions, provider: ProviderType | None = None
-) -> tuple[Agent[AgentDeps, AgentResponse], AgentDeps]:
+) -> tuple[ShotgunAgent, AgentDeps]:
     """Create a research agent with web search and artifact management capabilities.
 
     Args:
@@ -65,22 +64,25 @@ async def create_research_agent(
 
 
 async def run_research_agent(
-    agent: Agent[AgentDeps, AgentResponse],
-    query: str,
+    agent: ShotgunAgent,
+    prompt: str,
     deps: AgentDeps,
     message_history: list[ModelMessage] | None = None,
+    event_stream_handler: EventStreamHandler | None = None,
 ) -> AgentRunResult[AgentResponse]:
-    """Perform research on the given query and update research artifacts.
+    """Perform research on the given prompt and update research artifacts.
 
     Args:
         agent: The configured research agent
-        query: The research query to investigate
+        prompt: The research prompt to investigate
         deps: Agent dependencies
+        message_history: Optional message history for conversation continuity
+        event_stream_handler: Optional callback for streaming events
 
     Returns:
         Summary of research findings
     """
-    logger.debug("🔬 Starting research for query: %s", query)
+    logger.debug("🔬 Starting research for prompt: %s", prompt)
 
     message_history = await add_system_status_message(deps, message_history)
 
@@ -90,10 +92,11 @@ async def run_research_agent(
 
         result = await run_agent(
             agent=agent,
-            prompt=query,
+            prompt=prompt,
             deps=deps,
             message_history=message_history,
             usage_limits=usage_limits,
+            event_stream_handler=event_stream_handler,
         )
 
         logger.debug("✅ Research completed successfully")

shotgun/agents/router/__init__.py

@@ -0,0 +1,47 @@
+"""Router Agent - The intelligent orchestrator for shotgun agents."""
+
+from shotgun.agents.router.models import (
+    CascadeScope,
+    CreatePlanInput,
+    DelegationInput,
+    DelegationResult,
+    ExecutionPlan,
+    ExecutionStep,
+    ExecutionStepInput,
+    MarkStepDoneInput,
+    PlanApprovalStatus,
+    RemoveStepInput,
+    RouterDeps,
+    RouterMode,
+    StepCheckpointAction,
+    SubAgentResult,
+    SubAgentResultStatus,
+    ToolResult,
+)
+from shotgun.agents.router.router import create_router_agent, run_router_agent
+
+__all__ = [
+    # Agent factory
+    "create_router_agent",
+    "run_router_agent",
+    # Enums
+    "RouterMode",
+    "PlanApprovalStatus",
+    "StepCheckpointAction",
+    "CascadeScope",
+    "SubAgentResultStatus",
+    # Plan models
+    "ExecutionStep",
+    "ExecutionPlan",
+    # Tool I/O models
+    "ExecutionStepInput",
+    "CreatePlanInput",
+    "MarkStepDoneInput",
+    "RemoveStepInput",
+    "DelegationInput",
+    "ToolResult",
+    "DelegationResult",
+    "SubAgentResult",
+    # Deps
+    "RouterDeps",
+]