add computer, bash, text editor tools

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
+            tool_uses.append(block)
-    return None
+    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:
+    # Main conversation loop
-    # "Extra inputs are not permitted"
+    while True:
-    messages.append({
+        response = send_prompt(api_key, messages)
-        'role': 'assistant',
-        'content': first_response.get('content', []),
-    })
         
-    tool_use = extract_tool_use(first_response)
+        # Add Claude's response to messages
-    if not tool_use:
+        messages.append({
-        print("No tool use required. Response:")
+            'role': 'assistant',
-        print(first_response["content"])
+            'content': response.get('content', []),
-        return
+        })
         
-    # print first response from Claude
+        # Print Claude's response
-    for block in first_response.get("content", []):
+        print("\n[CLAUDE RESPONSE]")
-        print(f'{block['type']}:')
+        for block in response.get("content", []):
-        if block["type"] == "text":
+            if block["type"] == "text":
-            print(f'\t{block["text"]}')
+                print(block["text"])
-        elif block["type"] == "thinking":
+            elif block["type"] == "thinking":
-            print(f'\t{block["thinking"]}')
+                print(f"[THINKING] {block['thinking']}")
-        elif block["type"] == "tool_use":
+            elif block["type"] == "tool_use":
-            print(f'\tTool use requested: {block["name"]}')
+                print(f"[TOOL REQUEST] {block['name']} - {block.get('input', {})}")
-            print(f'\tInput: {block["input"]}')
         
-    # this is a stub for the tool execution
+        # Extract all tool uses
-    fake_result_message = {
+        tool_uses = extract_tool_uses(response)
-        "role": "user",
+        
-        "content": [
+        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"],
+                "tool_use_id": tool_id,
-                "content": "(Stub) Executed tool successfully"
+                "content": [result_content] if isinstance(result_content, dict) else result_content
-            }
+            })
-        ]
-    }
-    messages.append(fake_result_message)
         
-    # send the fake result back to Claude
+        # Add all tool results as a user message
-    headers = HEADERS.copy()
+        messages.append({
-    headers["x-api-key"] = api_key
+            "role": "user",
+            "content": tool_results
+        })
         
-    tool_result_msg = {
+        print("\n[TOOL RESULTS SENT BACK TO CLAUDE]")
-        "model": MODEL,
-        "max_tokens": 2048,
-        "tools": [],
-        "messages": messages,
-        "thinking": {
-            "type": "enabled",
-            "budget_tokens": 1024
-        }
-    }
     
-    final_response = send_request(headers, tool_result_msg)
+    print("\n[CONVERSATION COMPLETE]")
-
-    print("Final Claude response:")
-    for block in final_response.get("content", []):
-        if block["type"] == "text":
-            print(block["text"])
 
 
 if __name__ == "__main__":