move kwargs towards user interaction from handlers

Author: filip-michalsky

stagehand/handlers/act_handler.py

@@ -35,83 +35,40 @@ def __init__(
         self.user_provided_instructions = user_provided_instructions
         self.self_heal = self_heal
 
-    async def act(
-        self, 
-        options_or_action: Union[ActOptions, ObserveResult, str, dict, None] = None,
-        **kwargs
-    ) -> ActResult:
+    async def act(self, options: Union[ActOptions, ObserveResult]) -> 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
         """
-        # Handle ObserveResult case first (legacy compatibility)
-        if isinstance(options_or_action, ObserveResult):
+        if "selector" in options and "method" in options:
+            options = ObserveResult(**options)
             return await self._act_from_observe_result(
-                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="",
+                options, self.stagehand.dom_settle_timeout_ms
             )
 
         # Start inference timer if available
         if hasattr(self.stagehand, "start_inference_timer"):
             self.stagehand.start_inference_timer()
 
-        action_task = options.action
+        action_task = options.get("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.variables,
+            variables=options.get("variables"),
         )
 
         observe_options_dict = {"instruction": prompt}
         # Add other observe options from ActOptions if they exist
-        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
+        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"
+            )
 
         observe_options = ObserveOptions(**observe_options_dict)
 
@@ -136,16 +93,16 @@ async def act(
             element_to_act_on = observe_results[0]
 
             # Substitute variables in arguments
-            if options.variables:
-                variables = options.variables
+            if options.get("variables"):
+                variables = options.get("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.dom_settle_timeout_ms
+            dom_settle_timeout_ms = options.get("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, Union
+from typing import Any, Optional
 
 from ..types.agent import (
     ActionExecutionResult,
@@ -33,49 +33,8 @@ 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_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",
-            }
-
+    async def perform_action(self, action: AgentAction) -> ActionExecutionResult:
+        """Execute a single action on the page."""
         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, Union
+from typing import Optional, TypeVar
 
 from pydantic import BaseModel
 
@@ -34,41 +34,19 @@ def __init__(
 
     async def extract(
         self,
-        options_or_instruction: Union[ExtractOptions, str, dict, None] = None,
+        options: Optional[ExtractOptions] = None,
         schema: Optional[type[BaseModel]] = None,
-        **kwargs,
     ) -> ExtractResult:
         """
         Execute an extraction operation locally.
 
         Args:
-            options_or_instruction: ExtractOptions, instruction string, dict, or None
+            options: ExtractOptions containing the instruction and other parameters
             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, Optional, Union
+from typing import Any
 
 from stagehand.a11y.utils import get_accessibility_tree, get_xpath_by_resolved_object_id
 from stagehand.llm.inference import observe as observe_inference
@@ -28,45 +28,22 @@ def __init__(
         self.logger = stagehand_client.logger
         self.user_provided_instructions = user_provided_instructions
 
-
+    # TODO: better kwargs
     async def observe(
         self,
-        options_or_instruction: Union[ObserveOptions, str, dict, None] = None,
+        options: ObserveOptions,
         from_act: bool = False,
-        **kwargs,
     ) -> list[ObserveResult]:
         """
         Execute an observation operation locally.
 
         Args:
-            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
+            options: ObserveOptions containing the instruction and other parameters
 
         Returns:
             list of ObserveResult instances
         """
-        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
+        instruction = options.instruction
         if not instruction:
             instruction = (
                 "Find elements that can be used for any future actions in the page. "
@@ -140,8 +117,7 @@ async def observe(
         )
 
         # Draw overlay if requested
-        draw_overlay = options.draw_overlay if options else False
-        if draw_overlay:
+        if options.draw_overlay:
             await draw_observe_overlay(self.stagehand_page, elements_with_selectors)
 
         # Return the list of results without trying to attach _llm_response