update llmclient logging

Author: miguelg719

stagehand/__init__.py

@@ -1,11 +1,11 @@
-"""Stagehand - Browser automation with AI"""
+"""Stagehand - The AI Browser Automation Framework"""
 
 from .agent import Agent
 from .client import Stagehand
 from .config import StagehandConfig
 from .handlers.observe_handler import ObserveHandler
 from .llm import LLMClient
-from .logging import LogConfig, configure_logging, get_logger
+from .logging import LogConfig, configure_logging
 from .metrics import StagehandFunctionName, StagehandMetrics
 from .page import StagehandPage
 from .schemas import (
@@ -21,7 +21,7 @@
     ObserveResult,
 )
 
-__version__ = "0.1.0"
+__version__ = "0.0.1"
 
 __all__ = [
     "Stagehand",
@@ -44,5 +44,4 @@
     "StagehandFunctionName",
     "StagehandMetrics",
     "LogConfig",
-    "get_logger",
 ]

stagehand/client.py

@@ -222,6 +222,7 @@ 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,

stagehand/llm/client.py

@@ -1,14 +1,13 @@
 """LLM client for model interactions."""
 
-import logging
-from typing import Any, Callable, Optional
+from typing import TYPE_CHECKING, Any, Callable, Optional
 
 import litellm
 
 from stagehand.metrics import get_inference_time_ms, start_inference_timer
 
-# Configure logger for the module
-logger = logging.getLogger(__name__)
+if TYPE_CHECKING:
+    from ..logging import StagehandLogger
 
 
 class LLMClient:
@@ -19,6 +18,7 @@ class LLMClient:
 
     def __init__(
         self,
+        stagehand_logger: "StagehandLogger",
         api_key: Optional[str] = None,
         default_model: Optional[str] = None,
         metrics_callback: Optional[Callable[[Any, int, Optional[str]], None]] = None,
@@ -28,6 +28,7 @@ def __init__(
         Initializes the LiteLLMClient.
 
         Args:
+            stagehand_logger: StagehandLogger instance for centralized logging
             api_key: An API key for the default provider, if required.
                      It's often better to set provider-specific environment variables
                      (e.g., OPENAI_API_KEY, ANTHROPIC_API_KEY) which litellm reads automatically.
@@ -39,6 +40,7 @@ def __init__(
             **kwargs: Additional global settings for litellm (e.g., api_base).
                       See litellm documentation for available settings.
         """
+        self.logger = stagehand_logger
         self.default_model = default_model
         self.metrics_callback = metrics_callback
 
@@ -50,11 +52,13 @@ def __init__(
         for key, value in kwargs.items():
             if hasattr(litellm, key):
                 setattr(litellm, key, value)
-                logger.debug(f"Set global litellm.{key}")
+                self.logger.debug(f"Set global litellm.{key}", category="llm")
             # Handle common aliases or expected config names if necessary
             elif key == "api_base":  # Example: map api_base if needed
                 litellm.api_base = value
-                logger.debug(f"Set global litellm.api_base to {value}")
+                self.logger.debug(
+                    f"Set global litellm.api_base to {value}", category="llm"
+                )
 
     def create_response(
         self,
@@ -107,9 +111,11 @@ def create_response(
             k: v for k, v in params.items() if v is not None or k in kwargs
         }
 
-        logger.debug(
-            f"Calling litellm.completion with model={completion_model} and params: {filtered_params}"
+        self.logger.debug(
+            f"Calling litellm.completion with model={completion_model} and params: {filtered_params}",
+            category="llm",
         )
+
         try:
             # Start tracking inference time
             start_time = start_inference_timer()
@@ -127,6 +133,6 @@ def create_response(
             return response
 
         except Exception as e:
-            logger.error(f"Error calling litellm.completion: {e}")
+            self.logger.error(f"Error calling litellm.completion: {e}", category="llm")
             # Consider more specific exception handling based on litellm errors
             raise

stagehand/logging.py

@@ -49,15 +49,6 @@ def get_remote_verbose(self) -> int:
         """
         return self.verbose
 
-    def get_python_log_level(self) -> int:
-        """Get the Python logging level based on verbose setting."""
-        level_map = {
-            0: logging.ERROR,
-            1: logging.INFO,
-            2: logging.DEBUG,
-        }
-        return level_map.get(self.verbose, logging.INFO)
-
     def should_log(self, level: int) -> bool:
         """Check if a message at the given level should be logged."""
         # Always log errors (level 0)
@@ -214,24 +205,17 @@ def __init__(
             config: Optional LogConfig instance. If provided, overrides other parameters.
         """
         if config:
-            self.verbose = config.verbose
-            self.external_logger = config.external_logger
-            self.use_rich = config.use_rich
             self.config = config
         else:
-            self.verbose = verbose
-            self.external_logger = external_logger
-            self.use_rich = use_rich
             self.config = LogConfig(
                 verbose=verbose,
                 use_rich=use_rich,
                 external_logger=external_logger,
             )
 
-        self.console = get_console(self.use_rich)
+        self.console = get_console(self.config.use_rich)
 
         # Map our verbosity levels to Python's logging levels
-        # Now using only 3 levels to match the remote Fastify server
         self.level_map = {
             0: logging.ERROR,  # Critical errors only
             1: logging.INFO,  # Standard information
@@ -242,11 +226,26 @@ def __init__(
         self.level_style = {0: "error", 1: "info", 2: "debug"}
 
         # Update logger level based on verbosity
-        self._set_verbosity(self.verbose)
+        self._set_verbosity(self.config.verbose)
+
+    @property
+    def verbose(self):
+        """Get verbose level from config"""
+        return self.config.verbose
+
+    @property
+    def use_rich(self):
+        """Get use_rich setting from config"""
+        return self.config.use_rich
+
+    @property
+    def external_logger(self):
+        """Get external logger from config"""
+        return self.config.external_logger
 
     def _set_verbosity(self, level: int):
         """Set the logger verbosity level"""
-        self.verbose = level
+        self.config.verbose = level
         logger.setLevel(self.level_map.get(level, logging.INFO))
 
     def _format_json(self, data: dict) -> str:
@@ -445,7 +444,7 @@ def log(
             auxiliary: Optional dictionary of auxiliary data
         """
         # Skip logging if below current verbosity level
-        if level > self.verbose and level != 0:  # Always log errors (level 0)
+        if not self.config.should_log(level):
             return
 
         # Call external logger if provided (handle async function)
@@ -650,20 +649,6 @@ def debug(
         self.log(message, level=2, category=category, auxiliary=auxiliary)
 
 
-def get_logger(name: str, config: LogConfig) -> StagehandLogger:
-    """
-    Factory function to get a configured logger instance for a module.
-
-    Args:
-        name: The name of the module requesting the logger
-        config: LogConfig instance with logging configuration
-
-    Returns:
-        StagehandLogger: Configured logger instance
-    """
-    return StagehandLogger(config=config)
-
-
 # Create a synchronous wrapper for the async default_log_handler
 def sync_log_handler(log_data: dict[str, Any]) -> None:
     """