146 lines
4.5 KiB
Python
146 lines
4.5 KiB
Python
import os
|
|
import re
|
|
import base64
|
|
from desktop_env.envs.desktop_env import Action, MouseClick
|
|
import json
|
|
import requests
|
|
from mm_agents.gpt_4v_prompt import SYS_PROMPT
|
|
|
|
# Function to encode the image
|
|
def encode_image(image_path):
|
|
with open(image_path, "rb") as image_file:
|
|
return base64.b64encode(image_file.read()).decode('utf-8')
|
|
|
|
|
|
def parse_action_from_string(input_string):
|
|
# Search for a JSON string within the input string
|
|
matches = re.findall(r'```json\s+(.*?)\s+```', input_string, re.DOTALL)
|
|
if matches:
|
|
# Assuming there's only one match, parse the JSON string into a dictionary
|
|
try:
|
|
action_dict = json.loads(matches[0])
|
|
return action_dict
|
|
except json.JSONDecodeError as e:
|
|
return f"Failed to parse JSON: {e}"
|
|
else:
|
|
matches = re.findall(r'```\s+(.*?)\s+```', input_string, re.DOTALL)
|
|
if matches:
|
|
# Assuming there's only one match, parse the JSON string into a dictionary
|
|
try:
|
|
action_dict = json.loads(matches[0])
|
|
return action_dict
|
|
except json.JSONDecodeError as e:
|
|
return f"Failed to parse JSON: {e}"
|
|
else:
|
|
try:
|
|
action_dict = json.loads(input_string)
|
|
return action_dict
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError("Invalid response format: " + input_string)
|
|
|
|
class GPT4v_Agent:
|
|
def __init__(self, api_key, instruction, model="gpt-4-vision-preview", max_tokens=300):
|
|
self.instruction = instruction
|
|
self.model = model
|
|
self.max_tokens = max_tokens
|
|
|
|
self.headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {api_key}"
|
|
}
|
|
|
|
self.trajectory = [
|
|
{
|
|
"role": "system",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": SYS_PROMPT
|
|
},
|
|
]
|
|
}
|
|
]
|
|
|
|
def predict(self, obs):
|
|
base64_image = encode_image(obs)
|
|
self.trajectory.append({
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": "What's the next step for instruction '{}'?".format(self.instruction)
|
|
},
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {
|
|
"url": f"data:image/jpeg;base64,{base64_image}"
|
|
}
|
|
}
|
|
]
|
|
})
|
|
payload = {
|
|
"model": self.model,
|
|
"messages": self.trajectory,
|
|
"max_tokens": self.max_tokens
|
|
}
|
|
response = requests.post("https://api.openai.com/v1/chat/completions", headers=self.headers, json=payload)
|
|
|
|
action = self.parse_action(response.json()['choices'][0]['message']['content'])
|
|
|
|
return action
|
|
|
|
def parse_action(self, response: str):
|
|
# response example
|
|
"""
|
|
```json
|
|
{
|
|
"action_type": "CLICK",
|
|
"click_type": "RIGHT"
|
|
}
|
|
```
|
|
"""
|
|
|
|
# parse from the response
|
|
action = parse_action_from_string(response)
|
|
|
|
# add action into the trajectory
|
|
self.trajectory.append({
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": response
|
|
},
|
|
]
|
|
})
|
|
|
|
# parse action
|
|
parsed_action = {}
|
|
action_type = Action[action['action_type']].value
|
|
parsed_action["action_type"] = action_type
|
|
|
|
if action_type == Action.CLICK.value or action_type == Action.MOUSE_DOWN.value or action_type == Action.MOUSE_UP.value:
|
|
parsed_action["click_type"] = MouseClick[action['click_type']].value
|
|
|
|
if action_type == Action.MOUSE_MOVE.value:
|
|
parsed_action["x"] = action["x"]
|
|
parsed_action["y"] = action["y"]
|
|
|
|
# fixme: could these two actions be merged??
|
|
if action_type == Action.KEY.value:
|
|
parsed_action["key"] = [ord(c) for c in action["key"]]
|
|
|
|
if action_type == Action.TYPE.value:
|
|
parsed_action["text"] = [ord(c) for c in action["text"]]
|
|
|
|
return parsed_action
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# OpenAI API Key
|
|
api_key = os.environ.get("OPENAI_API_KEY")
|
|
|
|
agent = GPT4v_Agent(api_key=api_key, instruction="Open Google Sheet")
|
|
print(agent.predict(obs="stackoverflow.png"))
|
|
|