Merge pull request #77 from browserbase/miguel/stg-463-consolidate-client-base

consolidate base-client

Author: miguelg719

stagehand/api.py

@@ -1,169 +0,0 @@
-import logging
-import os
-from abc import ABC, abstractmethod
-from collections.abc import Awaitable
-from typing import Any, Callable, Optional
-
-from browserbase.types import SessionCreateParams as BrowserbaseSessionCreateParams
-
-from .config import StagehandConfig
-from .page import StagehandPage
-from .utils import default_log_handler
-
-logger = logging.getLogger(__name__)
-
-
-class StagehandBase(ABC):
-    """
-    Base class for Stagehand client implementations.
-    Defines the common interface and functionality.
-    """
-
-    def __init__(
-        self,
-        config: Optional["StagehandConfig"] = None,
-        server_url: Optional[str] = None,
-        session_id: Optional[str] = None,
-        browserbase_api_key: Optional[str] = None,
-        browserbase_project_id: Optional[str] = None,
-        model_api_key: Optional[str] = None,
-        on_log: Optional[
-            Callable[[dict[str, Any]], Awaitable[None]]
-        ] = default_log_handler,
-        verbose: int = 1,
-        model_name: Optional[str] = None,
-        dom_settle_timeout_ms: Optional[int] = None,
-        timeout_settings: Optional[Any] = None,
-        stream_response: Optional[bool] = None,
-        model_client_options: Optional[dict[str, Any]] = None,
-        self_heal: Optional[bool] = None,
-        wait_for_captcha_solves: Optional[bool] = None,
-        system_prompt: Optional[str] = None,
-        browserbase_session_create_params: Optional[
-            BrowserbaseSessionCreateParams
-        ] = None,
-        enable_caching: Optional[bool] = None,
-    ):
-        """
-        Initialize the Stagehand client with common configuration.
-
-        Args:
-            config (Optional[StagehandConfig]): Configuration object that can provide all settings.
-            server_url (Optional[str]): URL of the Stagehand server.
-            session_id (Optional[str]): Existing session ID to resume.
-            browserbase_api_key (Optional[str]): Browserbase API key.
-            browserbase_project_id (Optional[str]): Browserbase project ID.
-            model_api_key (Optional[str]): Model provider API key.
-            on_log (Optional[Callable]): Callback for log events.
-            verbose (int): Verbosity level.
-            model_name (Optional[str]): Name of the model to use.
-            dom_settle_timeout_ms (Optional[int]): Time for DOM to settle in ms.
-            debug_dom (Optional[bool]): Whether to enable DOM debugging.
-            timeout_settings (Optional[float]): Request timeout in seconds.
-            stream_response (Optional[bool]): Whether to stream responses.
-            model_client_options (Optional[dict]): Options for the model client.
-            self_heal (Optional[bool]): Whether to enable self-healing.
-            wait_for_captcha_solves (Optional[bool]): Whether to wait for CAPTCHA solves.
-            act_timeout_ms (Optional[int]): Timeout for act commands in ms.
-            system_prompt (Optional[str]): System prompt for LLM interactions.
-            browserbase_session_create_params (Optional[BrowserbaseSessionCreateParams]): Params for Browserbase session creation.
-            enable_caching (Optional[bool]): Whether to enable caching functionality.
-        """
-        self.server_url = server_url or os.getenv("STAGEHAND_SERVER_URL")
-
-        if config:
-            self.browserbase_api_key = (
-                config.api_key
-                or browserbase_api_key
-                or os.getenv("BROWSERBASE_API_KEY")
-            )
-            self.browserbase_project_id = (
-                config.project_id
-                or browserbase_project_id
-                or os.getenv("BROWSERBASE_PROJECT_ID")
-            )
-            self.session_id = config.browserbase_session_id or session_id
-            self.model_name = config.model_name or model_name
-            self.dom_settle_timeout_ms = (
-                config.dom_settle_timeout_ms or dom_settle_timeout_ms
-            )
-            self.self_heal = (
-                config.self_heal if config.self_heal is not None else self_heal
-            )
-            self.wait_for_captcha_solves = (
-                config.wait_for_captcha_solves
-                if config.wait_for_captcha_solves is not None
-                else wait_for_captcha_solves
-            )
-            self.system_prompt = config.system_prompt or system_prompt
-            self.browserbase_session_create_params = (
-                config.browserbase_session_create_params
-                or browserbase_session_create_params
-            )
-            self.enable_caching = (
-                config.enable_caching
-                if config.enable_caching is not None
-                else enable_caching
-            )
-            self.verbose = config.verbose if config.verbose is not None else verbose
-        else:
-            self.browserbase_api_key = browserbase_api_key or os.getenv(
-                "BROWSERBASE_API_KEY"
-            )
-            self.browserbase_project_id = browserbase_project_id or os.getenv(
-                "BROWSERBASE_PROJECT_ID"
-            )
-            self.session_id = session_id
-            self.model_name = model_name
-            self.dom_settle_timeout_ms = dom_settle_timeout_ms
-            self.self_heal = self_heal
-            self.wait_for_captcha_solves = wait_for_captcha_solves
-            self.system_prompt = system_prompt
-            self.browserbase_session_create_params = browserbase_session_create_params
-            self.enable_caching = enable_caching
-            self.verbose = verbose
-
-        # Handle model-related settings directly
-        self.model_api_key = model_api_key or os.getenv("MODEL_API_KEY")
-        self.model_client_options = model_client_options or {}
-        if self.model_api_key and "apiKey" not in self.model_client_options:
-            self.model_client_options["apiKey"] = self.model_api_key
-
-        # Handle streaming response setting directly
-        self.streamed_response = (
-            stream_response if stream_response is not None else True
-        )
-
-        self.on_log = on_log
-        self.timeout_settings = timeout_settings or 180.0
-
-        self._initialized = False
-        self._closed = False
-        self.page: Optional[StagehandPage] = None
-
-        # Validate essential fields if session_id was provided
-        if self.session_id:
-            if not self.browserbase_api_key:
-                raise ValueError(
-                    "browserbase_api_key is required (or set BROWSERBASE_API_KEY in env)."
-                )
-            if not self.browserbase_project_id:
-                raise ValueError(
-                    "browserbase_project_id is required (or set BROWSERBASE_PROJECT_ID in env)."
-                )
-
-    @abstractmethod
-    def init(self):
-        """
-        Initialize the Stagehand client.
-        Must be implemented by subclasses.
-        """
-        pass
-
-    @abstractmethod
-    def close(self):
-        """
-        Clean up resources.
-        Must be implemented by subclasses.
-        """
-        pass

stagehand/base.py

@@ -10,6 +10,7 @@
 
 import httpx
 from browserbase import Browserbase
+from browserbase.types import SessionCreateParams as BrowserbaseSessionCreateParams
 from dotenv import load_dotenv
 from playwright.async_api import (
     BrowserContext,
@@ -18,18 +19,17 @@
 )
 from playwright.async_api import Page as PlaywrightPage
 
-from .base import StagehandBase
 from .config import StagehandConfig
 from .context import StagehandContext
 from .llm import LLMClient
 from .metrics import StagehandFunctionName, StagehandMetrics
 from .page import StagehandPage
-from .utils import StagehandLogger, convert_dict_keys_to_camel_case
+from .utils import StagehandLogger, convert_dict_keys_to_camel_case, default_log_handler
 
 load_dotenv()
 
 
-class Stagehand(StagehandBase):
+class Stagehand:
     """
     Python client for interacting with a running Stagehand server and Browserbase remote headless browser.
 
@@ -42,6 +42,7 @@ class Stagehand(StagehandBase):
 
     def __init__(
         self,
+        *,
         config: Optional[StagehandConfig] = None,
         server_url: Optional[str] = None,
         session_id: Optional[str] = None,
@@ -62,6 +63,7 @@ def __init__(
         use_rich_logging: bool = True,
         env: Literal["BROWSERBASE", "LOCAL"] = None,
         local_browser_launch_options: Optional[dict[str, Any]] = None,
+        browserbase_session_create_params: Optional[BrowserbaseSessionCreateParams] = None,
     ):
         """
         Initialize the Stagehand client.
@@ -89,31 +91,81 @@ def __init__(
             local_browser_launch_options (Optional[dict[str, Any]]): Options for launching the local browser context
                 when env="LOCAL". See Playwright's launch_persistent_context documentation.
                 Common keys: 'headless', 'user_data_dir', 'downloads_path', 'viewport', 'locale', 'proxy', 'args', 'cdp_url'.
+            browserbase_session_create_params (Optional[BrowserbaseSessionCreateParams]): Params for Browserbase session creation.
         """
-        super().__init__(
-            config=config,
-            server_url=server_url,
-            session_id=session_id,
-            browserbase_api_key=browserbase_api_key,
-            browserbase_project_id=browserbase_project_id,
-            model_api_key=model_api_key,
-            on_log=on_log,
-            verbose=verbose,
-            model_name=model_name,
-            dom_settle_timeout_ms=dom_settle_timeout_ms,
-            timeout_settings=timeout_settings,
-            stream_response=stream_response,
-            model_client_options=model_client_options,
-            self_heal=self_heal,
-            wait_for_captcha_solves=wait_for_captcha_solves,
-            system_prompt=system_prompt,
+        # Initialize configuration from config object or individual parameters
+        self.server_url = server_url or os.getenv("STAGEHAND_SERVER_URL")
+
+        if config:
+            self.browserbase_api_key = (
+                config.api_key
+                or browserbase_api_key
+                or os.getenv("BROWSERBASE_API_KEY")
+            )
+            self.browserbase_project_id = (
+                config.project_id
+                or browserbase_project_id
+                or os.getenv("BROWSERBASE_PROJECT_ID")
+            )
+            self.session_id = config.browserbase_session_id or session_id
+            self.model_name = config.model_name or model_name
+            self.dom_settle_timeout_ms = (
+                config.dom_settle_timeout_ms or dom_settle_timeout_ms
+            )
+            self.self_heal = (
+                config.self_heal if config.self_heal is not None else self_heal
+            )
+            self.wait_for_captcha_solves = (
+                config.wait_for_captcha_solves
+                if config.wait_for_captcha_solves is not None
+                else wait_for_captcha_solves
+            )
+            self.system_prompt = config.system_prompt or system_prompt
+            self.browserbase_session_create_params = (
+                config.browserbase_session_create_params
+                or browserbase_session_create_params
+            )
+            self.verbose = config.verbose if config.verbose is not None else verbose
+        else:
+            self.browserbase_api_key = browserbase_api_key or os.getenv(
+                "BROWSERBASE_API_KEY"
+            )
+            self.browserbase_project_id = browserbase_project_id or os.getenv(
+                "BROWSERBASE_PROJECT_ID"
+            )
+            self.session_id = session_id
+            self.model_name = model_name
+            self.dom_settle_timeout_ms = dom_settle_timeout_ms
+            self.self_heal = self_heal
+            self.wait_for_captcha_solves = wait_for_captcha_solves
+            self.system_prompt = system_prompt
+            self.browserbase_session_create_params = browserbase_session_create_params
+            self.verbose = verbose
+
+        # Handle model-related settings directly
+        self.model_api_key = model_api_key or os.getenv("MODEL_API_KEY")
+        self.model_client_options = model_client_options or {}
+        if self.model_api_key and "apiKey" not in self.model_client_options:
+            self.model_client_options["apiKey"] = self.model_api_key
+
+        # Handle streaming response setting directly
+        self.streamed_response = (
+            stream_response if stream_response is not None else True
+        )
+
+        self.on_log = on_log or default_log_handler
+        self.timeout_settings = timeout_settings or httpx.Timeout(
+            connect=180.0,
+            read=180.0,
+            write=180.0,
+            pool=180.0,
         )
 
         self.env = env.upper() if env else "BROWSERBASE"
         self.local_browser_launch_options = (
             getattr(config, "local_browser_launch_options", {})
             if config
-            else local_browser_launch_options
+            else local_browser_launch_options if local_browser_launch_options else {}
         )
         self._local_user_data_dir_temp: Optional[Path] = (
             None  # To store path if created temporarily
@@ -127,6 +179,11 @@ def __init__(
         if self.env not in ["BROWSERBASE", "LOCAL"]:
             raise ValueError("env must be either 'BROWSERBASE' or 'LOCAL'")
 
+        # Initialize the centralized logger with the specified verbosity
+        self.logger = StagehandLogger(
+            verbose=self.verbose, external_logger=on_log, use_rich=use_rich_logging
+        )
+
         # If using BROWSERBASE, session_id or creation params are needed
         if self.env == "BROWSERBASE":
             if not self.session_id:
@@ -155,18 +212,7 @@ def __init__(
                         "browserbase_project_id is required for BROWSERBASE env with existing session_id (or set BROWSERBASE_PROJECT_ID in env)."
                     )
 
-        # Initialize the centralized logger with the specified verbosity
-        self.logger = StagehandLogger(
-            verbose=self.verbose, external_logger=on_log, use_rich=use_rich_logging
-        )
-
         self.httpx_client = httpx_client
-        self.timeout_settings = timeout_settings or httpx.Timeout(
-            connect=180.0,
-            read=180.0,
-            write=180.0,
-            pool=180.0,
-        )
         self._client: Optional[httpx.AsyncClient] = (
             None  # Used for server communication in BROWSERBASE
         )

stagehand/browser.py

@@ -77,7 +77,7 @@ class StagehandConfig(BaseModel):
         description="System prompt to use for LLM interactions",
     )
     local_browser_launch_options: Optional[dict[str, Any]] = Field(
-        None,
+        {},
         alias="localBrowserLaunchOptions",
         description="Local browser launch options",
     )

stagehand/client.py

@@ -87,12 +87,14 @@ async def wrapped_new_page(*args, **kwargs):
 
             return wrapped_new_page
         elif name == "pages":
-            # Wrap the pages method to return StagehandPage objects
-            def wrapped_pages(*args, **kwargs):
-                pw_pages = self._context.pages(*args, **kwargs)
-                # This is synchronous, so we can't await here
-                # We'll return the unwrapped pages and let the caller handle wrapping if needed
-                return pw_pages
+            async def wrapped_pages():
+                pw_pages = self._context.pages
+                # Return StagehandPage objects
+                result = []
+                for pw_page in pw_pages:
+                    stagehand_page = await self.get_stagehand_page(pw_page)
+                    result.append(stagehand_page)
+                return result
 
             return wrapped_pages
         return attr

stagehand/config.py

@@ -444,7 +444,7 @@ async def handle_possible_page_navigation(
             },
         )
 
-    if new_opened_tab:
+    if new_opened_tab and new_opened_tab.url != 'about:blank':
         logger.info(
             message="new page detected (new tab) with URL",
             category="action",