Merge branch 'main' into fm/stg-474-more-tests

merge main

Author: filip-michalsky

stagehand/__init__.py

@@ -1,7 +1,7 @@
 from .agent import Agent
-from .client import Stagehand
 from .config import StagehandConfig, default_config
 from .handlers.observe_handler import ObserveHandler
+from .main import Stagehand
 from .metrics import StagehandFunctionName, StagehandMetrics
 from .page import StagehandPage
 from .schemas import (

stagehand/api.py

@@ -73,7 +73,7 @@ async def _create_session(self):
         "x-language": "python",
     }
 
-    client = self.httpx_client or httpx.AsyncClient(timeout=self.timeout_settings)
+    client = httpx.AsyncClient(timeout=self.timeout_settings)
     async with client:
         resp = await client.post(
             f"{self.api_url}/sessions/start",
@@ -109,7 +109,7 @@ async def _execute(self, method: str, payload: dict[str, Any]) -> Any:
     # Convert snake_case keys to camelCase for the API
     modified_payload = convert_dict_keys_to_camel_case(payload)
 
-    client = self.httpx_client or httpx.AsyncClient(timeout=self.timeout_settings)
+    client = httpx.AsyncClient(timeout=self.timeout_settings)
 
     async with client:
         try:

stagehand/config.py

@@ -1,4 +1,4 @@
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Literal, Optional
 
 from browserbase.types import SessionCreateParams as BrowserbaseSessionCreateParams
 from pydantic import BaseModel, ConfigDict, Field
@@ -12,37 +12,49 @@ class StagehandConfig(BaseModel):
 
     Attributes:
         env (str): Environment type. 'BROWSERBASE' for remote usage
-        api_key (Optional[str]): API key for authentication.
-        project_id (Optional[str]): Project identifier.
-        headless (bool): Run browser in headless mode.
-        logger (Optional[Callable[[Any], None]]): Custom logging function.
-        dom_settle_timeout_ms (Optional[int]): Timeout for DOM to settle (in milliseconds).
+        api_key (Optional[str]): BrowserbaseAPI key for authentication.
+        project_id (Optional[str]): Browserbase Project identifier.
+        api_url (Optional[str]): Stagehand API URL.
         browserbase_session_create_params (Optional[BrowserbaseSessionCreateParams]): Browserbase session create params.
-        enable_caching (Optional[bool]): Enable caching functionality.
         browserbase_session_id (Optional[str]): Session ID for resuming Browserbase sessions.
         model_name (Optional[str]): Name of the model to use.
+        model_api_key (Optional[str]): Model API key.
+        logger (Optional[Callable[[Any], None]]): Custom logging function.
+        verbose (Optional[int]): Verbosity level for logs (1=minimal, 2=medium, 3=detailed).
+        use_rich_logging (bool): Whether to use Rich for colorized logging.
+        dom_settle_timeout_ms (Optional[int]): Timeout for DOM to settle (in milliseconds).
+        enable_caching (Optional[bool]): Enable caching functionality.
         self_heal (Optional[bool]): Enable self-healing functionality.
         wait_for_captcha_solves (Optional[bool]): Whether to wait for CAPTCHA to be solved.
         act_timeout_ms (Optional[int]): Timeout for act commands (in milliseconds).
+        headless (bool): Run browser in headless mode
         system_prompt (Optional[str]): System prompt to use for LLM interactions.
-        verbose (Optional[int]): Verbosity level for logs (1=minimal, 2=medium, 3=detailed).
         local_browser_launch_options (Optional[dict[str, Any]]): Local browser launch options.
     """
 
-    env: str = "BROWSERBASE"
+    env: Literal["BROWSERBASE", "LOCAL"] = "BROWSERBASE"
     api_key: Optional[str] = Field(
         None, alias="apiKey", description="Browserbase API key for authentication"
     )
     project_id: Optional[str] = Field(
         None, alias="projectId", description="Browserbase project ID"
     )
+    api_url: Optional[str] = Field(
+        None, alias="apiUrl", description="Stagehand API URL"
+    )  # might add a default value here
+    model_api_key: Optional[str] = Field(
+        None, alias="modelApiKey", description="Model API key"
+    )
     verbose: Optional[int] = Field(
         1,
-        description="Verbosity level for logs: 1=minimal (INFO), 2=medium (WARNING), 3=detailed (DEBUG)",
+        description="Verbosity level for logs: 0=minimal (ERROR), 1=medium (INFO), 2=detailed (DEBUG)",
     )
     logger: Optional[Callable[[Any], None]] = Field(
         None, description="Custom logging function"
     )
+    use_rich_logging: Optional[bool] = Field(
+        True, description="Whether to use Rich for colorized logging"
+    )
     dom_settle_timeout_ms: Optional[int] = Field(
         3000,
         alias="domSettleTimeoutMs",

stagehand/client.py stagehand/main.pystagehand/client.py renamed to stagehand/main.py

@@ -4,7 +4,7 @@
 import sys
 import time
 from pathlib import Path
-from typing import Any, Literal, Optional
+from typing import Any, Optional
 
 import httpx
 from dotenv import load_dotenv
@@ -39,62 +39,27 @@
 
 class Stagehand:
     """
-    Python client for interacting with a running Stagehand server and Browserbase remote headless browser.
-
-    Now supports automatically creating a new session if no session_id is provided.
-    You can provide a configuration via the 'config' parameter, or use individual parameters to override
-    the default configuration values.
+    Main Stagehand class.
     """
 
-    # Dictionary to store one lock per session_id
     _session_locks = {}
-
-    # Flag to track if cleanup has been called
     _cleanup_called = False
 
     def __init__(
         self,
-        config: Optional[StagehandConfig] = None,
-        *,
-        api_url: Optional[str] = None,
-        model_api_key: Optional[str] = None,
-        session_id: Optional[str] = None,
-        env: Optional[Literal["BROWSERBASE", "LOCAL"]] = None,
-        httpx_client: Optional[httpx.AsyncClient] = None,
-        timeout_settings: Optional[httpx.Timeout] = None,
-        use_rich_logging: bool = True,
+        config: StagehandConfig = default_config,
         **config_overrides,
     ):
         """
         Initialize the Stagehand client.
 
         Args:
             config (Optional[StagehandConfig]): Configuration object. If not provided, uses default_config.
-            api_url (Optional[str]): The running Stagehand server URL. Overrides config if provided.
-            model_api_key (Optional[str]): Your model API key (e.g. OpenAI, Anthropic, etc.). Overrides config if provided.
-            session_id (Optional[str]): Existing Browserbase session ID to connect to. Overrides config if provided.
-            env (Optional[Literal["BROWSERBASE", "LOCAL"]]): Environment to run in. Overrides config if provided.
-            httpx_client (Optional[httpx.AsyncClient]): Optional custom httpx.AsyncClient instance.
-            timeout_settings (Optional[httpx.Timeout]): Optional custom timeout settings for httpx.
-            use_rich_logging (bool): Whether to use Rich for colorized logging.
             **config_overrides: Additional configuration overrides to apply to the config.
         """
-        # Start with provided config or default config
-        if config is None:
-            config = default_config
 
         # Apply any overrides
         overrides = {}
-        if api_url is not None:
-            # api_url isn't in config, handle separately
-            pass
-        if model_api_key is not None:
-            # model_api_key isn't in config, handle separately
-            pass
-        if session_id is not None:
-            overrides["browserbase_session_id"] = session_id
-        if env is not None:
-            overrides["env"] = env
 
         # Add any additional config overrides
         overrides.update(config_overrides)
@@ -106,8 +71,9 @@ def __init__(
             self.config = config
 
         # Handle non-config parameters
-        self.api_url = api_url or os.getenv("STAGEHAND_API_URL")
-        self.model_api_key = model_api_key or os.getenv("MODEL_API_KEY")
+        self.api_url = self.config.api_url or os.getenv("STAGEHAND_API_URL")
+        self.model_api_key = self.config.model_api_key or os.getenv("MODEL_API_KEY")
+        self.model_name = self.config.model_name
 
         # Extract frequently used values from config for convenience
         self.browserbase_api_key = self.config.api_key or os.getenv(
@@ -117,7 +83,6 @@ def __init__(
             "BROWSERBASE_PROJECT_ID"
         )
         self.session_id = self.config.browserbase_session_id
-        self.model_name = self.config.model_name
         self.dom_settle_timeout_ms = self.config.dom_settle_timeout_ms
         self.self_heal = self.config.self_heal
         self.wait_for_captcha_solves = self.config.wait_for_captcha_solves
@@ -161,8 +126,7 @@ def __init__(
         # Handle streaming response setting
         self.streamed_response = True
 
-        self.httpx_client = httpx_client
-        self.timeout_settings = timeout_settings or httpx.Timeout(
+        self.timeout_settings = httpx.Timeout(
             connect=180.0,
             read=180.0,
             write=180.0,
@@ -184,7 +148,9 @@ def __init__(
         # Initialize the centralized logger with the specified verbosity
         self.on_log = self.config.logger or default_log_handler
         self.logger = StagehandLogger(
-            verbose=self.verbose, external_logger=self.on_log, use_rich=use_rich_logging
+            verbose=self.verbose,
+            external_logger=self.on_log,
+            use_rich=self.config.use_rich_logging,
         )
 
         # If using BROWSERBASE, session_id or creation params are needed
@@ -456,9 +422,7 @@ async def init(self):
 
         if self.env == "BROWSERBASE":
             if not self._client:
-                self._client = self.httpx_client or httpx.AsyncClient(
-                    timeout=self.timeout_settings
-                )
+                self._client = httpx.AsyncClient(timeout=self.timeout_settings)
 
             # Create session if we don't have one
             if not self.session_id:
@@ -570,8 +534,7 @@ async def close(self):
                     "Cannot end server session: HTTP client not available."
                 )
 
-            # Close internal HTTPX client if it was created by Stagehand
-            if self._client and not self.httpx_client:
+            if self._client:
                 self.logger.debug("Closing the internal HTTPX client...")
                 await self._client.aclose()
                 self._client = None