coding agent

coding_agent.py

# /// script
# dependencies = [
#   "pydantic-ai",
#   "python-dotenv",
#   "termcolor",
# ]
# ///

# To run:
# export OPEN_AI_API_KEY=???
# uv run coding_agent.py

from pathlib import Path
from typing import Any

from dotenv import load_dotenv
from pydantic_ai import Agent
from termcolor import colored

load_dotenv()

agent = Agent("openai:gpt-4o", system_prompt="You are a coding assistant whose goal is to help solve coding tasks.")

ENCODING = "utf-8"
YOU_COLOR = "cyan"
ASSISTANT_COLOR = "yellow"


def _resolve_path(path_str: str, action: str) -> Path:
    if not (path := Path(path_str).expanduser()).is_absolute():
        path = (Path.cwd() / path).resolve()
    print(f"{action}: {path}")
    return path


@agent.tool_plain
def read_file(filename: str) -> dict[str, Any]:
    """Gets the full content of a file provided by the user."""
    full_path = _resolve_path(filename, "Reading")
    return {"file_path": str(full_path), "content": full_path.read_text(encoding=ENCODING)}


@agent.tool_plain
def list_files(path: str = ".") -> dict[str, Any]:
    """Lists the files in a directory provided by the user."""
    full_path = _resolve_path(path, "Listing")
    return {
        "path": str(full_path),
        "files": [{"filename": item.name, "type": "file" if item.is_file() else "dir"} for item in full_path.iterdir()]
    }


@agent.tool_plain
def edit_file(path: str, new_str: str, old_str: str = "") -> dict[str, Any]:
    """
    Replaces first occurrence of old_str with new_str in file.
    If old_str is empty, create/overwrite file with new_str.
    """
    full_path = _resolve_path(path, "Editting")

    if old_str == "":
        full_path.write_text(new_str, encoding=ENCODING)
        action = "created_file"
    elif old_str not in (original := full_path.read_text(encoding=ENCODING)):
        action = "old_str not found"
    else:
        full_path.write_text(original.replace(old_str, new_str, 1), encoding=ENCODING)
        action = "edited"

    print(f"\t{action}: {full_path}")
    return {"path": str(full_path), "action": action}


def main():
    message_history = []
    while True:
        try:
            user_input = input(colored("You: ", YOU_COLOR)).strip()
        except (KeyboardInterrupt, EOFError):
            break
        result = agent.run_sync(user_input, message_history=message_history)
        message_history = result.all_messages()
        print(colored(f"Assistant: {result.output}", ASSISTANT_COLOR))



if __name__ == "__main__":
    main()