Merge pull request #110 from linusaltacc/gemini-pro-vision-support

Added Gemini Pro Vision Support to Self Operating Computer

Author: joshbickett

.example.env

@@ -1 +1,2 @@
-OPENAI_API_KEY='your-key-here'
+OPENAI_API_KEY='your-key-here'
+GOOGLE_API_KEY='your-key-here'

README.md

@@ -20,7 +20,7 @@
 
 ## Key Features
 - **Compatibility**: Designed for various multimodal models.
-- **Integration**: Currently integrated with **GPT-4v** as the default model.
+- **Integration**: Currently integrated with **GPT-4v** as the default model, with extended support for Gemini Pro Vision.
 - **Future Plans**: Support for additional models.
 
 ## Current Challenges
@@ -75,6 +75,12 @@ mv .example.env .env
 ```
 OPENAI_API_KEY='your-key-here'
 ```
+OR 
+
+  **Add your Google AI Studio API key to your new .env file. If you don't have one, you can obtain a key [here](https://makersuite.google.com/app/apikey) after setting up your Google AI Studio account**:
+```
+GOOGLE_API_KEY='your-key-here'
+```
 8. **Run it**!
 ```
 operate

operate/main.py

@@ -14,7 +14,7 @@
 import Xlib.display
 import Xlib.X
 import Xlib.Xutil  # not sure if Xutil is necessary
-
+import google.generativeai as genai
 from prompt_toolkit import prompt
 from prompt_toolkit.shortcuts import message_dialog
 from prompt_toolkit.styles import Style as PromptStyle
@@ -29,9 +29,14 @@
 
 DEBUG = False
 
-client = OpenAI()
-client.api_key = os.getenv("OPENAI_API_KEY")
-client.base_url = os.getenv("OPENAI_API_BASE_URL", client.base_url)
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
+
+if OPENAI_API_KEY:
+    client = OpenAI()
+    client.api_key = OPENAI_API_KEY
+    client.base_url = os.getenv("OPENAI_API_BASE_URL", client.base_url)
+
 
 monitor_size = {
     "width": 1920,
@@ -194,12 +199,37 @@ def supports_ansi():
     ANSI_BRIGHT_MAGENTA = ""
 
 
+def validation(
+    model,
+    accurate_mode,
+    voice_mode,
+):
+    if accurate_mode and model != "gpt-4-vision-preview":
+        print("To use accuracy mode, please use gpt-4-vision-preview")
+        sys.exit(1)
+
+    if voice_mode and not OPENAI_API_KEY:
+        print("To use voice mode, please add an OpenAI API key")
+        sys.exit(1)
+
+    if model == "gpt-4-vision-preview" and not OPENAI_API_KEY:
+        print("To use `gpt-4-vision-preview` add an OpenAI API key")
+        sys.exit(1)
+
+    if model == "gemini-pro-vision" and not GOOGLE_API_KEY:
+        print("To use `gemini-pro-vision` add a Google API key")
+        sys.exit(1)
+
+
 def main(model, accurate_mode, terminal_prompt, voice_mode=False):
     """
     Main function for the Self-Operating Computer
     """
     mic = None
     # Initialize WhisperMic if voice_mode is True if voice_mode is True
+
+    validation(model, accurate_mode, voice_mode)
+
     if voice_mode:
         try:
             from whisper_mic import WhisperMic
@@ -259,7 +289,8 @@ def main(model, accurate_mode, terminal_prompt, voice_mode=False):
             print("[loop] messages before next action:\n\n\n", messages[1:])
         try:
             response = get_next_action(model, messages, objective, accurate_mode)
-            action = parse_oai_response(response)
+
+            action = parse_response(response)
             action_type = action.get("type")
             action_detail = action.get("data")
 
@@ -358,6 +389,11 @@ def get_next_action(model, messages, objective, accurate_mode):
         return content
     elif model == "agent-1":
         return "coming soon"
+    elif model == "gemini-pro-vision":
+        content = get_next_action_from_gemini_pro_vision(
+            messages, objective, accurate_mode
+        )
+        return content
 
     raise ModelNotRecognizedException(model)
 
@@ -376,10 +412,11 @@ def get_last_assistant_message(messages):
     return None  # Return None if no assistant message is found
 
 
-def accurate_mode_double_check(pseudo_messages, prev_x, prev_y):
+def accurate_mode_double_check(model, pseudo_messages, prev_x, prev_y):
     """
     Reprompt OAI with additional screenshot of a mini screenshot centered around the cursor for further finetuning of clicked location
     """
+    print("[get_next_action_from_gemini_pro_vision] accurate_mode_double_check")
     try:
         screenshot_filename = os.path.join("screenshots", "screenshot_mini.png")
         capture_mini_screenshot_with_cursor(
@@ -394,31 +431,37 @@ def accurate_mode_double_check(pseudo_messages, prev_x, prev_y):
             img_base64 = base64.b64encode(img_file.read()).decode("utf-8")
 
         accurate_vision_prompt = format_accurate_mode_vision_prompt(prev_x, prev_y)
+        if model == "gpt-4-vision-preview":
+            accurate_mode_message = {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": accurate_vision_prompt},
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/jpeg;base64,{img_base64}"},
+                    },
+                ],
+            }
 
-        accurate_mode_message = {
-            "role": "user",
-            "content": [
-                {"type": "text", "text": accurate_vision_prompt},
-                {
-                    "type": "image_url",
-                    "image_url": {"url": f"data:image/jpeg;base64,{img_base64}"},
-                },
-            ],
-        }
-
-        pseudo_messages.append(accurate_mode_message)
-
-        response = client.chat.completions.create(
-            model="gpt-4-vision-preview",
-            messages=pseudo_messages,
-            presence_penalty=1,
-            frequency_penalty=1,
-            temperature=0.7,
-            max_tokens=300,
-        )
+            pseudo_messages.append(accurate_mode_message)
 
-        content = response.choices[0].message.content
+            response = client.chat.completions.create(
+                model="gpt-4-vision-preview",
+                messages=pseudo_messages,
+                presence_penalty=1,
+                frequency_penalty=1,
+                temperature=0.7,
+                max_tokens=300,
+            )
 
+            content = response.choices[0].message.content
+        elif model == "gemini-pro-vision":
+            model = genai.GenerativeModel("gemini-pro-vision")
+            response = model.generate_content(
+                [accurate_vision_prompt, Image.open(new_screenshot_filename)]
+            )
+            content = response.text[1:]
+            print(content)
         return content
     except Exception as e:
         print(f"Error reprompting model for accurate_mode: {e}")
@@ -501,7 +544,9 @@ def get_next_action_from_openai(messages, objective, accurate_mode):
                     print(
                         f"Previous coords before accurate tuning: prev_x {prev_x} prev_y {prev_y}"
                     )
-                content = accurate_mode_double_check(pseudo_messages, prev_x, prev_y)
+                content = accurate_mode_double_check(
+                    "gpt-4-vision-preview", pseudo_messages, prev_x, prev_y
+                )
                 assert content != "ERROR", "ERROR: accurate_mode_double_check failed"
 
         return content
@@ -511,7 +556,69 @@ def get_next_action_from_openai(messages, objective, accurate_mode):
         return "Failed take action after looking at the screenshot"
 
 
-def parse_oai_response(response):
+def get_next_action_from_gemini_pro_vision(messages, objective, accurate_mode):
+    """
+    Get the next action for Self-Operating Computer using Gemini Pro Vision
+    """
+    print("[get_next_action_from_gemini_pro_vision] messages", messages)
+    # sleep for a second
+    time.sleep(1)
+    try:
+        screenshots_dir = "screenshots"
+        if not os.path.exists(screenshots_dir):
+            os.makedirs(screenshots_dir)
+
+        screenshot_filename = os.path.join(screenshots_dir, "screenshot.png")
+        # Call the function to capture the screen with the cursor
+        capture_screen_with_cursor(screenshot_filename)
+
+        new_screenshot_filename = os.path.join(
+            "screenshots", "screenshot_with_grid.png"
+        )
+
+        add_grid_to_image(screenshot_filename, new_screenshot_filename, 500)
+        # sleep for a second
+        time.sleep(1)
+
+        with open(new_screenshot_filename, "rb") as img_file:
+            img_base64 = base64.b64encode(img_file.read()).decode("utf-8")
+
+        previous_action = get_last_assistant_message(messages)
+
+        vision_prompt = format_vision_prompt(objective, previous_action)
+
+        model = genai.GenerativeModel("gemini-pro-vision")
+        print("[get_next_action_from_gemini_pro_vision]  model.generate_content")
+
+        response = model.generate_content(
+            [vision_prompt, Image.open(new_screenshot_filename)]
+        )
+
+        # create a copy of messages and save to pseudo_messages
+        pseudo_messages = messages.copy()
+        pseudo_messages.append(response.text)
+        print(
+            "[get_next_action_from_gemini_pro_vision] pseudo_messages.append(response.text)",
+            response.text,
+        )
+
+        messages.append(
+            {
+                "role": "user",
+                "content": "`screenshot.png`",
+            }
+        )
+        content = response.text[1:]
+
+        return content
+
+    except Exception as e:
+        print(f"Error parsing JSON: {e}")
+        return "Failed take action after looking at the screenshot"
+
+
+def parse_response(response):
+    print("[parse_response] response", response)
     if response == "DONE":
         return {"type": "DONE", "data": None}
     elif response.startswith("CLICK"):
@@ -522,12 +629,18 @@ def parse_oai_response(response):
 
     elif response.startswith("TYPE"):
         # Extract the text to type
-        type_data = re.search(r'TYPE "(.+)"', response, re.DOTALL).group(1)
+        try:
+            type_data = re.search(r"TYPE (.+)", response, re.DOTALL).group(1)
+        except:
+            type_data = re.search(r'TYPE "(.+)"', response, re.DOTALL).group(1)
         return {"type": "TYPE", "data": type_data}
 
     elif response.startswith("SEARCH"):
         # Extract the search query
-        search_data = re.search(r'SEARCH "(.+)"', response).group(1)
+        try:
+            search_data = re.search(r'SEARCH "(.+)"', response).group(1)
+        except:
+            search_data = re.search(r"SEARCH (.+)", response).group(1)
         return {"type": "SEARCH", "data": search_data}
 
     return {"type": "UNKNOWN", "data": response}

requirements.txt

@@ -47,3 +47,4 @@ typing_extensions==4.8.0
 urllib3==2.0.7
 wcwidth==0.2.9
 zipp==3.17.0
+google-generativeai==0.3.0