shotgun-sh 0.1.0.dev12__py3-none-any.whl → 0.1.0.dev13__py3-none-any.whl

shotgun/agents/common.py

@@ -2,7 +2,6 @@
 
 import asyncio
 from collections.abc import Callable
-from pathlib import Path
 from typing import Any
 
 from pydantic_ai import (
@@ -15,7 +14,9 @@ from pydantic_ai import (
 from pydantic_ai.agent import AgentRunResult
 from pydantic_ai.messages import (
     ModelMessage,
+    ModelRequest,
     ModelResponse,
+    SystemPromptPart,
     TextPart,
 )
 
@@ -38,6 +39,14 @@ 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,
+)
 
 logger = get_logger(__name__)
 
@@ -45,70 +54,6 @@ logger = get_logger(__name__)
 prompt_loader = PromptLoader()
 
 
-def ensure_file_exists(filename: str, header: str) -> str:
-    """Ensure a markdown file exists with proper header and return its content.
-
-    Args:
-        filename: Name of the file (e.g., "research.md")
-        header: Header to add if file is empty (e.g., "# Research")
-
-    Returns:
-        Current file content
-    """
-    shotgun_dir = Path.cwd() / ".shotgun"
-    file_path = shotgun_dir / filename
-
-    try:
-        if file_path.exists():
-            content = file_path.read_text(encoding="utf-8")
-            if not content.strip():
-                # File exists but is empty, add header
-                header_content = f"{header}\n\n"
-                file_path.write_text(header_content, encoding="utf-8")
-                return header_content
-            return content
-        else:
-            # File doesn't exist, create it with header
-            shotgun_dir.mkdir(exist_ok=True)
-            header_content = f"{header}\n\n"
-            file_path.write_text(header_content, encoding="utf-8")
-            return header_content
-    except Exception as e:
-        logger.error("Failed to initialize %s: %s", filename, str(e))
-        return f"{header}\n\n"
-
-
-def register_common_tools(
-    agent: Agent[AgentDeps], additional_tools: list[Any], interactive_mode: bool
-) -> None:
-    """Register common tools with an agent.
-
-    Args:
-        agent: The Pydantic AI agent to register tools with
-        additional_tools: List of additional tools specific to this agent
-        interactive_mode: Whether to register interactive tools
-    """
-    logger.debug("📌 Registering tools with agent")
-
-    # Register additional tools first (agent-specific)
-    for tool in additional_tools:
-        agent.tool_plain(tool)
-
-    # Register interactive tool if enabled
-    if interactive_mode:
-        agent.tool(ask_user)
-        logger.debug("📞 User interaction tool registered")
-    else:
-        logger.debug("🚫 User interaction disabled (non-interactive mode)")
-
-    # Register common file management tools
-    agent.tool_plain(read_file)
-    agent.tool_plain(write_file)
-    agent.tool_plain(append_file)
-
-    logger.debug("✅ Tool registration complete")
-
-
 async def add_system_status_message(
     deps: AgentDeps,
     message_history: list[ModelMessage] | None = None,
@@ -128,7 +73,6 @@ async def add_system_status_message(
     system_state = prompt_loader.render(
         "agents/state/system_state.j2",
         codebase_understanding_graphs=codebase_understanding_graphs,
-        context="system state",
     )
     message_history.append(
         ModelResponse(
@@ -179,6 +123,7 @@ def create_base_agent(
             **agent_runtime_options.model_dump(),
             llm_model=model_config,
             codebase_service=codebase_service,
+            system_prompt_fn=system_prompt_fn,
         )
 
     except Exception as e:
@@ -194,8 +139,9 @@ def create_base_agent(
         history_processors=[token_limit_compactor],
     )
 
-    # Decorate the system prompt function
-    agent.system_prompt(system_prompt_fn)
+    # System prompt function is stored in deps and will be called manually in run_agent
+    func_name = getattr(system_prompt_fn, "__name__", str(system_prompt_fn))
+    logger.debug("🔧 System prompt function stored: %s", func_name)
 
     # Register additional tools first (agent-specific)
     for tool in additional_tools or []:
@@ -211,7 +157,15 @@ def create_base_agent(
     agent.tool_plain(write_file)
     agent.tool_plain(append_file)
 
-    # Register codebase understanding tools (always available)
+    # Register artifact management tools (always available)
+    agent.tool_plain(create_artifact)
+    agent.tool_plain(list_artifacts)
+    agent.tool_plain(list_artifact_templates)
+    agent.tool_plain(read_artifact)
+    agent.tool_plain(read_artifact_section)
+    agent.tool_plain(write_artifact_section)
+
+    # Register codebase understanding tools (conditional)
     if load_codebase_understanding_tools:
         agent.tool(query_graph)
         agent.tool(retrieve_code)
@@ -222,10 +176,47 @@ def create_base_agent(
     else:
         logger.debug("🚫🧠 Codebase understanding tools not registered")
 
-    logger.debug("✅ Agent creation complete")
+    logger.debug("✅ Agent creation complete with artifact and codebase tools")
     return agent, deps
 
 
+def build_agent_system_prompt(
+    agent_type: str,
+    ctx: RunContext[AgentDeps],
+    context_name: str | None = None,
+) -> str:
+    """Build system prompt for any agent type.
+
+    Args:
+        agent_type: Type of agent ('research', 'plan', 'tasks')
+        ctx: RunContext containing AgentDeps
+        context_name: Optional context name for template rendering
+
+    Returns:
+        Rendered system prompt
+    """
+    prompt_loader = PromptLoader()
+
+    # Add logging if research agent
+    if agent_type == "research":
+        logger.debug("🔧 Building research agent system prompt...")
+        logger.debug("Interactive mode: %s", ctx.deps.interactive_mode)
+
+    result = prompt_loader.render(
+        f"agents/{agent_type}.j2",
+        interactive_mode=ctx.deps.interactive_mode,
+        mode=agent_type,
+    )
+
+    if agent_type == "research":
+        logger.debug(
+            "✅ Research system prompt built successfully (length: %d chars)",
+            len(result),
+        )
+
+    return result
+
+
 def create_usage_limits() -> UsageLimits:
     """Create reasonable usage limits for agent runs.
 
@@ -238,20 +229,41 @@ def create_usage_limits() -> UsageLimits:
     )
 
 
-def get_file_history(filename: str) -> str:
-    """Get the history content from a file.
+async def add_system_prompt_message(
+    deps: AgentDeps,
+    message_history: list[ModelMessage] | None = None,
+) -> list[ModelMessage]:
+    """Add the system prompt as the first message in the message history.
 
     Args:
-        filename: Name of the file (e.g., "research.md")
+        deps: Agent dependencies containing system_prompt_fn
+        message_history: Existing message history
 
     Returns:
-        File content or fallback message
+        Updated message history with system prompt prepended as first message
     """
-    try:
-        return read_file(filename)
-    except Exception as e:
-        logger.debug("Could not load %s history: %s", filename, str(e))
-        return f"No {filename.replace('.md', '')} history available."
+    message_history = message_history or []
+
+    # Create a minimal RunContext to call the system prompt function
+    # We'll pass None for model and usage since they're not used by our system prompt functions
+    context = type(
+        "RunContext", (), {"deps": deps, "retry": 0, "model": None, "usage": None}
+    )()
+
+    # Render the system prompt using the stored function
+    system_prompt_content = deps.system_prompt_fn(context)
+    logger.debug(
+        "🎯 Rendered system prompt (length: %d chars)", len(system_prompt_content)
+    )
+
+    # Create system message and prepend to message history
+    system_message = ModelRequest(
+        parts=[SystemPromptPart(content=system_prompt_content)]
+    )
+    message_history.insert(0, system_message)
+    logger.debug("✅ System prompt prepended as first message")
+
+    return message_history
 
 
 async def run_agent(
@@ -261,6 +273,9 @@ async def run_agent(
     message_history: list[ModelMessage] | None = None,
     usage_limits: UsageLimits | None = None,
 ) -> AgentRunResult[str | DeferredToolRequests]:
+    # Add system prompt as first message
+    message_history = await add_system_prompt_message(deps, message_history)
+
     result = await agent.run(
         prompt,
         deps=deps,

shotgun/agents/config/constants.py

@@ -0,0 +1,18 @@
+"""Configuration constants for Shotgun agents."""
+
+# Field names
+API_KEY_FIELD = "api_key"
+MODEL_NAME_FIELD = "model_name"
+DEFAULT_PROVIDER_FIELD = "default_provider"
+USER_ID_FIELD = "user_id"
+CONFIG_VERSION_FIELD = "config_version"
+
+# Provider names (for consistency with data dict keys)
+OPENAI_PROVIDER = "openai"
+ANTHROPIC_PROVIDER = "anthropic"
+GOOGLE_PROVIDER = "google"
+
+# Environment variable names
+OPENAI_API_KEY_ENV = "OPENAI_API_KEY"
+ANTHROPIC_API_KEY_ENV = "ANTHROPIC_API_KEY"
+GEMINI_API_KEY_ENV = "GEMINI_API_KEY"

shotgun/agents/config/manager.py

@@ -1,6 +1,7 @@
 """Configuration manager for Shotgun CLI."""
 
 import json
+import os
 import uuid
 from pathlib import Path
 from typing import Any
@@ -10,6 +11,15 @@ from pydantic import SecretStr
 from shotgun.logging_config import get_logger
 from shotgun.utils import get_shotgun_home
 
+from .constants import (
+    ANTHROPIC_API_KEY_ENV,
+    ANTHROPIC_PROVIDER,
+    API_KEY_FIELD,
+    GEMINI_API_KEY_ENV,
+    GOOGLE_PROVIDER,
+    OPENAI_API_KEY_ENV,
+    OPENAI_PROVIDER,
+)
 from .models import ProviderType, ShotgunConfig
 
 logger = get_logger(__name__)
@@ -58,6 +68,22 @@ class ConfigManager:
 
             self._config = ShotgunConfig.model_validate(data)
             logger.debug("Configuration loaded successfully from %s", self.config_path)
+
+            # Check if the default provider has a key, if not find one that does
+            if not self.has_provider_key(self._config.default_provider):
+                original_default = self._config.default_provider
+                # Find first provider with a configured key
+                for provider in ProviderType:
+                    if self.has_provider_key(provider):
+                        logger.info(
+                            "Default provider %s has no API key, updating to %s",
+                            original_default.value,
+                            provider.value,
+                        )
+                        self._config.default_provider = provider
+                        self.save(self._config)
+                        break
+
             return self._config
 
         except Exception as e:
@@ -114,17 +140,25 @@ class ConfigManager:
         provider_config = self._get_provider_config(config, provider_enum)
 
         # Only support api_key updates
-        if "api_key" in kwargs:
-            api_key_value = kwargs["api_key"]
+        if API_KEY_FIELD in kwargs:
+            api_key_value = kwargs[API_KEY_FIELD]
             provider_config.api_key = (
                 SecretStr(api_key_value) if api_key_value is not None else None
             )
 
         # Reject other fields
-        unsupported_fields = set(kwargs.keys()) - {"api_key"}
+        unsupported_fields = set(kwargs.keys()) - {API_KEY_FIELD}
         if unsupported_fields:
             raise ValueError(f"Unsupported configuration fields: {unsupported_fields}")
 
+        # If no other providers have keys configured and we just added one,
+        # set this provider as the default
+        if API_KEY_FIELD in kwargs and api_key_value is not None:
+            other_providers = [p for p in ProviderType if p != provider_enum]
+            has_other_keys = any(self.has_provider_key(p) for p in other_providers)
+            if not has_other_keys:
+                config.default_provider = provider_enum
+
         self.save(config)
 
     def clear_provider_key(self, provider: ProviderType | str) -> None:
@@ -136,11 +170,27 @@ class ConfigManager:
         self.save(config)
 
     def has_provider_key(self, provider: ProviderType | str) -> bool:
-        """Check if the given provider has a non-empty API key configured."""
+        """Check if the given provider has a non-empty API key configured.
+
+        This checks both the configuration file and environment variables.
+        """
         config = self.load()
         provider_enum = self._ensure_provider_enum(provider)
         provider_config = self._get_provider_config(config, provider_enum)
-        return self._provider_has_api_key(provider_config)
+
+        # Check config first
+        if self._provider_has_api_key(provider_config):
+            return True
+
+        # Check environment variable
+        if provider_enum == ProviderType.OPENAI:
+            return bool(os.getenv(OPENAI_API_KEY_ENV))
+        elif provider_enum == ProviderType.ANTHROPIC:
+            return bool(os.getenv(ANTHROPIC_API_KEY_ENV))
+        elif provider_enum == ProviderType.GOOGLE:
+            return bool(os.getenv(GEMINI_API_KEY_ENV))
+
+        return False
 
     def has_any_provider_key(self) -> bool:
         """Determine whether any provider has a configured API key."""
@@ -175,25 +225,27 @@ class ConfigManager:
 
     def _convert_secrets_to_secretstr(self, data: dict[str, Any]) -> None:
         """Convert plain text secrets in data to SecretStr objects."""
-        for provider in ["openai", "anthropic", "google"]:
+        for provider in [OPENAI_PROVIDER, ANTHROPIC_PROVIDER, GOOGLE_PROVIDER]:
             if provider in data and isinstance(data[provider], dict):
                 if (
-                    "api_key" in data[provider]
-                    and data[provider]["api_key"] is not None
+                    API_KEY_FIELD in data[provider]
+                    and data[provider][API_KEY_FIELD] is not None
                 ):
-                    data[provider]["api_key"] = SecretStr(data[provider]["api_key"])
+                    data[provider][API_KEY_FIELD] = SecretStr(
+                        data[provider][API_KEY_FIELD]
+                    )
 
     def _convert_secretstr_to_plain(self, data: dict[str, Any]) -> None:
         """Convert SecretStr objects in data to plain text for JSON serialization."""
-        for provider in ["openai", "anthropic", "google"]:
+        for provider in [OPENAI_PROVIDER, ANTHROPIC_PROVIDER, GOOGLE_PROVIDER]:
             if provider in data and isinstance(data[provider], dict):
                 if (
-                    "api_key" in data[provider]
-                    and data[provider]["api_key"] is not None
+                    API_KEY_FIELD in data[provider]
+                    and data[provider][API_KEY_FIELD] is not None
                 ):
-                    if hasattr(data[provider]["api_key"], "get_secret_value"):
-                        data[provider]["api_key"] = data[provider][
-                            "api_key"
+                    if hasattr(data[provider][API_KEY_FIELD], "get_secret_value"):
+                        data[provider][API_KEY_FIELD] = data[provider][
+                            API_KEY_FIELD
                         ].get_secret_value()
 
     def _ensure_provider_enum(self, provider: ProviderType | str) -> ProviderType:
@@ -216,7 +268,7 @@ class ConfigManager:
 
     def _provider_has_api_key(self, provider_config: Any) -> bool:
         """Return True if the provider config contains a usable API key."""
-        api_key = getattr(provider_config, "api_key", None)
+        api_key = getattr(provider_config, API_KEY_FIELD, None)
         if api_key is None:
             return False
 

shotgun/agents/config/provider.py

@@ -13,6 +13,11 @@ from pydantic_ai.providers.openai import OpenAIProvider
 
 from shotgun.logging_config import get_logger
 
+from .constants import (
+    ANTHROPIC_API_KEY_ENV,
+    GEMINI_API_KEY_ENV,
+    OPENAI_API_KEY_ENV,
+)
 from .manager import get_config_manager
 from .models import MODEL_SPECS, ModelConfig, ProviderType
 
@@ -86,10 +91,10 @@ def get_provider_model(provider: ProviderType | None = None) -> ModelConfig:
     )
 
     if provider_enum == ProviderType.OPENAI:
-        api_key = _get_api_key(config.openai.api_key, "OPENAI_API_KEY")
+        api_key = _get_api_key(config.openai.api_key, OPENAI_API_KEY_ENV)
         if not api_key:
             raise ValueError(
-                "OpenAI API key not configured. Set via environment variable OPENAI_API_KEY or config."
+                f"OpenAI API key not configured. Set via environment variable {OPENAI_API_KEY_ENV} or config."
             )
 
         # Get model spec
@@ -108,10 +113,10 @@ def get_provider_model(provider: ProviderType | None = None) -> ModelConfig:
         )
 
     elif provider_enum == ProviderType.ANTHROPIC:
-        api_key = _get_api_key(config.anthropic.api_key, "ANTHROPIC_API_KEY")
+        api_key = _get_api_key(config.anthropic.api_key, ANTHROPIC_API_KEY_ENV)
         if not api_key:
             raise ValueError(
-                "Anthropic API key not configured. Set via environment variable ANTHROPIC_API_KEY or config."
+                f"Anthropic API key not configured. Set via environment variable {ANTHROPIC_API_KEY_ENV} or config."
             )
 
         # Get model spec
@@ -130,10 +135,10 @@ def get_provider_model(provider: ProviderType | None = None) -> ModelConfig:
         )
 
     elif provider_enum == ProviderType.GOOGLE:
-        api_key = _get_api_key(config.google.api_key, "GEMINI_API_KEY")
+        api_key = _get_api_key(config.google.api_key, GEMINI_API_KEY_ENV)
         if not api_key:
             raise ValueError(
-                "Gemini API key not configured. Set via environment variable GEMINI_API_KEY or config."
+                f"Gemini API key not configured. Set via environment variable {GEMINI_API_KEY_ENV} or config."
             )
 
         # Get model spec

shotgun/agents/models.py

@@ -1,10 +1,12 @@
 """Pydantic models for agent dependencies and configuration."""
 
 from asyncio import Future, Queue
+from collections.abc import Callable
 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
 
@@ -83,6 +85,10 @@ class AgentDeps(AgentRuntimeOptions):
         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",
+    )
+
 
 # Rebuild model to resolve forward references after imports are available
 try:

shotgun/agents/plan.py

@@ -1,51 +1,33 @@
 """Plan agent factory and functions using Pydantic AI with file-based memory."""
 
+from functools import partial
+
 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,
+    build_agent_system_prompt,
     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.
+    """Create a plan agent with artifact management capabilities.
 
     Args:
         agent_runtime_options: Agent runtime options for the agent
@@ -55,8 +37,15 @@ def create_plan_agent(
         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(
-        _build_plan_agent_system_prompt, agent_runtime_options, provider=provider
+        system_prompt_fn,
+        agent_runtime_options,
+        load_codebase_understanding_tools=True,
+        additional_tools=None,
+        provider=provider,
     )
     return agent, deps
 
@@ -67,7 +56,7 @@ async def run_plan_agent(
     deps: AgentDeps,
     message_history: list[ModelMessage] | None = None,
 ) -> AgentRunResult[str | DeferredToolRequests]:
-    """Create or update a plan based on the given goal.
+    """Create or update a plan based on the given goal using artifacts.
 
     Args:
         agent: The configured plan agent
@@ -80,11 +69,9 @@ async def run_plan_agent(
     """
     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
+    # 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()
@@ -108,12 +95,3 @@ async def run_plan_agent(
         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")