Merge branch 'main' into miguel/stg-470-fix-logging

Author: miguelg719

pyproject.toml

@@ -24,6 +24,10 @@ dependencies = [
     "playwright>=1.42.1",
     "requests>=2.31.0",
     "browserbase>=1.4.0",
+    "rich>=14.0.0",
+    "openai>=1.83.0",
+    "anthropic>=0.51.0",
+    "litellm>=1.72.0",
 ]
 
 [project.optional-dependencies]
@@ -36,15 +40,18 @@ dev = [
     "isort>=5.12.0",
     "mypy>=1.3.0",
     "ruff",
-    "rich",  # Useful for development/debugging
 ]
 
 [project.urls]
 Homepage = "https://github.com/browserbase/stagehand-python"
 Repository = "https://github.com/browserbase/stagehand-python"
 
 [tool.setuptools]
-packages = ["stagehand"]
+# Omit explicit package list so we can use the nested `.packages.find` table below.
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["stagehand*"]
 
 [tool.setuptools.package-data]
 stagehand = ["domScripts.js"]

stagehand/__init__.py

@@ -1,9 +1,9 @@
 """Stagehand - The AI Browser Automation Framework"""
 
 from .agent import Agent
-from .client import Stagehand
-from .config import StagehandConfig
+from .config import StagehandConfig, default_config
 from .handlers.observe_handler import ObserveHandler
+from .main import Stagehand
 from .llm import LLMClient
 from .logging import LogConfig, configure_logging
 from .metrics import StagehandFunctionName, StagehandMetrics

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 (0=errors, 1=info, 2=debug).
         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: 0=errors only, 1=info, 2=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 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
@@ -25,73 +25,41 @@
 from .config import StagehandConfig, default_config
 from .context import StagehandContext
 from .llm import LLMClient
-from .logging import LogConfig, StagehandLogger, default_log_handler
 from .metrics import StagehandFunctionName, StagehandMetrics
 from .page import StagehandPage
 from .schemas import AgentConfig
-from .utils import make_serializable
+from .utils import (
+    StagehandLogger,
+    default_log_handler,
+    make_serializable,
+)
 
 load_dotenv()
 
 
 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)
@@ -103,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(
@@ -114,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
@@ -138,8 +106,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,
@@ -158,19 +125,14 @@ def __init__(
         if self.env not in ["BROWSERBASE", "LOCAL"]:
             raise ValueError("env must be either 'BROWSERBASE' or 'LOCAL'")
 
-        # Create centralized log configuration
-        self.log_config = LogConfig(
+        # Initialize the centralized logger with the specified verbosity
+        self.on_log = self.config.logger or default_log_handler
+        self.logger = StagehandLogger(
             verbose=self.verbose,
-            use_rich=use_rich_logging,
-            env=self.env,
-            external_logger=self.config.logger or default_log_handler,
-            quiet_dependencies=True,
+            external_logger=self.on_log,
+            use_rich=self.config.use_rich_logging,
         )
 
-        # Initialize the centralized logger with the LogConfig
-        self.on_log = self.log_config.external_logger
-        self.logger = StagehandLogger(config=self.log_config)
-
         # If using BROWSERBASE, session_id or creation params are needed
         if self.env == "BROWSERBASE":
             if not self.session_id:
@@ -185,7 +147,7 @@ def __init__(
                     )
                 if not self.model_api_key:
                     # Model API key needed if Stagehand server creates the session
-                    self.logger.info(
+                    self.logger.warning(
                         "model_api_key is recommended when creating a new BROWSERBASE session to configure the Stagehand server's LLM."
                     )
             elif self.session_id:
@@ -221,7 +183,6 @@ def __init__(
         self.llm = None
         if self.env == "LOCAL":
             self.llm = LLMClient(
-                stagehand_logger=self.logger,
                 api_key=self.model_api_key,
                 default_model=self.model_name,
                 metrics_callback=self._handle_llm_metrics,
@@ -403,11 +364,13 @@ def _get_lock_for_session(self) -> asyncio.Lock:
         return self._session_locks[self.session_id]
 
     async def __aenter__(self):
+        self.logger.debug("Entering Stagehand context manager (__aenter__)...")
         # Just call init() if not already done
         await self.init()
         return self
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self.logger.debug("Exiting Stagehand context manager (__aexit__)...")
         await self.close()
 
     async def init(self):
@@ -428,18 +391,16 @@ 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:
                 await self._create_session()  # Uses self._client and api_url
-                self.logger.info(
-                    f"Created new Browserbase session. Session ID: {self.session_id}"
+                self.logger.debug(
+                    f"Created new Browserbase session via Stagehand server: {self.session_id}"
                 )
             else:
-                self.logger.info(
+                self.logger.debug(
                     f"Using existing Browserbase session: {self.session_id}"
                 )
 
@@ -538,12 +499,11 @@ async def close(self):
                         f"Error ending server session {self.session_id}: {str(e)}"
                     )
             elif self.session_id:
-                self.logger.debug(
+                self.logger.warning(
                     "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
@@ -580,16 +540,17 @@ async def _handle_log(self, msg: dict[str, Any]):
 
             # Map level strings to internal levels
             level_map = {
-                "debug": 2,
+                "debug": 3,
                 "info": 1,
+                "warning": 2,
                 "error": 0,
             }
 
             # Convert string level to int if needed
             if isinstance(level_str, str):
                 internal_level = level_map.get(level_str.lower(), 1)
             else:
-                internal_level = min(level_str, 2)  # Ensure level is between 0-2
+                internal_level = min(level_str, 3)  # Ensure level is between 0-3
 
             # Handle the case where message itself might be a JSON-like object
             if isinstance(message, dict):
@@ -623,7 +584,7 @@ def _log(
 
         Args:
             message: The message to log
-            level: Verbosity level (0=error, 1=info, 2=debug)
+            level: Verbosity level (0=error, 1=info, 2=detailed, 3=debug)
             category: Optional category for the message
             auxiliary: Optional auxiliary data to include
         """