shotgun-sh 0.1.0__py3-none-any.whl

shotgun/agents/tools/codebase/models.py

@@ -0,0 +1,252 @@
+"""Pydantic models for codebase tool outputs."""
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class QueryGraphResult(BaseModel):
+    """Result from graph query with formatted string output."""
+
+    success: bool = Field(description="Whether the query was successful")
+    query: str = Field(description="Original natural language query")
+    cypher_query: str | None = Field(None, description="Generated Cypher query")
+    column_names: list[str] = Field(
+        default_factory=list, description="Result column names"
+    )
+    results: list[dict[str, Any]] = Field(
+        default_factory=list, description="Query results"
+    )
+    row_count: int = Field(0, description="Number of result rows")
+    execution_time_ms: float = Field(
+        0.0, description="Query execution time in milliseconds"
+    )
+    error: str | None = Field(default=None, description="Error message if failed")
+
+    def __str__(self) -> str:
+        """Format query result for LLM consumption."""
+        if not self.success:
+            return f"**Query Failed**: {self.error}"
+
+        if not self.results:
+            return f"**No Results Found** for query: {self.query}"
+
+        output_lines = []
+        if self.cypher_query:
+            output_lines.append(f"**Generated Cypher**: `{self.cypher_query}`")
+            output_lines.append("")
+
+        output_lines.append(
+            f"**Results** ({self.row_count} rows, {self.execution_time_ms:.1f}ms):"
+        )
+        output_lines.append("")
+
+        if self.column_names and self.results:
+            # Create markdown table
+            headers = self.column_names
+            output_lines.append("| " + " | ".join(headers) + " |")
+            output_lines.append("|" + "|".join([" --- " for _ in headers]) + "|")
+
+            # Limit to first 10 rows to avoid overwhelming output
+            rows_to_show = min(10, len(self.results))
+            for row in self.results[:rows_to_show]:
+                row_values = []
+                for col in headers:
+                    value = row.get(col, "")
+                    # Convert to string and truncate if too long
+                    str_value = str(value) if value is not None else ""
+                    if len(str_value) > 50:
+                        str_value = str_value[:47] + "..."
+                    row_values.append(str_value)
+                output_lines.append("| " + " | ".join(row_values) + " |")
+
+            if self.row_count > rows_to_show:
+                output_lines.append(
+                    f"... and {self.row_count - rows_to_show} more rows"
+                )
+
+        return "\n".join(output_lines)
+
+
+class CodeSnippetResult(BaseModel):
+    """Result from code retrieval with formatted output."""
+
+    found: bool = Field(description="Whether the code entity was found")
+    qualified_name: str = Field(description="Fully qualified name searched for")
+    file_path: str | None = Field(None, description="Path to source file")
+    line_start: int | None = Field(None, description="Starting line number")
+    line_end: int | None = Field(None, description="Ending line number")
+    source_code: str | None = Field(None, description="Source code content")
+    docstring: str | None = Field(None, description="Docstring if available")
+    language: str = Field(
+        default="", description="Programming language for syntax highlighting"
+    )
+    error: str | None = Field(None, description="Error message if not found")
+
+    def __str__(self) -> str:
+        """Format code snippet for LLM consumption."""
+        if not self.found:
+            error_msg = (
+                self.error or f"Entity '{self.qualified_name}' not found in graph"
+            )
+            return f"**Not Found**: {error_msg}\n\nTry using `query_graph` to search for similar entities or check the qualified name."
+
+        output_lines = []
+        output_lines.append(f"**Qualified Name**: `{self.qualified_name}`")
+
+        if self.file_path:
+            output_lines.append(f"**File**: `{self.file_path}`")
+
+        if self.line_start and self.line_end:
+            output_lines.append(f"**Lines**: {self.line_start}-{self.line_end}")
+
+        if self.docstring:
+            output_lines.append(f"**Docstring**: {self.docstring}")
+
+        if self.source_code:
+            output_lines.append("")
+            output_lines.append("**Source Code**:")
+            language_tag = self.language if self.language else ""
+            output_lines.append(f"```{language_tag}")
+            output_lines.append(self.source_code)
+            output_lines.append("```")
+
+        return "\n".join(output_lines)
+
+
+class FileReadResult(BaseModel):
+    """Result from file reading with content output."""
+
+    success: bool = Field(description="Whether file was read successfully")
+    file_path: str = Field(description="Path to file that was read")
+    content: str | None = Field(None, description="File content")
+    encoding: str = Field("utf-8", description="Encoding used to read file")
+    size_bytes: int = Field(0, description="File size in bytes")
+    language: str = Field(
+        default="", description="Programming language for syntax highlighting"
+    )
+    error: str | None = Field(default=None, description="Error message if failed")
+
+    def __str__(self) -> str:
+        """Return file content or error message."""
+        if not self.success:
+            return f"**Error reading file `{self.file_path}`**: {self.error}"
+
+        output_lines = []
+        output_lines.append(f"**File**: `{self.file_path}`")
+        output_lines.append(f"**Size**: {self.size_bytes} bytes")
+
+        if self.encoding != "utf-8":
+            output_lines.append(f"**Encoding**: {self.encoding}")
+
+        output_lines.append("")
+        output_lines.append("**Content**:")
+        language_tag = self.language if self.language else ""
+        output_lines.append(f"```{language_tag}")
+        output_lines.append(self.content or "")
+        output_lines.append("```")
+
+        return "\n".join(output_lines)
+
+
+class DirectoryListResult(BaseModel):
+    """Result from directory listing with structured output."""
+
+    success: bool = Field(description="Whether directory was listed successfully")
+    directory: str = Field(description="Directory path that was listed")
+    full_path: str = Field(description="Absolute path to directory")
+    directories: list[str] = Field(
+        default_factory=list, description="Subdirectory names"
+    )
+    files: list[tuple[str, int]] = Field(
+        default_factory=list, description="Files as (name, size_bytes) tuples"
+    )
+    error: str | None = Field(default=None, description="Error message if failed")
+
+    def __str__(self) -> str:
+        """Format directory listing for LLM consumption."""
+        if not self.success:
+            return f"**Error listing directory `{self.directory}`**: {self.error}"
+
+        output_lines = []
+        output_lines.append(f"**Directory**: `{self.directory}`")
+        output_lines.append(f"**Full Path**: `{self.full_path}`")
+        output_lines.append("")
+
+        if not self.directories and not self.files:
+            return "\n".join(output_lines + ["Directory is empty"])
+
+        if self.directories:
+            output_lines.append("**Directories**:")
+            for dir_name in self.directories:
+                output_lines.append(f"  📁 {dir_name}/")
+            output_lines.append("")
+
+        if self.files:
+            output_lines.append("**Files**:")
+            for file_name, size_bytes in self.files:
+                if size_bytes < 1024:
+                    size_str = f"{size_bytes}B"
+                elif size_bytes < 1024 * 1024:
+                    size_str = f"{size_bytes / 1024:.1f}KB"
+                else:
+                    size_str = f"{size_bytes / (1024 * 1024):.1f}MB"
+                output_lines.append(f"  📄 {file_name} ({size_str})")
+
+        output_lines.append("")
+        output_lines.append(
+            f"**Total**: {len(self.directories)} directories, {len(self.files)} files"
+        )
+
+        return "\n".join(output_lines)
+
+
+class ShellCommandResult(BaseModel):
+    """Result from shell command execution with formatted output."""
+
+    success: bool = Field(description="Whether command executed without errors")
+    command: str = Field(description="Command that was executed")
+    args: list[str] = Field(default_factory=list, description="Command arguments")
+    stdout: str = Field(default="", description="Standard output")
+    stderr: str = Field(default="", description="Standard error output")
+    return_code: int = Field(default=0, description="Process return code")
+    execution_time_ms: float = Field(
+        default=0.0, description="Execution time in milliseconds"
+    )
+    error: str | None = Field(
+        default=None, description="Error message if execution failed"
+    )
+
+    def __str__(self) -> str:
+        """Format command output for LLM consumption."""
+        if self.error:
+            return f"**Command Failed**: {self.error}"
+
+        output_lines = []
+        cmd_str = f"{self.command} {' '.join(self.args)}".strip()
+        output_lines.append(f"**Command**: `{cmd_str}`")
+        output_lines.append(f"**Execution Time**: {self.execution_time_ms:.1f}ms")
+
+        if self.stdout:
+            output_lines.append("")
+            output_lines.append("**Output**:")
+            output_lines.append("```")
+            output_lines.append(self.stdout.rstrip())
+            output_lines.append("```")
+
+        if self.stderr:
+            output_lines.append("")
+            output_lines.append("**Error Output**:")
+            output_lines.append("```")
+            output_lines.append(self.stderr.rstrip())
+            output_lines.append("```")
+
+        if self.return_code != 0:
+            output_lines.append("")
+            output_lines.append(f"**Exit Code**: {self.return_code}")
+
+        if not self.stdout and not self.stderr:
+            output_lines.append("")
+            output_lines.append("Command executed successfully with no output")
+
+        return "\n".join(output_lines)

shotgun/agents/tools/codebase/query_graph.py

@@ -0,0 +1,67 @@
+"""Query codebase knowledge graph using natural language."""
+
+from pydantic_ai import RunContext
+
+from shotgun.agents.models import AgentDeps
+from shotgun.codebase.models import QueryType
+from shotgun.logging_config import get_logger
+
+from .models import QueryGraphResult
+
+logger = get_logger(__name__)
+
+
+async def query_graph(
+    ctx: RunContext[AgentDeps], graph_id: str, query: str
+) -> QueryGraphResult:
+    """Query codebase knowledge graph using natural language.
+
+    Args:
+        ctx: RunContext containing AgentDeps with codebase service
+        graph_id: Graph ID to query (use the ID, not the name)
+        query: Natural language question about the codebase
+
+    Returns:
+        QueryGraphResult with formatted output via __str__
+    """
+    logger.debug("🔧 Querying graph %s with query: %s", graph_id, query)
+
+    try:
+        if not ctx.deps.codebase_service:
+            return QueryGraphResult(
+                success=False,
+                query=query,
+                error="No codebase indexed",
+            )
+
+        # Execute natural language query
+        result = await ctx.deps.codebase_service.execute_query(
+            graph_id=graph_id,
+            query=query,
+            query_type=QueryType.NATURAL_LANGUAGE,
+        )
+
+        # Create QueryGraphResult from service result
+        graph_result = QueryGraphResult(
+            success=result.success,
+            query=query,
+            cypher_query=result.cypher_query,
+            column_names=result.column_names,
+            results=result.results,
+            row_count=result.row_count,
+            execution_time_ms=result.execution_time_ms,
+            error=result.error,
+        )
+
+        logger.debug(
+            "📄 Query completed: %s with %d results",
+            "success" if graph_result.success else "failed",
+            graph_result.row_count,
+        )
+
+        return graph_result
+
+    except Exception as e:
+        error_msg = f"Error querying graph: {str(e)}"
+        logger.error("❌ Query graph failed: %s", str(e))
+        return QueryGraphResult(success=False, query=query, error=error_msg)

shotgun/agents/tools/codebase/retrieve_code.py

@@ -0,0 +1,81 @@
+"""Retrieve source code by qualified name from codebase."""
+
+from pathlib import Path
+
+from pydantic_ai import RunContext
+
+from shotgun.agents.models import AgentDeps
+from shotgun.codebase.core.code_retrieval import retrieve_code_by_qualified_name
+from shotgun.codebase.core.language_config import get_language_config
+from shotgun.logging_config import get_logger
+
+from .models import CodeSnippetResult
+
+logger = get_logger(__name__)
+
+
+async def retrieve_code(
+    ctx: RunContext[AgentDeps], graph_id: str, qualified_name: str
+) -> CodeSnippetResult:
+    """Get source code by fully qualified name.
+
+    Args:
+        ctx: RunContext containing AgentDeps with codebase service
+        graph_id: Graph ID to query (use the ID, not the name)
+        qualified_name: Fully qualified name like "module.Class.method"
+
+    Returns:
+        CodeSnippetResult with formatted output via __str__
+    """
+    logger.debug("🔧 Retrieving code for: %s in graph %s", qualified_name, graph_id)
+
+    try:
+        if not ctx.deps.codebase_service:
+            return CodeSnippetResult(
+                found=False,
+                qualified_name=qualified_name,
+                error="No codebase indexed",
+            )
+
+        # Use the existing code retrieval functionality
+        code_snippet = await retrieve_code_by_qualified_name(
+            manager=ctx.deps.codebase_service.manager,
+            graph_id=graph_id,
+            qualified_name=qualified_name,
+        )
+
+        # Detect language from file extension
+        language = ""
+        if code_snippet.file_path:
+            file_extension = Path(code_snippet.file_path).suffix
+            language_config = get_language_config(file_extension)
+            if language_config:
+                language = language_config.name
+
+        # Convert to our result model
+        result = CodeSnippetResult(
+            found=code_snippet.found,
+            qualified_name=code_snippet.qualified_name,
+            file_path=code_snippet.file_path if code_snippet.found else None,
+            line_start=code_snippet.line_start if code_snippet.found else None,
+            line_end=code_snippet.line_end if code_snippet.found else None,
+            source_code=code_snippet.source_code if code_snippet.found else None,
+            docstring=code_snippet.docstring,
+            language=language,
+            error=code_snippet.error_message if not code_snippet.found else None,
+        )
+
+        logger.debug(
+            "📄 Retrieved code: %s for %s",
+            "found" if result.found else "not found",
+            qualified_name,
+        )
+
+        return result
+
+    except Exception as e:
+        error_msg = f"Error retrieving code: {str(e)}"
+        logger.error("❌ Retrieve code failed: %s", str(e))
+        return CodeSnippetResult(
+            found=False, qualified_name=qualified_name, error=error_msg
+        )

shotgun/agents/tools/file_management.py

@@ -0,0 +1,218 @@
+"""File management tools for Pydantic AI agents.
+
+These tools are restricted to the .shotgun directory for security.
+"""
+
+from pathlib import Path
+from typing import Literal
+
+from pydantic_ai import RunContext
+
+from shotgun.agents.models import AgentDeps, AgentType, FileOperationType
+from shotgun.logging_config import get_logger
+from shotgun.utils.file_system_utils import get_shotgun_base_path
+
+logger = get_logger(__name__)
+
+# Map agent modes to their allowed directories/files (in workflow order)
+AGENT_DIRECTORIES = {
+    AgentType.RESEARCH: "research.md",
+    AgentType.SPECIFY: "specification.md",
+    AgentType.PLAN: "plan.md",
+    AgentType.TASKS: "tasks.md",
+    AgentType.EXPORT: "*",  # Export agent can write anywhere except protected files
+}
+
+# Files protected from export agent modifications
+PROTECTED_AGENT_FILES = {
+    "research.md",
+    "specification.md",
+    "plan.md",
+    "tasks.md",
+}
+
+
+def _validate_agent_scoped_path(filename: str, agent_mode: AgentType | None) -> Path:
+    """Validate and resolve a file path within the agent's scoped directory.
+
+    Args:
+        filename: Relative filename
+        agent_mode: The current agent mode
+
+    Returns:
+        Absolute path to the file within the agent's scoped directory
+
+    Raises:
+        ValueError: If the path attempts to access files outside the agent's scope
+    """
+    base_path = get_shotgun_base_path()
+
+    if agent_mode and agent_mode in AGENT_DIRECTORIES:
+        # For export mode, allow writing to any file except protected agent files
+        if agent_mode == AgentType.EXPORT:
+            # Check if trying to write to a protected file
+            if filename in PROTECTED_AGENT_FILES:
+                raise ValueError(
+                    f"Export agent cannot write to protected file '{filename}'. "
+                    f"Protected files are: {', '.join(sorted(PROTECTED_AGENT_FILES))}"
+                )
+
+            # Allow writing anywhere else in .shotgun directory
+            full_path = (base_path / filename).resolve()
+        else:
+            # For other agents, only allow writing to their specific file
+            allowed_file = AGENT_DIRECTORIES[agent_mode]
+            if filename != allowed_file:
+                raise ValueError(
+                    f"{agent_mode.value.capitalize()} agent can only write to '{allowed_file}'. "
+                    f"Attempted to write to '{filename}'"
+                )
+            full_path = (base_path / filename).resolve()
+    else:
+        # No agent mode specified, fall back to old validation
+        full_path = (base_path / filename).resolve()
+
+    # Ensure the resolved path is within the .shotgun directory
+    try:
+        full_path.relative_to(base_path.resolve())
+    except ValueError as e:
+        raise ValueError(
+            f"Access denied: Path '{filename}' is outside .shotgun directory"
+        ) from e
+
+    return full_path
+
+
+def _validate_shotgun_path(filename: str) -> Path:
+    """Validate and resolve a file path within the .shotgun directory.
+
+    Args:
+        filename: Relative filename within .shotgun directory
+
+    Returns:
+        Absolute path to the file within .shotgun directory
+
+    Raises:
+        ValueError: If the path attempts to access files outside .shotgun directory
+    """
+    base_path = get_shotgun_base_path()
+
+    # Create the full path
+    full_path = (base_path / filename).resolve()
+
+    # Ensure the resolved path is within the .shotgun directory
+    try:
+        full_path.relative_to(base_path.resolve())
+    except ValueError as e:
+        raise ValueError(
+            f"Access denied: Path '{filename}' is outside .shotgun directory"
+        ) from e
+
+    return full_path
+
+
+async def read_file(ctx: RunContext[AgentDeps], filename: str) -> str:
+    """Read a file from the .shotgun directory.
+
+    Args:
+        filename: Relative path to file within .shotgun directory
+
+    Returns:
+        File contents as string
+
+    Raises:
+        ValueError: If path is outside .shotgun directory
+        FileNotFoundError: If file does not exist
+    """
+    logger.debug("🔧 Reading file: %s", filename)
+
+    try:
+        file_path = _validate_shotgun_path(filename)
+
+        if not file_path.exists():
+            raise FileNotFoundError(f"File not found: {filename}")
+
+        content = file_path.read_text(encoding="utf-8")
+        logger.debug("📄 Read %d characters from %s", len(content), filename)
+        return content
+
+    except Exception as e:
+        error_msg = f"Error reading file '{filename}': {str(e)}"
+        logger.error("❌ File read failed: %s", error_msg)
+        return error_msg
+
+
+async def write_file(
+    ctx: RunContext[AgentDeps],
+    filename: str,
+    content: str,
+    mode: Literal["w", "a"] = "w",
+) -> str:
+    """Write content to a file in the .shotgun directory.
+
+    Args:
+        filename: Relative path to file within .shotgun directory
+        content: Content to write to the file
+        mode: Write mode - 'w' for overwrite, 'a' for append
+
+    Returns:
+        Success message or error message
+
+    Raises:
+        ValueError: If path is outside .shotgun directory or invalid mode
+    """
+    logger.debug("🔧 Writing file: %s (mode: %s)", filename, mode)
+
+    if mode not in ["w", "a"]:
+        raise ValueError(f"Invalid mode '{mode}'. Use 'w' for write or 'a' for append")
+
+    try:
+        # Use agent-scoped validation for write operations
+        file_path = _validate_agent_scoped_path(filename, ctx.deps.agent_mode)
+
+        # Determine operation type
+        if mode == "a":
+            operation = FileOperationType.UPDATED
+        else:
+            operation = (
+                FileOperationType.CREATED
+                if not file_path.exists()
+                else FileOperationType.UPDATED
+            )
+
+        # Ensure parent directory exists
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Write content
+        if mode == "a":
+            with open(file_path, "a", encoding="utf-8") as f:
+                f.write(content)
+            logger.debug("📄 Appended %d characters to %s", len(content), filename)
+            result = f"Successfully appended {len(content)} characters to {filename}"
+        else:
+            file_path.write_text(content, encoding="utf-8")
+            logger.debug("📄 Wrote %d characters to %s", len(content), filename)
+            result = f"Successfully wrote {len(content)} characters to {filename}"
+
+        # Track the file operation
+        ctx.deps.file_tracker.add_operation(file_path, operation)
+
+        return result
+
+    except Exception as e:
+        error_msg = f"Error writing file '{filename}': {str(e)}"
+        logger.error("❌ File write failed: %s", error_msg)
+        return error_msg
+
+
+async def append_file(ctx: RunContext[AgentDeps], filename: str, content: str) -> str:
+    """Append content to a file in the .shotgun directory.
+
+    Args:
+        filename: Relative path to file within .shotgun directory
+        content: Content to append to the file
+
+    Returns:
+        Success message or error message
+    """
+    return await write_file(ctx, filename, content, mode="a")

shotgun/agents/tools/user_interaction.py

@@ -0,0 +1,37 @@
+"""User interaction tools for Pydantic AI agents."""
+
+from asyncio import get_running_loop
+
+from pydantic_ai import CallDeferred, RunContext
+
+from shotgun.agents.models import AgentDeps, UserQuestion
+from shotgun.logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+async def ask_user(ctx: RunContext[AgentDeps], question: str) -> str:
+    """Ask the human a question and return the answer.
+
+
+    Args:
+        question: The question to ask the user with a clear CTA at the end. Needs to be is readable, clear, and easy to understand. Use Markdown formatting. Make key phrases and words stand out.
+
+    Returns:
+        The user's response as a string
+    """
+    tool_call_id = ctx.tool_call_id
+    assert tool_call_id is not None  # noqa: S101
+
+    try:
+        logger.debug("\n👉 %s\n", question)
+        future = get_running_loop().create_future()
+        await ctx.deps.queue.put(
+            UserQuestion(question=question, tool_call_id=tool_call_id, result=future)
+        )
+        ctx.deps.tasks.append(future)
+        raise CallDeferred(question)
+
+    except (EOFError, KeyboardInterrupt):
+        logger.warning("User input interrupted or unavailable")
+        return "User input not available or interrupted"