Merge pull request #37 from ryanguo-google/main

Add FunctionResponsePart Support, and fix issues when using Vertex API

Author: ryanguo-google

agent.py

@@ -30,6 +30,24 @@
 
 from computers import EnvState, Computer
 
+MAX_RECENT_TURN_WITH_SCREENSHOTS = 3
+PREDEFINED_COMPUTER_USE_FUNCTIONS = [
+    "open_web_browser",
+    "click_at",
+    "hover_at",
+    "type_text_at",
+    "scroll_document",
+    "scroll_at",
+    "wait_5_seconds",
+    "go_back",
+    "go_forward",
+    "search",
+    "navigate",
+    "key_combination",
+    "drag_and_drop",
+]
+
+
 console = Console()
 
 # Built-in Computer Use tools will return "EnvState".
@@ -60,12 +78,6 @@ def __init__(
             vertexai=os.environ.get("USE_VERTEXAI", "0").lower() in ["true", "1"],
             project=os.environ.get("VERTEXAI_PROJECT"),
             location=os.environ.get("VERTEXAI_LOCATION"),
-            http_options=types.HttpOptions(
-                api_version="v1alpha",
-                base_url=os.environ.get(
-                    "GEMINI_API_SERVER", "https://generativelanguage.googleapis.com"
-                ),
-            ),
         )
         self._contents: list[Content] = [
             Content(
@@ -101,7 +113,6 @@ def __init__(
                 ),
                 types.Tool(function_declarations=custom_functions),
             ],
-            thinking_config=types.ThinkingConfig(include_thoughts=True),
         )
 
     def handle_action(self, action: types.FunctionCall) -> FunctionResponseT:
@@ -321,9 +332,9 @@ def run_one_iteration(self) -> Literal["COMPLETE", "CONTINUE"]:
                             "url": fc_result.url,
                             **extra_fr_fields,
                         },
-                        data=[
-                            types.Part(
-                                inline_data=types.Blob(
+                        parts=[
+                            types.FunctionResponsePart(
+                                inline_data=types.FunctionResponseBlob(
                                     mime_type="image/png", data=fc_result.screenshot
                                 )
                             )
@@ -342,6 +353,35 @@ def run_one_iteration(self) -> Literal["COMPLETE", "CONTINUE"]:
             )
         )
 
+        # only keep screenshots in the few most recent turns, remove the screenshot images from the old turns.
+        turn_with_screenshots_found = 0
+        for content in reversed(self._contents):
+            if content.role == "user" and content.parts:
+                # check if content has screenshot of the predefined computer use functions.
+                has_screenshot = False
+                for part in content.parts:
+                    if (
+                        part.function_response
+                        and part.function_response.parts
+                        and part.function_response.name
+                        in PREDEFINED_COMPUTER_USE_FUNCTIONS
+                    ):
+                        has_screenshot = True
+                        break
+
+                if has_screenshot:
+                    turn_with_screenshots_found += 1
+                    # remove the screenshot image if the number of screenshots exceed the limit.
+                    if turn_with_screenshots_found > MAX_RECENT_TURN_WITH_SCREENSHOTS:
+                        for part in content.parts:
+                            if (
+                                part.function_response
+                                and part.function_response.parts
+                                and part.function_response.name
+                                in PREDEFINED_COMPUTER_USE_FUNCTIONS
+                            ):
+                                part.function_response.parts = None
+
         return "CONTINUE"
 
     def _get_safety_confirmation(

requirements.txt

@@ -1,6 +1,6 @@
 termcolor==3.1.0
 pydantic==2.11.4
-./sdk/google_genai-1.27.0-py3-none-any.whl
+google-genai==1.39.0
 playwright==1.52.0
 browserbase==1.3.0
 rich