llm agent using openai api

agent.py

import subprocess
from openai.types.chat import (
    ChatCompletionMessageToolCallUnion,
    ChatCompletionMessageFunctionToolCall,
    ChatCompletionMessageParam,
)
import argparse
import os
import sys
import json

from openai import OpenAI

API_KEY = os.getenv("OPENROUTER_API_KEY")
BASE_URL = os.getenv("OPENROUTER_BASE_URL", default="https://openrouter.ai/api/v1")


def write_tool(filepath: str, content: str):
    try:
        with open(filepath, "w+") as fd:
            fd.write(content)
        return f"{filepath} written successfully"
    except Exception as e:
        return f"Could not write to {filepath}; error: {str(e)}"


def read_tool(filepath: str):
    with open(filepath) as fd:
        return fd.read()


def bash_tool(command: str):
    process = subprocess.run(command, shell=True, capture_output=True)
    output = json.dumps(
        {
            "stderr": process.stderr.decode(),
            "stdout": process.stdout.decode(),
            "returncode": process.returncode,
        },
    )
    return output


write_tool_spec: ChatCompletionMessageFunctionToolCall = {
    "type": "function",
    "function": {
        "name": "Write",
        "description": "Write contents to a file",
        "parameters": {
            "type": "object",
            "properties": {
                "file_path": {
                    "type": "string",
                    "description": "The path to the file",
                },
                "content": {
                    "type": "string",
                    "description": "content to be written to the file",
                },
            },
            "required": ["file_path", "content"],
        },
    },
}
read_tool_spec: ChatCompletionMessageFunctionToolCall = {
    "type": "function",
    "function": {
        "name": "Read",
        "description": "Read and return the contents of a file",
        "parameters": {
            "type": "object",
            "properties": {
                "file_path": {
                    "type": "string",
                    "description": "The path to the file to read",
                }
            },
            "required": ["file_path"],
        },
    },
}
bash_tool_spec: ChatCompletionMessageFunctionToolCall = {
    "type": "function",
    "function": {
        "name": "Bash",
        "description": "Execute a shell command",
        "parameters": {
            "type": "object",
            "required": ["command"],
            "properties": {
                "command": {
                    "type": "string",
                    "description": "The command to execute",
                }
            },
        },
    },
}


def main():
    p = argparse.ArgumentParser()
    p.add_argument("-p", required=True)
    args = p.parse_args()

    if not API_KEY:
        raise RuntimeError("OPENROUTER_API_KEY is not set")

    client = OpenAI(api_key=API_KEY, base_url=BASE_URL)

    tools_map = {"Read": read_tool, "Write": write_tool, "Bash": bash_tool}
    messages: list[ChatCompletionMessageParam] = [{"role": "user", "content": args.p}]
    tools: list[ChatCompletionMessageToolCallUnion] = [
        read_tool_spec,
        write_tool_spec,
        bash_tool_spec,
    ]

    while True:
        chat = client.chat.completions.create(
            model="anthropic/claude-haiku-4.5",
            messages=messages,
            tools=tools,  # ty:ignore[invalid-argument-type]
        )

        # Break out if there is no response
        if not chat.choices or len(chat.choices) == 0:
            raise RuntimeError("no choices in response")

        tool_calls: list[ChatCompletionMessageToolCallUnion] | None = chat.choices[
            0
        ].message.tool_calls

        # We need to record the LLM's response in messages as well !

        # log("Adding llm response back to `messages` !")
        messages.append(chat.choices[0].message.model_dump())

        if tool_calls is None or len(tool_calls) == 0:
            print(chat.choices[0].message.content)
            break
        else:
            # tool call exists
            for tool_call in tool_calls:
                if tool_call.type == "function":
                    log("Got a function tool call")
                    func_name = tool_call.function.name
                    func_args = json.loads(tool_call.function.arguments)
                    func = tools_map.get(func_name)
                    if not func:
                        raise RuntimeError(f"tool call not found {func_name=}")

                    log(f"[TOOL_CALL]: {func_name=} {func_args=}")

                    result = ""
                    if func_name == "Read":
                        file_path = func_args.get("file_path")
                        result = func(file_path)
                    elif func_name == "Write":
                        file_path = func_args.get("file_path")
                        content = func_args.get("content")
                        result = func(file_path, content)
                    elif func_name == "Bash":
                        command = func_args.get("command")
                        result = func(command)

                    messages.append(
                        {
                            "role": "tool",  # pyright ignore
                            "tool_call_id": tool_call.id,
                            "content": result,
                        }
                    )
                    # log("Added tool output to `messages` !")
                else:
                    tool_call_type = tool_call.type
                    tool_call = tool_call.model_dump_json()
                    raise RuntimeError(
                        f"Expected function tool call, found {tool_call_type} tool call {tool_call=}"
                    )


def log(s: str):
    # You can use print statements as follows for debugging, they'll be visible when running tests.
    print(s, file=sys.stderr)


if __name__ == "__main__":
    main()