diff --git a/main.py b/main.py index cfd6b26..b28db37 100644 --- a/main.py +++ b/main.py @@ -1,7 +1,11 @@ -import os import sys import httpx import json +import base64 +import pyautogui +import subprocess +from io import BytesIO +from PIL import Image ANTHROPIC_API_URL = "https://api.anthropic.com/v1/messages" MODEL = "claude-3-7-sonnet-20250219" @@ -32,6 +36,190 @@ def send_request (headers, body): sys.exit(1) +def execute_computer_tool(tool_input): + """Execute computer tool actions like screenshot, click, type, etc.""" + action = tool_input.get("action") + + if action == "screenshot": + # Take a screenshot + screenshot = pyautogui.screenshot() + + # Convert to base64 + buffered = BytesIO() + screenshot.save(buffered, format="PNG") + + # Check size and resize if needed (5MB = 5242880 bytes) + if buffered.tell() > 4242880: + # Reset buffer + buffered.seek(0) + screenshot = Image.open(buffered) + + # Resize to 80% repeatedly until under 5MB + while True: + buffered = BytesIO() + new_size = (int(screenshot.width * 0.8), int(screenshot.height * 0.8)) + screenshot = screenshot.resize(new_size, Image.Resampling.LANCZOS) + screenshot.save(buffered, format="PNG", optimize=True) + if buffered.tell() <= 4242880: + break + + # Convert to base64 + buffered.seek(0) + img_base64 = base64.b64encode(buffered.read()).decode('utf-8') + + return { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": img_base64 + } + } + + elif action == "click": + # Get coordinates + x = tool_input.get("coordinate_x") + y = tool_input.get("coordinate_y") + if x is None or y is None: + return {"type": "text", "text": "Error: Missing coordinates for click action"} + + # Perform click + pyautogui.click(x, y) + return {"type": "text", "text": f"Clicked at coordinates ({x}, {y})"} + + elif action == "double_click": + # Get coordinates + x = tool_input.get("coordinate_x") + y = tool_input.get("coordinate_y") + if x is None or y is None: + return {"type": "text", "text": "Error: Missing coordinates for double click action"} + + # Perform double click + pyautogui.doubleClick(x, y) + return {"type": "text", "text": f"Double-clicked at coordinates ({x}, {y})"} + + elif action == "type": + # Get text to type + text = tool_input.get("text", "") + pyautogui.typewrite(text) + return {"type": "text", "text": f"Typed: {text}"} + + elif action == "key": + # Press a key or key combination + key = tool_input.get("key", "") + try: + pyautogui.press(key) + return {"type": "text", "text": f"Pressed key: {key}"} + except Exception as e: + return {"type": "text", "text": f"Error pressing key {key}: {str(e)}"} + + elif action == "scroll": + # Scroll action + direction = tool_input.get("direction", "down") + amount = tool_input.get("amount", 3) + + scroll_amount = -amount if direction == "up" else amount + pyautogui.scroll(scroll_amount) + return {"type": "text", "text": f"Scrolled {direction} by {amount}"} + + else: + return {"type": "text", "text": f"Unknown computer action: {action}"} + + +def execute_bash_tool(command): + """Execute bash commands safely.""" + try: + # Execute the command + result = subprocess.run( + command, + shell=True, + capture_output=True, + text=True, + timeout=30 # 30 second timeout + ) + + output = result.stdout + if result.stderr: + output += f"\n[STDERR]: {result.stderr}" + + return {"type": "text", "text": output if output else "(Command executed successfully with no output)"} + + except subprocess.TimeoutExpired: + return {"type": "text", "text": "Error: Command timed out after 30 seconds"} + except Exception as e: + return {"type": "text", "text": f"Error executing command: {str(e)}"} + + +def execute_text_editor_tool(command, path, **kwargs): + """Execute text editor actions.""" + + if command == "view": + try: + with open(path, 'r') as f: + content = f.read() + return {"type": "text", "text": content} + except Exception as e: + return {"type": "text", "text": f"Error reading file: {str(e)}"} + + elif command == "str_replace": + old_str = kwargs.get("old_str", "") + new_str = kwargs.get("new_str", "") + + try: + with open(path, 'r') as f: + content = f.read() + + if old_str not in content: + return {"type": "text", "text": f"Error: '{old_str}' not found in file"} + + # Replace the string + new_content = content.replace(old_str, new_str) + + # Write back to file + with open(path, 'w') as f: + f.write(new_content) + + return {"type": "text", "text": "String replaced successfully"} + + except Exception as e: + return {"type": "text", "text": f"Error modifying file: {str(e)}"} + + elif command == "create": + content = kwargs.get("content", "") + try: + with open(path, 'w') as f: + f.write(content) + return {"type": "text", "text": f"File created: {path}"} + except Exception as e: + return {"type": "text", "text": f"Error creating file: {str(e)}"} + + else: + return {"type": "text", "text": f"Unknown text editor command: {command}"} + + +def execute_tool(tool_name, tool_input): + """Main function to execute any tool based on its name and input.""" + + print(f"\n[EXECUTING] Tool: {tool_name}") + print(f"[INPUT] {json.dumps(tool_input, indent=2)}") + + if tool_name == "computer": + # Pass the entire tool_input to execute_computer_tool + return execute_computer_tool(tool_input) + + elif tool_name == "bash": + command = tool_input.get("command", "") + return execute_bash_tool(command) + + elif tool_name == "str_replace_editor": + command = tool_input.get("command") + path = tool_input.get("path", "") + return execute_text_editor_tool(command, path, **tool_input) + + else: + return {"type": "text", "text": f"Unknown tool: {tool_name}"} + + def send_prompt(api_key, messages): headers = HEADERS.copy() headers["x-api-key"] = api_key @@ -66,11 +254,13 @@ def send_prompt(api_key, messages): return send_request(headers, body) -def extract_tool_use(response): +def extract_tool_uses(response): + """Extract all tool use requests from the response.""" + tool_uses = [] for block in response.get("content", []): if block.get("type") == "tool_use": - return block - return None + tool_uses.append(block) + return tool_uses def run(): @@ -88,66 +278,60 @@ def run(): }) print(f"Sending prompt: {prompt}") - first_response = send_prompt(api_key, messages) - - # We can't just append the whole first response, otherwise we get 400 bad req: - # "Extra inputs are not permitted" - messages.append({ - 'role': 'assistant', - 'content': first_response.get('content', []), - }) - - tool_use = extract_tool_use(first_response) - if not tool_use: - print("No tool use required. Response:") - print(first_response["content"]) - return - - # print first response from Claude - for block in first_response.get("content", []): - print(f'{block['type']}:') - if block["type"] == "text": - print(f'\t{block["text"]}') - elif block["type"] == "thinking": - print(f'\t{block["thinking"]}') - elif block["type"] == "tool_use": - print(f'\tTool use requested: {block["name"]}') - print(f'\tInput: {block["input"]}') - - # this is a stub for the tool execution - fake_result_message = { - "role": "user", - "content": [ - { + + # Main conversation loop + while True: + response = send_prompt(api_key, messages) + + # Add Claude's response to messages + messages.append({ + 'role': 'assistant', + 'content': response.get('content', []), + }) + + # Print Claude's response + print("\n[CLAUDE RESPONSE]") + for block in response.get("content", []): + if block["type"] == "text": + print(block["text"]) + elif block["type"] == "thinking": + print(f"[THINKING] {block['thinking']}") + elif block["type"] == "tool_use": + print(f"[TOOL REQUEST] {block['name']} - {block.get('input', {})}") + + # Extract all tool uses + tool_uses = extract_tool_uses(response) + + if not tool_uses: + # No more tool uses, we're done + break + + # Execute all requested tools and collect results + tool_results = [] + for tool_use in tool_uses: + tool_name = tool_use["name"] + tool_input = tool_use.get("input", {}) + tool_id = tool_use["id"] + + # Execute the tool + result_content = execute_tool(tool_name, tool_input) + + # Add to results + tool_results.append({ "type": "tool_result", - "tool_use_id": tool_use["id"], - "content": "(Stub) Executed tool successfully" - } - ] - } - messages.append(fake_result_message) - - # send the fake result back to Claude - headers = HEADERS.copy() - headers["x-api-key"] = api_key - - tool_result_msg = { - "model": MODEL, - "max_tokens": 2048, - "tools": [], - "messages": messages, - "thinking": { - "type": "enabled", - "budget_tokens": 1024 - } - } - - final_response = send_request(headers, tool_result_msg) - - print("Final Claude response:") - for block in final_response.get("content", []): - if block["type"] == "text": - print(block["text"]) + "tool_use_id": tool_id, + "content": [result_content] if isinstance(result_content, dict) else result_content + }) + + # Add all tool results as a user message + messages.append({ + "role": "user", + "content": tool_results + }) + + print("\n[TOOL RESULTS SENT BACK TO CLAUDE]") + + print("\n[CONVERSATION COMPLETE]") if __name__ == "__main__":