|
14 | 14 | import asyncio |
15 | 15 | import time |
16 | 16 | from typing import Literal |
| 17 | +from typing import Optional |
17 | 18 |
|
18 | 19 | from google.adk.tools.computer_use.base_computer import BaseComputer |
19 | 20 | from google.adk.tools.computer_use.base_computer import ComputerEnvironment |
@@ -79,29 +80,61 @@ def __init__( |
79 | 80 | initial_url: str = "https://www.google.com", |
80 | 81 | search_engine_url: str = "https://www.google.com", |
81 | 82 | highlight_mouse: bool = False, |
| 83 | + user_data_dir: Optional[str] = None, |
82 | 84 | ): |
83 | 85 | self._initial_url = initial_url |
84 | 86 | self._screen_size = screen_size |
85 | 87 | self._search_engine_url = search_engine_url |
86 | 88 | self._highlight_mouse = highlight_mouse |
| 89 | + self._user_data_dir = user_data_dir |
87 | 90 |
|
88 | 91 | @override |
89 | 92 | async def initialize(self): |
90 | 93 | print("Creating session...") |
91 | 94 | self._playwright = await async_playwright().start() |
92 | | - self._browser = await self._playwright.chromium.launch( |
93 | | - args=["--disable-blink-features=AutomationControlled"], |
94 | | - headless=False, |
95 | | - ) |
96 | | - self._context = await self._browser.new_context( |
97 | | - viewport={ |
98 | | - "width": self._screen_size[0], |
99 | | - "height": self._screen_size[1], |
100 | | - } |
101 | | - ) |
102 | | - self._page = await self._context.new_page() |
103 | | - await self._page.goto(self._initial_url) |
104 | 95 |
|
| 96 | + # Define common arguments for both launch types |
| 97 | + browser_args = [ |
| 98 | + "--disable-blink-features=AutomationControlled", |
| 99 | + "--disable-gpu", |
| 100 | + ] |
| 101 | + |
| 102 | + if self._user_data_dir: |
| 103 | + termcolor.cprint( |
| 104 | + f"Starting playwright with persistent profile: {self._user_data_dir}", |
| 105 | + color="yellow", |
| 106 | + attrs=["bold"], |
| 107 | + ) |
| 108 | + # Use a persistent context if user_data_dir is provided |
| 109 | + self._context = await self._playwright.chromium.launch_persistent_context( |
| 110 | + self._user_data_dir, |
| 111 | + headless=False, |
| 112 | + args=browser_args, |
| 113 | + ) |
| 114 | + self._browser = self._context.browser |
| 115 | + else: |
| 116 | + termcolor.cprint( |
| 117 | + "Starting playwright with a temporary profile.", |
| 118 | + color="yellow", |
| 119 | + attrs=["bold"], |
| 120 | + ) |
| 121 | + # Launch a temporary browser instance if user_data_dir is not provided |
| 122 | + self._browser = await self._playwright.chromium.launch( |
| 123 | + args=browser_args, |
| 124 | + headless=False, |
| 125 | + ) |
| 126 | + self._context = await self._browser.new_context() |
| 127 | + |
| 128 | + if not self._context.pages: |
| 129 | + self._page = await self._context.new_page() |
| 130 | + await self._page.goto(self._initial_url) |
| 131 | + else: |
| 132 | + self._page = self._context.pages[0] # Use existing page if any |
| 133 | + |
| 134 | + await self._page.set_viewport_size({ |
| 135 | + "width": self._screen_size[0], |
| 136 | + "height": self._screen_size[1], |
| 137 | + }) |
105 | 138 | termcolor.cprint( |
106 | 139 | f"Started local playwright.", |
107 | 140 | color="green", |
|
0 commit comments