|
| 1 | +import os |
| 2 | +from typing import Tuple, Optional |
| 3 | +from playwright.sync_api import Browser, Page, Error as PlaywrightError |
| 4 | +from .base_playwright import BasePlaywrightComputer |
| 5 | +from dotenv import load_dotenv |
| 6 | +import base64 |
| 7 | +from steel import Steel |
| 8 | + |
| 9 | +load_dotenv() |
| 10 | + |
| 11 | + |
| 12 | +class SteelBrowser(BasePlaywrightComputer): |
| 13 | + """ |
| 14 | + Steel is an open-source browser API purpose-built for AI agents. |
| 15 | + It provides a simple API for controlling browsers remotely, with features like: |
| 16 | + - Session recording and replay |
| 17 | + - Built-in proxy support with residential IPs |
| 18 | + - Anti-bot protection |
| 19 | + - Browser context management |
| 20 | +
|
| 21 | + To get started: |
| 22 | + 1. Sign up at https://steel.dev |
| 23 | + 2. Get your API key from the dashboard |
| 24 | + 3. Add STEEL_API_KEY to your .env file |
| 25 | +
|
| 26 | + If you're running Steel locally or self-hosted, add the following to your .env file: |
| 27 | + STEEL_API_KEY=your_api_key |
| 28 | + STEEL_API_URL=http://localhost:3000 (or your self-hosted URL) |
| 29 | +
|
| 30 | + For more information, visit: https://docs.steel.dev |
| 31 | +
|
| 32 | + IMPORTANT: The `goto` tool, as defined in playwright_with_custom_functions.py, is strongly recommended when using the Steel computer. |
| 33 | + Make sure to include this tool in your configuration when using the Steel computer. |
| 34 | + """ |
| 35 | + |
| 36 | + def __init__( |
| 37 | + self, |
| 38 | + width: int = 1024, |
| 39 | + height: int = 768, |
| 40 | + proxy: bool = False, |
| 41 | + solve_captcha: bool = False, |
| 42 | + virtual_mouse: bool = True, |
| 43 | + session_timeout: int = 900000, # 15 minutes default |
| 44 | + ad_blocker: bool = True, |
| 45 | + start_url: str = "https://bing.com" |
| 46 | + ): |
| 47 | + """ |
| 48 | + Initialize the Steel browser instance. |
| 49 | +
|
| 50 | + Args: |
| 51 | + width (int): Browser viewport width. Default is 1024. |
| 52 | + height (int): Browser viewport height. Default is 768. |
| 53 | + use_proxy (bool): Whether to use Steel's proxy network (residential IPs). Default is False. |
| 54 | + solve_captcha (bool): Whether to enable automatic CAPTCHA solving. Default is False. |
| 55 | + virtual_mouse (bool): Whether to show a virtual mouse cursor. Default is True. |
| 56 | + session_timeout (int): Session timeout in milliseconds. Default is 5 minutes. |
| 57 | + ad_blocker (bool): Whether to enable ad blocking. Default is True. |
| 58 | + start_url (str): The initial URL to navigate to. Default is "https://bing.com". |
| 59 | + """ |
| 60 | + super().__init__() |
| 61 | + |
| 62 | + # Initialize Steel client |
| 63 | + self.client = Steel(steel_api_key=os.getenv("STEEL_API_KEY")) |
| 64 | + self.dimensions = (width, height) |
| 65 | + self.proxy = proxy |
| 66 | + self.solve_captcha = solve_captcha |
| 67 | + self.virtual_mouse = virtual_mouse |
| 68 | + self.session_timeout = session_timeout |
| 69 | + self.ad_blocker = ad_blocker |
| 70 | + self.start_url = start_url |
| 71 | + self.session = None |
| 72 | + |
| 73 | + def _get_browser_and_page(self) -> Tuple[Browser, Page]: |
| 74 | + """ |
| 75 | + Create a Steel browser session and connect to it. |
| 76 | +
|
| 77 | + Returns: |
| 78 | + Tuple[Browser, Page]: A tuple containing the connected browser and page objects. |
| 79 | + """ |
| 80 | + # Create Steel session |
| 81 | + width, height = self.dimensions |
| 82 | + self.session = self.client.sessions.create( |
| 83 | + use_proxy=self.proxy, |
| 84 | + solve_captcha=self.solve_captcha, |
| 85 | + api_timeout=self.session_timeout, |
| 86 | + block_ads=self.ad_blocker, |
| 87 | + dimensions={"width": width, "height": height} |
| 88 | + ) |
| 89 | + |
| 90 | + print("Steel Session created successfully!") |
| 91 | + print(f"View live session at: {self.session.session_viewer_url}") |
| 92 | + |
| 93 | + # Connect to the remote browser using Steel's connection URL |
| 94 | + browser = self._playwright.chromium.connect_over_cdp( |
| 95 | + f"wss://connect.steel.dev?apiKey={os.getenv('STEEL_API_KEY')}&sessionId={self.session.id}" |
| 96 | + ) |
| 97 | + context = browser.contexts[0] |
| 98 | + |
| 99 | + # Set up page event handlers |
| 100 | + context.on("page", self._handle_new_page) |
| 101 | + |
| 102 | + # Add virtual mouse cursor if enabled |
| 103 | + if self.virtual_mouse: |
| 104 | + context.add_init_script(""" |
| 105 | + // Only run in the top frame |
| 106 | + if (window.self === window.top) { |
| 107 | + function initCursor() { |
| 108 | + const CURSOR_ID = '__cursor__'; |
| 109 | + if (document.getElementById(CURSOR_ID)) return; |
| 110 | +
|
| 111 | + const cursor = document.createElement('div'); |
| 112 | + cursor.id = CURSOR_ID; |
| 113 | + Object.assign(cursor.style, { |
| 114 | + position: 'fixed', |
| 115 | + top: '0px', |
| 116 | + left: '0px', |
| 117 | + width: '20px', |
| 118 | + height: '20px', |
| 119 | + backgroundImage: 'url("data:image/svg+xml;utf8,<svg width=\\'16\\' height=\\'16\\' viewBox=\\'0 0 20 20\\' fill=\\'black\\' outline=\\'white\\' xmlns=\\'http://www.w3.org/2000/svg\\'><path d=\\'M15.8089 7.22221C15.9333 7.00888 15.9911 6.78221 15.9822 6.54221C15.9733 6.29333 15.8978 6.06667 15.7555 5.86221C15.6133 5.66667 15.4311 5.52445 15.2089 5.43555L1.70222 0.0888888C1.47111 0 1.23555 -0.0222222 0.995555 0.0222222C0.746667 0.0755555 0.537779 0.186667 0.368888 0.355555C0.191111 0.533333 0.0755555 0.746667 0.0222222 0.995555C-0.0222222 1.23555 0 1.47111 0.0888888 1.70222L5.43555 15.2222C5.52445 15.4445 5.66667 15.6267 5.86221 15.7689C6.06667 15.9111 6.28888 15.9867 6.52888 15.9955H6.58221C6.82221 15.9955 7.04445 15.9333 7.24888 15.8089C7.44445 15.6845 7.59555 15.52 7.70221 15.3155L10.2089 10.2222L15.3022 7.70221C15.5155 7.59555 15.6845 7.43555 15.8089 7.22221Z\\' ></path></svg>")', |
| 120 | + backgroundSize: 'cover', |
| 121 | + pointerEvents: 'none', |
| 122 | + zIndex: '99999', |
| 123 | + transform: 'translate(-2px, -2px)', |
| 124 | + }); |
| 125 | +
|
| 126 | + document.body.appendChild(cursor); |
| 127 | + document.addEventListener("mousemove", (e) => { |
| 128 | + cursor.style.top = e.clientY + "px"; |
| 129 | + cursor.style.left = e.clientX + "px"; |
| 130 | + }); |
| 131 | + } |
| 132 | +
|
| 133 | + requestAnimationFrame(function checkBody() { |
| 134 | + if (document.body) { |
| 135 | + initCursor(); |
| 136 | + } else { |
| 137 | + requestAnimationFrame(checkBody); |
| 138 | + } |
| 139 | + }); |
| 140 | + } |
| 141 | + """) |
| 142 | + |
| 143 | + page = context.pages[0] |
| 144 | + page.on("close", self._handle_page_close) |
| 145 | + |
| 146 | + # Navigate to start URL |
| 147 | + page.goto(self.start_url) |
| 148 | + |
| 149 | + return browser, page |
| 150 | + |
| 151 | + def _handle_new_page(self, page: Page): |
| 152 | + """Handle creation of a new page.""" |
| 153 | + print("New page created") |
| 154 | + self._page = page |
| 155 | + page.on("close", self._handle_page_close) |
| 156 | + |
| 157 | + def _handle_page_close(self, page: Page): |
| 158 | + """Handle page closure.""" |
| 159 | + print("Page closed") |
| 160 | + if self._page == page: |
| 161 | + if self._browser.contexts[0].pages: |
| 162 | + self._page = self._browser.contexts[0].pages[-1] |
| 163 | + else: |
| 164 | + print("Warning: All pages have been closed.") |
| 165 | + self._page = None |
| 166 | + |
| 167 | + def __exit__(self, exc_type, exc_val, exc_tb): |
| 168 | + """Clean up resources when exiting.""" |
| 169 | + if self._page: |
| 170 | + self._page.close() |
| 171 | + if self._browser: |
| 172 | + self._browser.close() |
| 173 | + if self._playwright: |
| 174 | + self._playwright.stop() |
| 175 | + |
| 176 | + # Release the Steel session |
| 177 | + if self.session: |
| 178 | + print("Releasing Steel session...") |
| 179 | + self.client.sessions.release(self.session.id) |
| 180 | + print( |
| 181 | + f"Session completed. View replay at {self.session.session_viewer_url}") |
| 182 | + |
| 183 | + def screenshot(self) -> str: |
| 184 | + """ |
| 185 | + Capture a screenshot of the current viewport using CDP. |
| 186 | +
|
| 187 | + Returns: |
| 188 | + str: Base64 encoded screenshot data |
| 189 | + """ |
| 190 | + try: |
| 191 | + cdp_session = self._page.context.new_cdp_session(self._page) |
| 192 | + result = cdp_session.send("Page.captureScreenshot", { |
| 193 | + "format": "png", |
| 194 | + "fromSurface": True |
| 195 | + }) |
| 196 | + return result['data'] |
| 197 | + except PlaywrightError as error: |
| 198 | + print( |
| 199 | + f"CDP screenshot failed, falling back to standard screenshot: {error}") |
| 200 | + return super().screenshot() |
0 commit comments