Update BrowserUse and move browser initialization to utils

whoisarpit · whoisarpit · commit 61818151d336 · 2025-03-26T13:16:14.000+08:00
diff --git a/patchwork/common/utils/browser_initializer.py b/patchwork/common/utils/browser_initializer.py
@@ -0,0 +1,167 @@
+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)
+
+
+async def close_current_tab(browser: BrowserContext):
+    await browser.close_current_tab()
+    msg = "🔄  Closed current tab"
+    logger.info(msg)
+    return ActionResult(extracted_content=msg, include_in_memory=True)
+
+
+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[BrowserConfig], downloads_path: Optional[str] = None):
+        """
+        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
+
+        if downloads_path and not os.path.exists(downloads_path):
+            os.makedirs(downloads_path)
+
+        context_config = BrowserContextConfig(
+            # cookies_file=cookies_file,
+            browser_window_size={"width": 1920, "height": 1080},
+        )
+        browser = cls.init_browser(config=config)
+
+        class BrowserContextWithDownloadHandling(BrowserContext):
+            async def handle_download(self, download):
+                suggested_filename = download.suggested_filename
+                unique_filename = await self._get_unique_filename(downloads_path, suggested_filename)
+                download_path = os.path.join(downloads_path, unique_filename)
+                await download.save_as(download_path)
+                logger.info(f"Downloaded file saved to {download_path}")
+
+            async def _initialize_session(self):
+                async def _download_listener(download):
+                    logger.info("[BUD] Download event triggered")
+                    await self.handle_download(download)
+                    return download
+
+                def _new_page_listener(page):
+                    logger.info("[BUD] Adding download event listener to page")
+                    page.on("download", _download_listener)
+                    return page
+
+                await super()._initialize_session()
+
+                logger.info("[BUD] Adding page event listener to context")
+                self.session.context.on("page", _new_page_listener)
+
+                logger.info(f"[BUD] Adding download event listener to {len(self.session.context.pages)} existing pages")
+                for page in self.session.context.pages:
+                    page.on("download", _download_listener)
+
+        cls._browser_context = (
+            BrowserContextWithDownloadHandling(browser=browser, config=context_config)
+            if downloads_path
+            else 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()
+
+        controller.action(
+            "Set the value of a file input to the given path or list of paths",
+        )(set_file_input)
+
+        controller.action(
+            description="Close the tab that is currently active",
+        )(close_current_tab)
+
+        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):
     """
@@ -155,9 +50,13 @@ def __init__(self, inputs):
                 api_key=self.inputs["anthropic_api_key"],
             )
 
+        gifs_base_path = os.path.join(os.path.dirname(__file__), "../../../tmp/gifs")
+        if not os.path.exists(gifs_base_path):
+            os.makedirs(gifs_base_path)
+
         # Configure GIF generation for debugging/visualization
         self.generate_gif = (
-            f"agent_history_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.gif"
+            f"{gifs_base_path}/agent_history_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.gif"
             if ("generate_gif" in self.inputs and self.inputs["generate_gif"])
             or ("debug" in self.inputs and self.inputs["debug"])
             else False
@@ -173,21 +72,37 @@ def run(self) -> dict:
         Returns:
             dict: Results of the browser automation task
         """
-        from browser_use import Agent
+        from browser_use import Agent, BrowserConfig
+
+        from patchwork.common.utils.browser_initializer import BrowserInitializer
+
+        browser_config = BrowserConfig(
+            headless=self.inputs.get("headless", True),
+            disable_security=True,
+        )
+        browser_context = BrowserInitializer.init_browser_context(
+            browser_config, self.inputs.get("downloads_path", None)
+        )
+        controller = BrowserInitializer.init_controller()
+        logger.info("Browser initialized")
 
-        browser, controller = init_browser()
         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),
+            use_vision=self.inputs.get("use_vision", True),
         )
 
         # 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,10 @@ 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]]]]
+    downloads_path: Optional[str]
+    use_vision: Optional[bool]
 
 
 class BrowserUseOutputs(TypedDict):
