aiagent.py

import requests
import os
import json

TOOL_SPECS = [
    {
        "type": "function",
        "function": {
            "name": "read_file",
            "description": "Read the contents of a file",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "The file path to read"
                    }
                },
                "required": ["path"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "write_file",
            "description": "Write content to a file",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "The file path to write"
                    },
                    "content": {
                        "type": "string",
                        "description": "The content to write to the file"
                    }
                },
                "required": ["path", "content"]
            }
        }
    },
    {
        "type": "function",
        "function": {
            "name": "list_files",
            "description": "List files in a directory",
            "parameters": {
                "type": "object",
                "properties": {
                    "directory": {
                        "type": "string",
                        "description": "The directory to list files from"
                    }
                },
                "required": ["directory"]
            }
        }
    }
]


def read_file(path):
    """Read the contents of a file"""
    try:
        with open(path, 'r') as f:
            content = f.read()
        return content
    except Exception as e:
        return f"Error reading file: {str(e)}"


def write_file(path, content):
    """Write content to a file"""
    try:
        with open(path, 'w') as f:
            f.write(content)
        return "File successfully written"
    except Exception as e:
        return f"Error writing file: {str(e)}"


def list_files(directory):
    """List files in a directory"""
    files = os.listdir(directory)
    return files


def handle_tool(tool_call):
    """Execute a single tool call and return the result"""
    tool_name = tool_call['function']['name']
    tool_args = json.loads(tool_call['function']['arguments'])
    
    print(f"[Executing {tool_name}...]")
    
    if tool_name == 'read_file':
        result = read_file(**tool_args)
    elif tool_name == 'write_file':
        result = write_file(**tool_args)
    elif tool_name == 'list_files':
        result = list_files(**tool_args)
    else:
        result = f"Unknown tool: {tool_name}"
    
    # Ensure result is always a string
    if not isinstance(result, str):
        result = str(result)
    
    return {
        "role": "tool",
        "tool_call_id": tool_call['id'],
        "content": result
    }


def llm(messages):
    api_key = "your_api_key"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    }
    data = {
        "model": "gpt-3.5-turbo",
        "messages": messages,
        "tools": TOOL_SPECS,
        "tool_choice": "auto"
    }
    url = "https://api.openai.com/v1/chat/completions"
    try: 
        response = requests.post(url, json=data, headers=headers)
        message = response.json()["choices"][0]['message']
        # print("Full message")
        # print(message)
        return message
    except:
        print("error")
        raise Exception("bad")


def fake_ai(messages):
    latest_user_message = messages[-1]["content"]
    return f"AI: You said {latest_user_message}... so insightful "


print("Press q to quit")
messages = []

def handle_message(messages, ai_message_obj):
    # Check if AI wants to use tools
    if 'tool_calls' in ai_message_obj and ai_message_obj['tool_calls']:
        # Add AI message with tool calls
        messages.append(ai_message_obj)
        
        # Execute each tool and add results
        for tool_call in ai_message_obj['tool_calls']:
            tool_result = handle_tool(tool_call)
            messages.append(tool_result)
        
        # Get final response from AI
        final_response = llm(messages)
        if final_response['content'] is not None:
            print(f"AI: {final_response['content']}")
        if 'tool_calls' in final_response and ai_message_obj['tool_calls']:
            handle_message(messages, final_response)
    else:
        print(f"AI: {ai_message_obj['content']}")
        messages.append(ai_message_obj)
        return


while True:
    user_message = input("You: ")
    if user_message == 'q':
        break

    messages.append({
        "role": "user",
        "content": user_message
    })

    ai_message_obj = llm(messages)

    handle_message(messages, ai_message_obj)