update kwargs

Author: filip-michalsky

stagehand/handlers/act_handler.py

@@ -35,40 +35,83 @@ def __init__(
         self.user_provided_instructions = user_provided_instructions
         self.self_heal = self_heal
 
-    async def act(self, options: Union[ActOptions, ObserveResult]) -> ActResult:
+    async def act(
+        self, 
+        options_or_action: Union[ActOptions, ObserveResult, str, dict, None] = None,
+        **kwargs
+    ) -> ActResult:
         """
         Perform an act based on an instruction.
         This method will observe the page and then perform the act on the first element returned.
+        
+        Args:
+            options_or_action: ActOptions, ObserveResult, action string, dict, or None
+            **kwargs: Additional options to be merged
+            
+        Returns:
+            ActResult instance
         """
-        if "selector" in options and "method" in options:
-            options = ObserveResult(**options)
+        # Handle ObserveResult case first (legacy compatibility)
+        if isinstance(options_or_action, ObserveResult):
             return await self._act_from_observe_result(
-                options, self.stagehand.dom_settle_timeout_ms
+                options_or_action, self.stagehand.dom_settle_timeout_ms
+            )
+        
+        # Handle the new flexible parameter format
+        options: Optional[ActOptions] = None
+        options_dict = {}
+
+        if isinstance(options_or_action, ActOptions):
+            options_dict = options_or_action.model_dump()
+        elif isinstance(options_or_action, dict):
+            # Check if it's actually an ObserveResult dict
+            if "selector" in options_or_action and "method" in options_or_action:
+                observe_result = ObserveResult(**options_or_action)
+                return await self._act_from_observe_result(
+                    observe_result, self.stagehand.dom_settle_timeout_ms
+                )
+            options_dict = options_or_action.copy()
+        elif isinstance(options_or_action, str):
+            options_dict["action"] = options_or_action
+
+        options_dict.update(kwargs)
+
+        # Validate options if we have any
+        if options_dict:
+            try:
+                options = ActOptions(**options_dict)
+            except Exception as e:
+                self.logger.error(f"Invalid act options: {e}")
+                raise
+        
+        if not options or not options.action:
+            return ActResult(
+                success=False,
+                message="No action provided for act operation",
+                action="",
             )
 
         # Start inference timer if available
         if hasattr(self.stagehand, "start_inference_timer"):
             self.stagehand.start_inference_timer()
 
-        action_task = options.get("action")
+        action_task = options.action
         self.logger.info(
             f"Starting action for task: '{action_task}'",
             category="act",
         )
         prompt = build_act_observe_prompt(
             action=action_task,
             supported_actions=list(method_handler_map.keys()),
-            variables=options.get("variables"),
+            variables=options.variables,
         )
 
         observe_options_dict = {"instruction": prompt}
         # Add other observe options from ActOptions if they exist
-        if options.get("model_name"):
-            observe_options_dict["model_name"] = options.get("model_name")
-        if options.get("model_client_options"):
-            observe_options_dict["model_client_options"] = options.get(
-                "model_client_options"
-            )
+        if options.model_name:
+            observe_options_dict["model_name"] = options.model_name
+        if options.model_client_options:
+            observe_options_dict["model_client_options"] = options.model_client_options
 
         observe_options = ObserveOptions(**observe_options_dict)
 
@@ -93,16 +136,16 @@ async def act(self, options: Union[ActOptions, ObserveResult]) -> ActResult:
             element_to_act_on = observe_results[0]
 
             # Substitute variables in arguments
-            if options.get("variables"):
-                variables = options.get("variables", {})
+            if options.variables:
+                variables = options.variables
                 element_to_act_on.arguments = [
                     str(arg).replace(f"%{key}%", str(value))
                     for arg in (element_to_act_on.arguments or [])
                     for key, value in variables.items()
                 ]
 
             # domSettleTimeoutMs might come from options if specified for act
-            dom_settle_timeout_ms = options.get("dom_settle_timeout_ms")
+            dom_settle_timeout_ms = options.dom_settle_timeout_ms
 
             try:
                 await self._perform_playwright_method(

stagehand/handlers/cua_handler.py

@@ -1,6 +1,6 @@
 import asyncio
 import base64
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 from ..types.agent import (
     ActionExecutionResult,
@@ -33,8 +33,49 @@ async def get_screenshot_base64(self) -> str:
         screenshot_bytes = await self.page.screenshot(full_page=False, type="png")
         return base64.b64encode(screenshot_bytes).decode()
 
-    async def perform_action(self, action: AgentAction) -> ActionExecutionResult:
-        """Execute a single action on the page."""
+    async def perform_action(
+        self, 
+        action_or_dict: Union[AgentAction, dict, None] = None,
+        **kwargs
+    ) -> ActionExecutionResult:
+        """
+        Execute a single action on the page.
+        
+        Args:
+            action_or_dict: AgentAction, dict, or None
+            **kwargs: Additional action parameters to be merged
+            
+        Returns:
+            ActionExecutionResult dict
+        """
+        action: Optional[AgentAction] = None
+        action_dict = {}
+
+        if isinstance(action_or_dict, AgentAction):
+            action_dict = action_or_dict.model_dump()
+        elif isinstance(action_or_dict, dict):
+            action_dict = action_or_dict.copy()
+
+        action_dict.update(kwargs)
+
+        # Validate action if we have any
+        if action_dict:
+            try:
+                action = AgentAction(**action_dict)
+            except Exception as e:
+                self.logger.error(f"Invalid agent action: {e}")
+                raise
+        
+        if not action:
+            self.logger.error(
+                "No action provided for perform_action",
+                category=StagehandFunctionName.AGENT,
+            )
+            return {
+                "success": False,
+                "error": "No action provided",
+            }
+
         self.logger.info(
             f"Performing action: {action.action.root if action.action else ''}",
             category=StagehandFunctionName.AGENT,

stagehand/handlers/extract_handler.py

@@ -1,6 +1,6 @@
 """Extract handler for performing data extraction from page elements using LLMs."""
 
-from typing import Optional, TypeVar
+from typing import Optional, TypeVar, Union
 
 from pydantic import BaseModel
 
@@ -34,19 +34,41 @@ def __init__(
 
     async def extract(
         self,
-        options: Optional[ExtractOptions] = None,
+        options_or_instruction: Union[ExtractOptions, str, dict, None] = None,
         schema: Optional[type[BaseModel]] = None,
+        **kwargs,
     ) -> ExtractResult:
         """
         Execute an extraction operation locally.
 
         Args:
-            options: ExtractOptions containing the instruction and other parameters
+            options_or_instruction: ExtractOptions, instruction string, dict, or None
             schema: Optional Pydantic model for structured output
+            **kwargs: Additional options to be merged
 
         Returns:
             ExtractResult instance
         """
+        options: Optional[ExtractOptions] = None
+        options_dict = {}
+
+        if isinstance(options_or_instruction, ExtractOptions):
+            options_dict = options_or_instruction.model_dump()
+        elif isinstance(options_or_instruction, dict):
+            options_dict = options_or_instruction.copy()
+        elif isinstance(options_or_instruction, str):
+            options_dict["instruction"] = options_or_instruction
+
+        options_dict.update(kwargs)
+
+        # Validate options if we have any
+        if options_dict:
+            try:
+                options = ExtractOptions(**options_dict)
+            except Exception as e:
+                self.logger.error(f"Invalid extract options: {e}")
+                raise
+
         if not options:
             # If no options provided, extract the entire page text
             self.logger.info("Extracting entire page text")

stagehand/handlers/observe_handler.py

@@ -1,6 +1,6 @@
 """Observe handler for performing observations of page elements using LLMs."""
 
-from typing import Any
+from typing import Any, Optional, Union
 
 from stagehand.a11y.utils import get_accessibility_tree, get_xpath_by_resolved_object_id
 from stagehand.llm.inference import observe as observe_inference
@@ -28,22 +28,45 @@ def __init__(
         self.logger = stagehand_client.logger
         self.user_provided_instructions = user_provided_instructions
 
-    # TODO: better kwargs
+
     async def observe(
         self,
-        options: ObserveOptions,
+        options_or_instruction: Union[ObserveOptions, str, dict, None] = None,
         from_act: bool = False,
+        **kwargs,
     ) -> list[ObserveResult]:
         """
         Execute an observation operation locally.
 
         Args:
-            options: ObserveOptions containing the instruction and other parameters
+            options_or_instruction: ObserveOptions, instruction string, dict, or None
+            from_act: Whether this observe is being called from act
+            **kwargs: Additional options to be merged
 
         Returns:
             list of ObserveResult instances
         """
-        instruction = options.instruction
+        options: Optional[ObserveOptions] = None
+        options_dict = {}
+
+        if isinstance(options_or_instruction, ObserveOptions):
+            options_dict = options_or_instruction.model_dump()
+        elif isinstance(options_or_instruction, dict):
+            options_dict = options_or_instruction.copy()
+        elif isinstance(options_or_instruction, str):
+            options_dict["instruction"] = options_or_instruction
+
+        options_dict.update(kwargs)
+
+        # Validate options if we have any
+        if options_dict:
+            try:
+                options = ObserveOptions(**options_dict)
+            except Exception as e:
+                self.logger.error(f"Invalid observe options: {e}")
+                raise
+
+        instruction = options.instruction if options else None
         if not instruction:
             instruction = (
                 "Find elements that can be used for any future actions in the page. "
@@ -117,7 +140,8 @@ async def observe(
         )
 
         # Draw overlay if requested
-        if options.draw_overlay:
+        draw_overlay = options.draw_overlay if options else False
+        if draw_overlay:
             await draw_observe_overlay(self.stagehand_page, elements_with_selectors)
 
         # Return the list of results without trying to attach _llm_response