Merge pull request #63 from browserbase/miguel/stg-444-google-cua

support for google

Author: miguelg719

requirements.txt

@@ -6,4 +6,5 @@ playwright>=1.42.1
 requests>=2.31.0
 rich
 browserbase
-litellm==1.67.1
+litellm==1.67.1
+./temp/google_genai-1.14.0-py3-none-any.whl

stagehand/agent/agent.py

@@ -9,12 +9,14 @@
 )
 from .anthropic_cua import AnthropicCUAClient
 from .client import AgentClient
+from .google_cua import GoogleCUAClient
 from .openai_cua import OpenAICUAClient
 
 MODEL_TO_CLIENT_CLASS_MAP: dict[str, type[AgentClient]] = {
     "computer-use-preview": OpenAICUAClient,
-    "claude-3-5-sonnet-20240620": AnthropicCUAClient,
-    "claude-3-7-sonnet-20250219": AnthropicCUAClient,
+    "claude-3-5-sonnet-latest": AnthropicCUAClient,
+    "claude-3-7-sonnet-latest": AnthropicCUAClient,
+    "models/computer-use-exp": GoogleCUAClient,
 }
 
 AGENT_METRIC_FUNCTION_NAME = "AGENT_EXECUTE_TASK"
@@ -39,6 +41,8 @@ def __init__(self, stagehand_client, **kwargs):
             stagehand=self.stagehand, page=self.stagehand.page._page, logger=self.logger
         )
 
+        self.viewport = self.stagehand.page._page.viewport_size
+        # self.viewport = {"width": 1024, "height": 768}
         self.client: AgentClient = self._get_client()
 
     def _get_client(self) -> AgentClient:
@@ -61,6 +65,7 @@ def _get_client(self) -> AgentClient:
             config=self.config,
             logger=self.logger,
             handler=self.cua_handler,
+            viewport=self.viewport,
         )
 
     async def execute(

stagehand/agent/anthropic_cua.py

@@ -50,21 +50,33 @@ def __init__(
         config: Optional[AgentConfig] = None,
         logger: Optional[Any] = None,
         handler: Optional[CUAHandler] = None,
+        viewport: Optional[dict[str, int]] = None,
         **kwargs,
     ):
         super().__init__(model, instructions, config, logger, handler)
         self.anthropic_sdk_client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
 
-        dimensions = [1024, 768]  # Default dimensions
+        dimensions = (
+            (viewport["width"], viewport["height"]) if viewport else (1024, 768)
+        )  # Default dimensions
         if self.config:
             if hasattr(self.config, "display_width") and self.config.display_width is not None:  # type: ignore
                 dimensions[0] = self.config.display_width  # type: ignore
             if hasattr(self.config, "display_height") and self.config.display_height is not None:  # type: ignore
                 dimensions[1] = self.config.display_height  # type: ignore
-
+        computer_tool_type = (
+            "computer_20250124"
+            if model == "claude-3-7-sonnet-latest"
+            else "computer_20241022"
+        )
+        self.beta_flag = (
+            ["computer-use-2025-01-24"]
+            if model == "claude-3-7-sonnet-latest"
+            else ["computer-use-2024-10-22"]
+        )
         self.tools = [
             {
-                "type": "computer_20250124",
+                "type": computer_tool_type,
                 "name": "computer",
                 "display_width_px": dimensions[0],
                 "display_height_px": dimensions[1],
@@ -155,7 +167,7 @@ async def run_task(
                     + "Remember to call the computer tools, and only goto or navigate_back if you need to. Screenshots, clicks, etc, will be parsed from computer tool calls",  # System prompt
                     messages=current_messages,
                     tools=self.tools,
-                    betas=["computer-use-2025-01-24"],
+                    betas=self.beta_flag,
                 )
                 end_time = asyncio.get_event_loop().time()
                 total_inference_time_ms += int((end_time - start_time) * 1000)