Merge pull request #1 from browserbase/anirudh/new-sdk

Use New SDK

Author: kamath

README.md

@@ -12,7 +12,10 @@
 </p>
 <br/>
 
-## Selenium (Python) with Browserbase
+## Selenium with Browserbase
+
+## Introduction
+
 Browserbase is the best developer platform to reliably run, manage, and monitor headless browsers.
 
 Get browsers' complete control and leverage Browserbase's
@@ -22,7 +25,6 @@ and LLM data retrievals.
 
 **Get started in under one minute** with Selenium.
 
-
 ## Setup
 
 ### 1. Install dependencies
@@ -31,21 +33,30 @@ and LLM data retrievals.
 pip install -r requirements.txt
 ```
 
+Alternatively, we suggest using [Poetry](https://python-poetry.org/) to manage your dependencies.
+
+```bash
+poetry install
+```
 
-### 2. Get your Browserbase API Key:
+### 2. Create `.env` file:
+
+```bash
+cp .env.example .env
+```
+
+### 3. Get your Browserbase API Key:
 
 - [Create an account](https://www.browserbase.com/sign-up) or [log in to Browserbase](https://www.browserbase.com/sign-in)
-- Copy your API Key and Project ID [from your Settings page](https://www.browserbase.com/settings)
+- Copy your API Key and Project ID [from your Settings page](https://www.browserbase.com/settings) into the `.env` file
 
-### 3. Run the script:
+### 4. Run the script:
 
 ```bash
-BROWSERBASE_API_KEY=xxxx BROWSERBASE_PROJECT_ID==xxxx python main.py
+python main.py # or poetry run python main.py
 ```
 
-
 ## Further reading
 
-- [See how to leverage the Session Debugger for faster development](https://docs.browserbase.com/guides/browser-remote-control#accelerate-your-local-development-with-remote-debugging)
-- [Learn more about Browserbase infrastructure](https://docs.browserbase.com/under-the-hood)
-- [Explore the Sessions API](https://docs.browserbase.com/api-reference/list-all-sessions)
+- [Explore the Browserbase Python SDK](https://docs.browserbase.com/sdk/python)
+- [Explore the Selenium Python API](https://selenium-python.readthedocs.io/api.html)

main.py

@@ -1,40 +1,67 @@
+from typing import Dict
+
 from selenium import webdriver
 from selenium.webdriver.remote.remote_connection import RemoteConnection
-from selenium.webdriver.common.desired_capabilities import DesiredCapabilities
-import requests
+from browserbase import Browserbase
+from dotenv import load_dotenv
 import os
 
-def create_session():
-    url = 'https://www.browserbase.com/v1/sessions'
-    headers = {'Content-Type': 'application/json', 'x-bb-api-key': os.environ["BROWSERBASE_API_KEY"]}
-    response = requests.post(url, json={ "projectId": os.environ["BROWSERBASE_PROJECT_ID"] }, headers=headers)
-    return response.json()['id']
+load_dotenv()
+
+BROWSERBASE_API_KEY = os.getenv("BROWSERBASE_API_KEY")
+BROWSERBASE_PROJECT_ID = os.getenv("BROWSERBASE_PROJECT_ID")
+
+bb = Browserbase(api_key=BROWSERBASE_API_KEY)
+
+
+class BrowserbaseConnection(RemoteConnection):
+    """
+    Manage a single session with Browserbase.
+    """
+
+    session_id: str
+
+    def __init__(self, session_id: str, *args, **kwargs):  # type: ignore
+        super().__init__(*args, **kwargs)  # type: ignore
+        self.session_id = session_id
+
+    def get_remote_connection_headers(  # type: ignore
+        self, parsed_url: str, keep_alive: bool = False
+    ) -> Dict[str, str]:
+        headers = super().get_remote_connection_headers(parsed_url, keep_alive)  # type: ignore
+
+        # Update headers to include the Browserbase required information
+        headers["x-bb-api-key"] = BROWSERBASE_API_KEY
+        headers["session-id"] = self.session_id
+
+        return headers  # type: ignore
 
 
-class CustomRemoteConnection(RemoteConnection):
-    _session_id = None
+def run() -> None:
+    # Use the custom class to create and connect to a new browser session
+    session = bb.sessions.create(project_id=BROWSERBASE_PROJECT_ID)
+    connection = BrowserbaseConnection(session.id, session.selenium_remote_url)
+    driver = webdriver.Remote(
+        command_executor=connection, options=webdriver.ChromeOptions()  # type: ignore
+    )
 
-    def __init__(self, remote_server_addr: str, session_id: str):
-        super().__init__(remote_server_addr)
-        self._session_id = session_id
+    # Print a bit of info about the browser we've connected to
+    print(
+        "Connected to Browserbase",
+        f"{driver.name} version {driver.caps['browserVersion']}",  # type: ignore
+    )
 
-    def get_remote_connection_headers(self, parsed_url, keep_alive=False):
-        headers = super().get_remote_connection_headers(parsed_url, keep_alive)
-        headers.update({'x-bb-api-key': os.environ["BROWSERBASE_API_KEY"]})
-        headers.update({'session-id': self._session_id})
-        return headers
+    try:
+        # Perform our browser commands
+        driver.get("https://www.sfmoma.org")
+        print(f"At URL: {driver.current_url} | Title: {driver.title}")
+        assert driver.current_url == "https://www.sfmoma.org/"
+        assert driver.title == "SFMOMA"
 
+    finally:
+        # Make sure to quit the driver so your session is ended!
+        driver.quit()
 
-def run():
-    session_id = create_session()
-    custom_conn = CustomRemoteConnection('http://connect.browserbase.com/webdriver', session_id)
-    options = webdriver.ChromeOptions()
-    options.debugger_address = "localhost:9223"
-    driver = webdriver.Remote(custom_conn, options=options)
-    driver.get("https://www.browserbase.com")
-    get_title = driver.title
-    print(get_title)
-    # Make sure to quit the driver so your session is ended!
-    driver.quit()
 
-run()
+if __name__ == "__main__":
+    run()