refactor context tree (Skyvern-AI#212)

feiofil · Apr 21, 2024 · cc6ae8b · cc6ae8b
1 parent 02db2a9
commit cc6ae8b
Show file tree

Hide file tree

Showing 4 changed files with 215 additions and 61 deletions.
diff --git a/skyvern/webeye/browser_factory.py b/skyvern/webeye/browser_factory.py
@@ -135,11 +135,13 @@ def __init__(
         browser_context: BrowserContext | None = None,
         page: Page | None = None,
         browser_artifacts: BrowserArtifacts = BrowserArtifacts(),
+        new_context_tree: bool = False,
     ):
         self.pw = pw
         self.browser_context = browser_context
         self.page = page
         self.browser_artifacts = browser_artifacts
+        self.new_context_tree = new_context_tree
 
     async def _close_all_other_pages(self) -> None:
         if not self.browser_context or not self.page:

diff --git a/skyvern/webeye/browser_manager.py b/skyvern/webeye/browser_manager.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import random
+
 import structlog
 from playwright.async_api import Browser, Playwright, async_playwright
 
@@ -23,13 +25,19 @@ def __new__(cls) -> BrowserManager:
 
     @staticmethod
     async def _create_browser_state(
-        proxy_location: ProxyLocation | None = None, url: str | None = None
+        proxy_location: ProxyLocation | None = None, url: str | None = None, new_context_tree: bool = False
     ) -> BrowserState:
         pw = await async_playwright().start()
         browser_context, browser_artifacts = await BrowserContextFactory.create_browser_context(
             pw, proxy_location=proxy_location, url=url
         )
-        return BrowserState(pw=pw, browser_context=browser_context, page=None, browser_artifacts=browser_artifacts)
+        return BrowserState(
+            pw=pw,
+            browser_context=browser_context,
+            page=None,
+            browser_artifacts=browser_artifacts,
+            new_context_tree=new_context_tree,
+        )
 
     async def get_or_create_for_task(self, task: Task) -> BrowserState:
         if task.task_id in self.pages:
@@ -42,8 +50,11 @@ async def get_or_create_for_task(self, task: Task) -> BrowserState:
             )
             self.pages[task.task_id] = self.pages[task.workflow_run_id]
             return self.pages[task.task_id]
-        LOG.info("Creating browser state for task", task_id=task.task_id)
-        browser_state = await self._create_browser_state(task.proxy_location, task.url)
+
+        # TODO: percentage to use new context tree, starting from 20%
+        new_ctx = random.choices([False, True], weights=[0.8, 0.2], k=1)[0]
+        LOG.info("Creating browser state for task", task_id=task.task_id, new_ctx=new_ctx)
+        browser_state = await self._create_browser_state(task.proxy_location, task.url, new_ctx)
 
         # The URL here is only used when creating a new page, and not when using an existing page.
         # This will make sure browser_state.page is not None.