Add docstrings for function arguments and return types in agent_utils, tool_use_agent, loop, study, and llm_utils modules

Author: amanjaiswal73892

src/agentlab/agents/agent_utils.py

@@ -96,6 +96,14 @@ def draw_mouse_pointer(image: Image.Image, x: int, y: int) -> Image.Image:
     """
     Draws a semi-transparent mouse pointer at (x, y) on the image.
     Returns a new image with the pointer drawn.
+
+    Args:
+        image: The image to draw the mouse pointer on.
+        x: The x coordinate for the mouse pointer.
+        y: The y coordinate for the mouse pointer.
+
+    Returns:
+        A new image with the mouse pointer drawn.
     """
     pointer_size = 20  # Length of the pointer
     overlay = image.convert("RGBA").copy()
@@ -134,6 +142,14 @@ def draw_click_indicator(image: Image.Image, x: int, y: int) -> Image.Image:
     """
     Draws a click indicator (+ shape with disconnected lines) at (x, y) on the image.
     Returns a new image with the click indicator drawn.
+
+    Args:
+        image: The image to draw the click indicator on.
+        x: The x coordinate for the click indicator.
+        y: The y coordinate for the click indicator.
+
+    Returns:
+        A new image with the click indicator drawn.
     """
     line_length = 10  # Length of each line segment
     gap = 4  # Gap from center point
@@ -188,7 +204,13 @@ def zoom_webpage(page: Page, zoom_factor: float = 1.5):
 
     Args:
         page: The Playwright Page object.
-        zoom_factor: The zoom factor to apply (default is 1.0, which means no zoom).
+        zoom_factor: The zoom factor to apply (default is 1.5).
+
+    Returns:
+        Page: The modified Playwright Page object.
+
+    Raises:
+        ValueError: If zoom_factor is less than or equal to 0.
     """
 
     if zoom_factor <= 0:

src/agentlab/agents/tool_use_agent/tool_use_agent.py

@@ -41,7 +41,11 @@ def _init(self):
     def make(self) -> "Block":
         """Returns a copy so the init can start adding some stuff to `self` without changing the
         original datatclass that should only contain a config.
-        The aim is avoid having 2 calss definition for each block, e.g. Block and BlockArgs."""
+        The aim is avoid having 2 calss definition for each block, e.g. Block and BlockArgs.
+
+        Returns:
+            Block: A copy of the current block instance with initialization applied.
+        """
         block = self.__class__(**asdict(self))
         block._init()
         return block

src/agentlab/experiments/loop.py

@@ -55,6 +55,7 @@ def make_env(
             action_mapping: overrides the action mapping of the environment.
             exp_dir: will set some environment parameters (e.g., record_video_dir) with respect to the directory where the experiment is running.
             exp_task_kwargs: use with caution! Will override task parameters to experiment-specific values. Useful to set different server configs for different experiments, or output file paths within the experiment's folder (e.g., assistantbench).
+            use_raw_page_output: if True, the environment will also return raw page output in the observation.
 
         Returns:
             env: the gym environment.

src/agentlab/experiments/study.py

@@ -729,6 +729,15 @@ def _convert_env_args(env_args_list):
     """Return a list where every element is the *new* EnvArgs.
 
     For backward compatibility, we need to convert the old EnvArgs to the new one.
+
+    Args:
+        env_args_list (list): list of EnvArgs objects to convert
+
+    Returns:
+        list: list of converted EnvArgs objects
+
+    Raises:
+        TypeError: If an element in env_args_list is not of expected type.
     """
     from bgym import EnvArgs as BGymEnvArgs
 

src/agentlab/llm/llm_utils.py

@@ -107,6 +107,24 @@ def generic_call_api_with_retries(
     Makes an API call with retries for transient failures, rate limiting,
     and responses deemed invalid by a custom validation function.
     (Refactored for improved readability with helper functions)
+
+    Args:
+        client_function: The API client function to call.
+        api_params: Parameters to pass to the client function.
+        is_response_valid_fn: Function to validate if the response is valid.
+        rate_limit_exceptions: Tuple of exception types for rate limiting.
+        api_error_exceptions: Tuple of exception types for API errors.
+        get_status_code_fn: Optional function to extract status code from exceptions.
+        max_retries: Maximum number of retry attempts.
+        initial_retry_delay_seconds: Initial delay between retries in seconds.
+        max_retry_delay_seconds: Maximum delay between retries in seconds.
+
+    Returns:
+        The API response if successful.
+
+    Raises:
+        Exception: For unexpected errors that are immediately re-raised.
+        RuntimeError: If API call fails after maximum retries.
     """
 
     def _calculate_delay(
@@ -231,6 +249,14 @@ def call_openai_api_with_retries(client_function, api_params, max_retries=10):
     Makes an OpenAI API call with retries for transient failures,
     rate limiting, and invalid or error-containing responses.
     (This is now a wrapper around generic_call_api_with_retries for OpenAI)
+
+    Args:
+        client_function: The OpenAI API client function to call.
+        api_params: Parameters to pass to the client function.
+        max_retries: Maximum number of retry attempts.
+
+    Returns:
+        The OpenAI API response if successful.
     """
 
     def is_openai_response_valid(response):
@@ -266,6 +292,14 @@ def call_anthropic_api_with_retries(client_function, api_params, max_retries=10)
     Makes an Anthropic API call with retries for transient failures,
     rate limiting, and invalid responses.
     (This is a wrapper around generic_call_api_with_retries for Anthropic)
+
+    Args:
+        client_function: The Anthropic API client function to call.
+        api_params: Parameters to pass to the client function.
+        max_retries: Maximum number of retry attempts.
+
+    Returns:
+        The Anthropic API response if successful.
     """
 
     def is_anthropic_response_valid(response):