Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions stagehand/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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:
Expand Down
52 changes: 11 additions & 41 deletions stagehand/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -54,47 +54,19 @@ class Stagehand:

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)
Expand All @@ -106,8 +78,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(
Expand All @@ -117,7 +90,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
Expand All @@ -141,8 +113,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,
Expand All @@ -164,7 +135,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
Expand Down Expand Up @@ -425,9 +398,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:
Expand Down Expand Up @@ -539,8 +510,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
Expand Down
32 changes: 22 additions & 10 deletions stagehand/config.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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",
Expand Down