400 line Code Editing Agent in Python

claude-agent.py

"""
This is a python adoptation of Thorsten Ball's `How To build an Agent` blogpost, all credits goes to Thorsten
https://ampcode.com/how-to-build-an-agent
"""

import sys
import json
from anthropic import Anthropic
from anthropic.types import MessageParam, Message
from typing import List, Dict, Callable, Optional, Any, Type
from pydantic import BaseModel, Field
import pathlib


# Define the input schema using Pydantic
class ReadFileInput(BaseModel):
    path: str = Field(
        ..., description="The relative path of a file in the working directory."
    )


# Define the actual function that implements the tool
def read_file(input_data: Dict[str, Any]) -> str:
    # Validate input using the Pydantic model
    try:
        validated_input = ReadFileInput(**input_data)
        file_path = pathlib.Path(validated_input.path)

        # Basic security check: prevent escaping the current directory
        # This is rudimentary; real applications need more robust checks.
        if not file_path.resolve().is_relative_to(pathlib.Path.cwd().resolve()):
            raise ValueError("File path is outside the allowed directory.")
        if not file_path.is_file():
            raise FileNotFoundError(
                f"File not found or is a directory: {validated_input.path}"
            )

        content = file_path.read_text()
        return content
    except FileNotFoundError:
        # Return error message that the LLM can understand
        return f"Error: File not found at path: {input_data.get('path', 'N/A')}"
    except Exception as e:
        # General error handling
        return f"Error reading file: {e}"


# Define a class to hold tool information
class ToolDefinition:
    def __init__(
        self,
        name: str,
        description: str,
        input_schema: Type[BaseModel],
        function: Callable,
    ):
        self.name = name
        self.description = description
        self.input_schema = input_schema
        self.function = function

    # Helper to generate the format Anthropic expects for the 'tools' parameter
    def get_anthropic_schema(self) -> Dict[str, Any]:
        # Pydantic v2+ uses model_json_schema()
        schema = self.input_schema.model_json_schema()
        # We need to remove the 'title' key if pydantic adds it,
        # and adjust the structure slightly for Anthropic
        return {
            "name": self.name,
            "description": self.description,
            "input_schema": {
                "type": "object",
                "properties": schema.get("properties", {}),
                "required": schema.get("required", []),
            },
        }


# Define the input schema using Pydantic
class EditFileInput(BaseModel):
    path: str = Field(..., description="The relative path to the file to edit.")
    old_str: str = Field(
        ...,
        description="The exact text content to search for in the file. Use an empty string to insert at the beginning or create a new file.",
    )
    new_str: str = Field(
        ...,
        description="The text content to replace old_str with. Cannot be the same as old_str unless old_str is empty.",
    )

    # Add a validator
    from pydantic import model_validator

    @model_validator(mode="after")
    def check_old_new_diff(self) -> "EditFileInput":
        if self.old_str != "" and self.old_str == self.new_str:
            raise ValueError(
                "old_str and new_str must be different if old_str is not empty"
            )
        return self


# Helper function to create a file with directories if needed
def create_new_file(file_path: pathlib.Path, content: str) -> str:
    try:
        # Create parent directories if they don't exist
        file_path.parent.mkdir(parents=True, exist_ok=True)
        file_path.write_text(content)
        return f"Successfully created file {file_path}"
    except Exception as e:
        raise OSError(f"Failed to create file {file_path}: {e}") from e


# Define the actual function that implements the tool
def edit_file(input_data: Dict[str, Any]) -> str:
    try:
        validated_input = EditFileInput(**input_data)
        file_path = pathlib.Path(validated_input.path)

        # Security check
        if not file_path.resolve().is_relative_to(pathlib.Path.cwd().resolve()):
            raise ValueError("File path is outside the allowed directory.")

        if file_path.is_dir():
            raise ValueError(
                f"Path points to a directory, not a file: {validated_input.path}"
            )

        if validated_input.old_str == "":
            # Create new file or prepend
            if file_path.exists():
                # Prepend content
                existing_content = file_path.read_text()
                new_content = validated_input.new_str + existing_content
                file_path.write_text(new_content)
                return f"OK - Prepended content to {validated_input.path}"
            else:
                # Create new file
                return create_new_file(file_path, validated_input.new_str)
        else:
            # Replace existing content
            if not file_path.exists():
                return f"Error: File not found at path: {validated_input.path} (cannot replace content)"

            old_content = file_path.read_text()
            # Use replace with count=1 to replace only the first occurrence?
            # The original Go code used -1 (replace all). Let's stick to that.
            # For more robust editing, diff/patch libraries or line-based editing would be better.
            if validated_input.old_str not in old_content:
                return f"Error: old_str '{validated_input.old_str[:50]}...' not found in the file."

            new_content = old_content.replace(
                validated_input.old_str, validated_input.new_str
            )

            if old_content == new_content:
                # This case should ideally be caught by Pydantic validator if old_str != "",
                # but double check doesn't hurt.
                return "Warning: No changes made. old_str might not have been found or new_str is identical."

            file_path.write_text(new_content)
            return "OK - File edited successfully."

    except Exception as e:
        return f"Error editing file: {e}"


# Create the ToolDefinition instance for edit_file
EditFileDefinition = ToolDefinition(
    name="edit_file",
    description="""Make edits to a text file by replacing content.

Replaces the first occurrence of 'old_str' with 'new_str' in the file specified by 'path'.
- 'old_str' and 'new_str' MUST be different unless 'old_str' is empty.
- If 'old_str' is an empty string (""), 'new_str' will be prepended to the file if it exists, or a new file will be created with 'new_str' as content if the file doesn't exist.
- Use with caution. This performs a simple string replacement.
""",
    input_schema=EditFileInput,
    function=edit_file,
)


# Define the input schema using Pydantic
class ListFilesInput(BaseModel):
    path: Optional[str] = Field(
        None,
        description="Optional relative path to list files from. Defaults to current directory if not provided.",
    )


# Define the actual function that implements the tool
def list_files(input_data: Dict[str, Any]) -> str:
    try:
        validated_input = ListFilesInput(**input_data)
        target_path_str = validated_input.path if validated_input.path else "."
        target_path = pathlib.Path(target_path_str).resolve()
        base_path = pathlib.Path.cwd().resolve()

        # Security check
        if not target_path.is_relative_to(base_path):
            raise ValueError("Listing path is outside the allowed directory.")

        if not target_path.is_dir():
            raise ValueError(f"Path is not a directory: {target_path_str}")

        files_list = []
        for item in target_path.iterdir():
            # Get path relative to the *requested* directory for cleaner output
            relative_item_path = item.relative_to(target_path)
            if item.is_dir():
                files_list.append(f"{relative_item_path}/")
            else:
                files_list.append(str(relative_item_path))

        # Return the list as a JSON string, as Claude handles structured data well
        return json.dumps(files_list)

    except Exception as e:
        return f"Error listing files: {e}"


# Create the ToolDefinition instance for list_files
ListFilesDefinition = ToolDefinition(
    name="list_files",
    description="List files and directories at a given relative path. If no path is provided, lists files in the current directory.",
    input_schema=ListFilesInput,
    function=list_files,
)


# Create the ToolDefinition instance for read_file
ReadFileDefinition = ToolDefinition(
    name="read_file",
    description="Read the contents of a given relative file path. Use this when you want to see what's inside a file. Do not use this with directory names.",
    input_schema=ReadFileInput,
    function=read_file,
)


# Helper function to get user input
def get_user_message() -> Optional[str]:
    try:
        return input()
    except EOFError:
        return None


class Agent:
    def __init__(
        self,
        client: Anthropic,
        get_user_message_func: Callable[[], Optional[str]],
        tools: List[ToolDefinition] = None,
    ):
        self.client = client
        self.get_user_message = get_user_message_func
        self.tools = {tool.name: tool for tool in tools} if tools else {}

    def _run_inference(self, conversation: List[MessageParam]) -> Message:
        anthropic_tools = [tool.get_anthropic_schema() for tool in self.tools.values()]
        # print(f"DEBUG: Sending tools: {json.dumps(anthropic_tools, indent=2)}") # Optional debug print
        message = self.client.messages.create(
            model="claude-3-7-sonnet-latest",
            max_tokens=1024,
            messages=conversation,
            tools=anthropic_tools if anthropic_tools else None,
            # Tool choice can be added here if needed, e.g., tool_choice={"type": "auto"}
        )
        # print(f"DEBUG: Received message: {message}") # Optional debug print
        return message

    def _execute_tool(
        self, tool_name: str, tool_input: Dict[str, Any], tool_use_id: str
    ) -> Dict[str, Any]:
        """Executes a tool and returns the result in the format Anthropic expects."""
        tool_result_content = ""
        is_error = False
        if tool_name in self.tools:
            tool_def = self.tools[tool_name]
            # ANSI escape codes for color
            print(f"\u001b[92mtool\u001b[0m: {tool_name}({json.dumps(tool_input)})")
            try:
                # Execute the tool's function
                tool_result_content = tool_def.function(tool_input)
            except Exception as e:
                print(f"Error executing tool {tool_name}: {e}")
                tool_result_content = f"Error executing tool {tool_name}: {e}"
                is_error = True
        else:
            print(f"Error: Tool '{tool_name}' not found.")
            tool_result_content = f"Error: Tool '{tool_name}' not found."
            is_error = True

        # Return the result in the format required for the next API call
        return {
            "type": "tool_result",
            "tool_use_id": tool_use_id,
            "content": str(tool_result_content),  # Content should be a string
            "is_error": is_error,  # Indicate if the tool execution failed
        }

    def run(self):
        conversation: List[MessageParam] = []
        print("Chat with Claude (use 'ctrl-d' or 'ctrl-c' to quit)")

        read_user_input = True
        while True:
            if read_user_input:
                # ANSI escape codes for color
                print("\u001b[94mYou\u001b[0m: ", end="", flush=True)
                user_input = self.get_user_message()
                if user_input is None:
                    print("\nExiting...")
                    break
                user_message: MessageParam = {"role": "user", "content": user_input}
                conversation.append(user_message)

            try:
                # print(f"DEBUG: Sending conversation: {json.dumps(conversation, indent=2)}") # Optional debug
                message = self._run_inference(conversation)

                # Append the assistant's response *before* processing tool calls
                # The SDK's message object might not be directly serializable,
                # so reconstruct the dictionary if needed or store the object
                # For simplicity here, let's reconstruct the dict structure Anthropic uses
                assistant_response_content = []
                if message.content:
                    for block in message.content:
                        assistant_response_content.append(
                            block.model_dump()
                        )  # Use model_dump() for Pydantic objects in SDK

                assistant_message: MessageParam = {
                    "role": message.role,
                    "content": assistant_response_content,
                }
                conversation.append(assistant_message)
                # print(f"DEBUG: Appended assistant message: {json.dumps(assistant_message, indent=2)}") # Optional debug

                tool_results = []
                assistant_said_something = False

                # Process message content (text and tool calls)
                if message.content:
                    for content_block in message.content:
                        if content_block.type == "text":
                            # ANSI escape codes for color
                            print(f"\u001b[93mClaude\u001b[0m: {content_block.text}")
                            assistant_said_something = True
                        elif content_block.type == "tool_use":
                            # It's requesting a tool
                            tool_name = content_block.name
                            tool_input = content_block.input or {}
                            tool_use_id = content_block.id
                            tool_result = self._execute_tool(
                                tool_name, tool_input, tool_use_id
                            )
                            tool_results.append(tool_result)

                # If there were tool calls, send results back
                if tool_results:
                    read_user_input = False  # Let the agent respond to the tool result
                    # Construct the user message containing tool results
                    tool_result_message: MessageParam = {
                        "role": "user",
                        "content": tool_results,  # Send list of tool result blocks
                    }
                    conversation.append(tool_result_message)
                    # print(f"DEBUG: Appended tool results message: {json.dumps(tool_result_message, indent=2)}") # Optional debug
                else:
                    # No tool calls, wait for next user input
                    read_user_input = True
                    if (
                        not assistant_said_something
                        and message.stop_reason == "tool_use"
                    ):
                        # Handle cases where the assistant *only* outputs tool calls and no text
                        # You might want to print a generic message or just continue
                        print("\u001b[93mClaude\u001b[0m: (Thinking...)")
                        pass  # Let the loop continue to process tool results

            except Exception as e:
                print(f"\nAn error occurred: {e}")
                # Decide how to handle errors, e.g., retry, ask user, exit
                # For simplicity, let's just allow the user to type again
                read_user_input = True


def main():
    # Client automatically looks for ANTHROPIC_API_KEY env var
    try:
        client = Anthropic()
    except Exception as e:
        print(f"Error creating Anthropic client: {e}")
        print("Please ensure the ANTHROPIC_API_KEY environment variable is set.")
        sys.exit(1)

    # Define tools 
    tools = [ReadFileDefinition, ListFilesDefinition, EditFileDefinition]
    agent = Agent(client, get_user_message, tools)
    try:
        agent.run()
    except Exception as e:
        print(f"\nError: {e}")


if __name__ == "__main__":
    main()