Update BrowserUse step

whoisarpit · whoisarpit · commit 0d85c0440747 · 2025-03-17T13:40:38.000+05:30
diff --git a/patchwork/common/utils/browser_initializer.py b/patchwork/common/utils/browser_initializer.py
@@ -0,0 +1,121 @@
+import logging
+import os
+from typing import List, Optional
+
+from browser_use import Browser, BrowserConfig, BrowserContextConfig, Controller
+from browser_use.agent.views import ActionResult
+from browser_use.browser.context import BrowserContext
+
+logger = logging.getLogger(__name__)
+
+
+async def set_file_input(index: int, paths: str | List[str], browser: BrowserContext):
+    """
+    Set the file input value to the given path or list of paths.
+
+    Args:
+        index: The DOM element index to target
+        paths: Local file path or list of local file paths to upload
+        browser: Browser context for interaction
+
+    Returns:
+        ActionResult: Result of the upload operation
+    """
+    if isinstance(paths, str):
+        paths = [paths]
+
+    for path in paths:
+        if not os.path.exists(path):
+            return ActionResult(error=f"File {path} does not exist")
+
+    dom_el = await browser.get_dom_element_by_index(index)
+    file_upload_dom_el = dom_el.get_file_upload_element()
+
+    if file_upload_dom_el is None:
+        msg = f"No file upload element found at index {index}. The element may be hidden or not an input type file"
+        logger.info(msg)
+        return ActionResult(error=msg)
+
+    file_upload_el = await browser.get_locate_element(file_upload_dom_el)
+
+    if file_upload_el is None:
+        msg = f"No file upload element found at index {index}. The element may be hidden or not an input type file"
+        logger.info(msg)
+        return ActionResult(error=msg)
+
+    try:
+        await file_upload_el.set_input_files(paths)
+        msg = f"Successfully set file input value to {paths}"
+        logger.info(msg)
+        return ActionResult(extracted_content=msg, include_in_memory=True)
+    except Exception as e:
+        msg = f"Failed to upload file to index {index}: {str(e)}"
+        logger.info(msg)
+        return ActionResult(error=msg)
+
+
+class BrowserInitializer:
+    """
+    Initialize and cache browser and controller instances.
+
+    This class uses a singleton pattern to ensure we only create one browser
+    instance throughout the application lifecycle, which saves resources.
+    """
+
+    _browser = None
+    _controller = None
+    _browser_context = None
+
+    @classmethod
+    def init_browser(cls, config=BrowserConfig()):
+        """
+        Initialize and cache the Browser instance.
+
+        Returns:
+            Browser: Browser instance for web automation
+        """
+        if cls._browser is not None:
+            return cls._browser
+
+        cls._browser = Browser(config=config)
+        return cls._browser
+
+    @classmethod
+    def init_browser_context(cls, config: Optional[BrowserContextConfig]):
+        """
+        Initialize and cache the BrowserContext instance.
+
+        Returns:
+            BrowserContext: BrowserContext instance for managing browser context
+        """
+        if cls._browser_context is not None:
+            return cls._browser_context
+
+        downloads_path = os.path.join(os.getcwd(), "downloads")
+        if not os.path.exists(downloads_path):
+            os.makedirs(downloads_path)
+
+        cookies_file = os.path.join(os.getcwd(), "cookies.json")
+
+        context_config = BrowserContextConfig(
+            save_downloads_path=downloads_path, cookies_file=cookies_file, _force_keep_context_alive=True
+        )
+        browser = cls.init_browser(config=config)
+        cls._browser_context = BrowserContext(browser=browser, config=context_config)
+        return cls._browser_context
+
+    @classmethod
+    def init_controller(cls):
+        """
+        Initialize and cache the Controller instance.
+
+        Returns:
+            Controller: Controller instance for managing browser actions
+        """
+        if cls._controller is not None:
+            return cls._controller
+
+        controller = Controller()
+
+        cls._controller = controller
+        return cls._controller
diff --git a/patchwork/steps/BrowserUse/BrowserUse.py b/patchwork/steps/BrowserUse/BrowserUse.py
@@ -10,111 +10,6 @@
 
 logger = logging.getLogger(__name__)
 
-# Global variables to cache browser initialization
-_browser = None
-_controller = None
-
-
-def init_browser():
-    """
-    Initialize and cache browser and controller instances.
-
-    This function uses a singleton pattern to ensure we only create one browser
-    instance throughout the application lifecycle, which saves resources.
-
-    Returns:
-        tuple: (Browser, Controller) instances for web automation
-    """
-    global _browser, _controller
-
-    # Return cached instances if already initialized
-    if _browser is not None and _controller is not None:
-        return _browser, _controller
-
-    from browser_use import Browser, BrowserConfig, BrowserContextConfig, Controller
-    from browser_use.agent.views import ActionResult
-    from browser_use.browser.context import BrowserContext
-
-    # Set up downloads directory for browser operations
-    downloads_path = os.path.join(os.getcwd(), "downloads")
-    if not os.path.exists(downloads_path):
-        os.makedirs(downloads_path)
-
-    context_config = BrowserContextConfig(save_downloads_path=downloads_path)
-    config = BrowserConfig(headless=True, disable_security=True, new_context_config=context_config)
-    controller = Controller()
-
-    # Register custom action to upload files to web elements
-    @controller.action(
-        description="Upload file to interactive element with file path",
-    )
-    async def upload_file(index: int, path: str, browser: BrowserContext):
-        """
-        Upload a file to a file input element identified by its index.
-
-        Args:
-            index: The DOM element index to target
-            path: Local file path to upload
-            browser: Browser context for interaction
-
-        Returns:
-            ActionResult: Result of the upload operation
-        """
-        if not os.path.exists(path):
-            return ActionResult(error=f"File {path} does not exist")
-
-        dom_el = await browser.get_dom_element_by_index(index)
-        file_upload_dom_el = dom_el.get_file_upload_element()
-
-        if file_upload_dom_el is None:
-            msg = f"No file upload element found at index {index}. The element may be hidden or not an input type file"
-            logger.info(msg)
-            return ActionResult(error=msg)
-
-        file_upload_el = await browser.get_locate_element(file_upload_dom_el)
-
-        if file_upload_el is None:
-            msg = f"No file upload element found at index {index}. The element may be hidden or not an input type file"
-            logger.info(msg)
-            return ActionResult(error=msg)
-
-        try:
-            await file_upload_el.set_input_files(path)
-            msg = f"Successfully uploaded file to index {index}"
-            logger.info(msg)
-            return ActionResult(extracted_content=msg, include_in_memory=True)
-        except Exception as e:
-            msg = f"Failed to upload file to index {index}: {str(e)}"
-            logger.info(msg)
-            return ActionResult(error=msg)
-
-    # Register custom action to read file contents
-    @controller.action(description="Read the file content of a file given a path")
-    async def read_file(path: str):
-        """
-        Read and return the contents of a file at the specified path.
-
-        Args:
-            path: Path to the file to read
-
-        Returns:
-            ActionResult: File contents or error message
-        """
-        if not os.path.exists(path):
-            return ActionResult(error=f"File {path} does not exist")
-
-        with open(path, "r") as f:
-            content = f.read()
-        msg = f"File content: {content}"
-        logger.info(msg)
-        return ActionResult(extracted_content=msg, include_in_memory=True)
-
-    # Cache the initialized instances
-    _browser = Browser(config=config)
-    _controller = controller
-
-    return _browser, _controller
-
 
 class BrowserUse(Step, input_class=BrowserUseInputs, output_class=BrowserUseOutputs):
     """
@@ -174,20 +69,29 @@ def run(self) -> dict:
             dict: Results of the browser automation task
         """
         from browser_use import Agent
+        from browser_use import BrowserConfig
+        from patchwork.common.utils.browser_initializer import BrowserInitializer
 
-        browser, controller = init_browser()
+        browser_config = BrowserConfig(headless=self.inputs.get("headless", True), disable_security=True)
+        browser_context = BrowserInitializer.init_browser_context(browser_config)
+        controller = BrowserInitializer.init_controller()
+        logger.info("Browser initialized")
         agent = Agent(
-            browser=browser,
+            browser_context=browser_context,
             controller=controller,
             task=mustache_render(self.inputs["task"], self.inputs["task_value"]),
             llm=self.llm,
             generate_gif=self.generate_gif,
             validate_output=True,
+            initial_actions=self.inputs.get("initial_actions", None),
         )
 
         # Run the agent in an event loop
         loop = asyncio.new_event_loop()
         self.history = loop.run_until_complete(agent.run())
+        loop.run_until_complete(browser_context.close())
+        loop.run_until_complete(browser_context.browser.close())
+        loop.close()
 
         # Format results as JSON if schema provided
         if "example_json" in self.inputs:
diff --git a/patchwork/steps/BrowserUse/typed.py b/patchwork/steps/BrowserUse/typed.py
@@ -1,4 +1,4 @@
-from typing_extensions import Annotated, Any, Dict, Optional, TypedDict
+from typing_extensions import Annotated, Any, Dict, Optional, TypedDict, List
 
 from patchwork.common.utils.step_typing import StepTypeConfig
 
@@ -14,6 +14,8 @@ class BrowserUseInputs(__BrowserUseInputsRequired, total=False):
     anthropic_api_key: Annotated[str, StepTypeConfig(or_op=["google_api_key", "openai_api_key"])]
     google_api_key: Annotated[str, StepTypeConfig(or_op=["openai_api_key", "anthropic_api_key"])]
     generate_gif: Optional[bool]
+    headless: Optional[bool]
+    initial_actions: Optional[List[Dict[str, Dict[str, Any]]]]
 
 
 class BrowserUseOutputs(TypedDict):
