add computer, bash, text editor tools (#9)
This commit was merged in pull request #9.
This commit is contained in:
@@ -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__":
|
||||
|
||||
Reference in New Issue
Block a user