separate logging from utils

Author: miguelg719

stagehand/__init__.py

@@ -2,6 +2,7 @@
 from .client import Stagehand
 from .config import StagehandConfig, default_config
 from .handlers.observe_handler import ObserveHandler
+from .logging import configure_logging
 from .metrics import StagehandFunctionName, StagehandMetrics
 from .page import StagehandPage
 from .schemas import (
@@ -17,7 +18,6 @@
     ObserveOptions,
     ObserveResult,
 )
-from .utils import configure_logging
 
 __version__ = "0.0.1"
 

stagehand/a11y/utils.py

@@ -7,13 +7,14 @@
 if TYPE_CHECKING:
     from stagehand.page import StagehandPage
 
+from ..logging import StagehandLogger
 from ..types.a11y import (
     AccessibilityNode,
     AXNode,
     CDPSession,
     TreeResult,
 )
-from ..utils import StagehandLogger, format_simplified_tree
+from ..utils import format_simplified_tree
 
 
 async def _clean_structural_nodes(

stagehand/client.py

@@ -2,9 +2,9 @@
 import json
 import os
 import shutil
-import tempfile
 import signal
 import sys
+import tempfile
 import time
 from pathlib import Path
 from typing import Any, Literal, Optional
@@ -23,13 +23,12 @@
 from .config import StagehandConfig, default_config
 from .context import StagehandContext
 from .llm import LLMClient
+from .logging import StagehandLogger, default_log_handler
 from .metrics import StagehandFunctionName, StagehandMetrics
 from .page import StagehandPage
 from .schemas import AgentConfig
 from .utils import (
-    StagehandLogger,
     convert_dict_keys_to_camel_case,
-    default_log_handler,
     make_serializable,
 )
 
@@ -47,7 +46,7 @@ class Stagehand:
 
     # Dictionary to store one lock per session_id
     _session_locks = {}
-    
+
     # Flag to track if cleanup has been called
     _cleanup_called = False
 
@@ -193,7 +192,7 @@ def __init__(
                     raise ValueError(
                         "browserbase_project_id is required for BROWSERBASE env with existing session_id (or set BROWSERBASE_PROJECT_ID in env)."
                     )
-            
+
         # Register signal handlers for graceful shutdown
         self._register_signal_handlers()
 
@@ -224,13 +223,16 @@ def __init__(
 
     def _register_signal_handlers(self):
         """Register signal handlers for SIGINT and SIGTERM to ensure proper cleanup."""
+
         def cleanup_handler(sig, frame):
             # Prevent multiple cleanup calls
             if self.__class__._cleanup_called:
                 return
 
             self.__class__._cleanup_called = True
-            print(f"\n[{signal.Signals(sig).name}] received. Ending Browserbase session...")
+            print(
+                f"\n[{signal.Signals(sig).name}] received. Ending Browserbase session..."
+            )
 
             try:
                 # Try to get the current event loop
@@ -252,11 +254,11 @@ def cleanup_handler(sig, frame):
                 def schedule_cleanup():
                     task = asyncio.create_task(self._async_cleanup())
                     # Shield the task to prevent it from being cancelled
-                    shielded = asyncio.shield(task)
+                    asyncio.shield(task)
                     # We don't need to await here since we're in call_soon_threadsafe
-                
+
                 loop.call_soon_threadsafe(schedule_cleanup)
-                
+
             except Exception as e:
                 print(f"Error during signal cleanup: {str(e)}")
                 sys.exit(1)