darglinting everything

Author: TLSDC

src/agentlab/agents/agent_args.py

@@ -1,5 +1,5 @@
-from bgym import AbstractAgentArgs
 import bgym
+from bgym import AbstractAgentArgs
 
 
 class AgentArgs(AbstractAgentArgs):
@@ -28,6 +28,9 @@ def set_reproducibility_mode(self):
         as possible e.g. setting the temperature of the model to 0.
 
         This is only called when reproducibility is requested.
+
+        Raises:
+            NotImplementedError: If the agent does not support reproducibility.
         """
         raise NotImplementedError(
             f"set_reproducibility_mode is not implemented for agent_args {self.__class__.__name__}"

src/agentlab/agents/generic_agent/agent_configs.py

@@ -193,12 +193,19 @@
     add_missparsed_messages=True,
 )
 
+
 AGENT_8B = GenericAgentArgs(
     chat_model_args=CHAT_MODEL_ARGS_DICT["meta-llama/Meta-Llama-3-8B-Instruct"],
     flags=FLAGS_8B,
 )
 
 
+AGENT_LLAMA31_8B = GenericAgentArgs(
+    chat_model_args=CHAT_MODEL_ARGS_DICT["openrouter/meta-llama/llama-3.1-8b-instruct"],
+    flags=FLAGS_8B,
+)
+
+
 # GPT-4o default config
 FLAGS_GPT_4o = GenericPromptFlags(
     obs=dp.ObsFlags(

src/agentlab/agents/generic_agent/reproducibility_agent.py

@@ -199,6 +199,8 @@ def make_repro_agent(agent_args: AgentArgs, exp_dir: Path | str):
         agent_args (AgentArgs): The original agent args.
         exp_dir (Path | str): The directory where the experiment was saved.
 
+    Returns:
+        ReproAgentArgs: The new agent args.
     """
     exp_dir = Path(exp_dir)
     assert isinstance(agent_args, GenericAgentArgs)

src/agentlab/agents/visualwebarena/agent.py

@@ -188,6 +188,12 @@ def get_action(self, obs: Any) -> tuple[str, dict]:
         Replica of VisualWebArena agent
         https://github.com/web-arena-x/visualwebarena/blob/89f5af29305c3d1e9f97ce4421462060a70c9a03/agent/prompts/prompt_constructor.py#L211
         https://github.com/web-arena-x/visualwebarena/blob/89f5af29305c3d1e9f97ce4421462060a70c9a03/agent/prompts/prompt_constructor.py#L272
+
+        Args:
+            obs (Any): Observation from the environment
+
+        Returns:
+            tuple[str, dict]: Action and AgentInfo
         """
         user_messages = []
 

src/agentlab/analyze/inspect_results.py

@@ -16,7 +16,6 @@
 from IPython.display import display
 from tqdm import tqdm
 
-
 from agentlab.analyze.error_categorization import (
     ERR_CLASS_MAP,
     is_critical_server_error,
@@ -83,9 +82,7 @@ def set_index_from_variables(
         index_black_list: List of wildard patterns to match variables that
             should be excluded from the index.
         task_key: The key to use as the first level of the index.
-        force_at_leaste_one_variable: If True, force at least one variable in the
-            index. If no variable is found, the index will be set to
-            task_key + "agent.agent_name".
+        add_agent_and_benchmark: If True, add agent.agent_name and env.benchmark
     """
     df.reset_index(inplace=True)
     constants, variables, _ = get_constants_and_variables(df)
@@ -141,6 +138,7 @@ def load_result_df(
             should be included in the index.
         index_black_list: List of wildard patterns to match variables that
             should be excluded from the index.
+        remove_args_suffix: If True, remove the _args suffix from the columns
 
     Returns:
         pd.DataFrame: The result dataframe
@@ -777,17 +775,13 @@ def _categorize_error(row):
 
 
 def _benchmark_from_task_name(task_name: str):
-    """Extract the benchmark from the task name.
-    TODO should be more robost, e.g. handle workarna.L1, workarena.L2, etc.
-    """
+    """Extract the benchmark from the task name."""
+    # TODO should be more robost, e.g. handle workarna.L1, workarena.L2, etc.
     return task_name.split(".")[0]
 
 
 def summarize_study(result_df: pd.DataFrame) -> pd.DataFrame:
-    """Create a summary of the study.
-
-    Similar to global report, but handles single agent differently.
-    """
+    """Create a summary of the study. Similar to global report, but handles single agent differently."""
 
     levels = list(range(result_df.index.nlevels))
     return result_df.groupby(level=levels[1:]).apply(summarize)

src/agentlab/experiments/exp_utils.py

@@ -95,6 +95,9 @@ def add_dependencies(exp_args_list: list[ExpArgs], task_dependencies: dict[str,
     Returns:
         list[ExpArgs]
             The modified exp_args_list with dependencies added.
+
+    Raises:
+        ValueError: If the task_dependencies are not valid.
     """
 
     if task_dependencies is None or all([len(dep) == 0 for dep in task_dependencies.values()]):

src/agentlab/experiments/graph_execution_ray.py

@@ -47,6 +47,17 @@ def poll_for_timeout(tasks: dict[str, ray.ObjectRef], timeout: float, poll_inter
 
     I tried various different methods for killing a job that hangs. so far it's
     the only one that seems to work reliably (hopefully)
+
+    Args:
+        tasks: dict[str, ray.ObjectRef]
+            Dictionary of task_id: task_ref
+        timeout: float
+            Timeout in seconds
+        poll_interval: float
+            Polling interval in seconds
+
+    Returns:
+        dict[str, Any]: Dictionary of task_id: result
     """
     task_list = list(tasks.values())
     task_ids = list(tasks.keys())

src/agentlab/experiments/launch_exp.py

@@ -4,6 +4,7 @@
 
 import bgym
 from browsergym.experiments.loop import ExpArgs, yield_all_exp_results
+
 from agentlab.experiments.exp_utils import run_exp
 
 
@@ -24,13 +25,16 @@ def run_experiments(
             Number of parallel jobs.
         exp_args_list: list[ExpArgs]
             List of ExpArgs objects.
-        exp_dir: Path
+        study_dir: Path
             Directory where the experiments will be saved.
         parallel_backend: str
             Parallel backend to use. Either "joblib", "ray" or "sequential".
             The only backend that supports webarena graph dependencies correctly is ray or sequential.
         avg_step_timeout: int
             Will raise a TimeoutError if the episode is not finished after env_args.max_steps * avg_step_timeout seconds.
+
+    Raises:
+        ValueError: If the parallel_backend is not recognized.
     """
 
     if len(exp_args_list) == 0:
@@ -110,6 +114,13 @@ def find_incomplete(study_dir: str | Path, include_errors=True):
             Find all incomplete experiments and relaunch them.
             - "incomplete_only": relaunch only the incomplete experiments.
             - "incomplete_or_error": relaunch incomplete or errors.
+
+    Returns:
+        list[ExpArgs]
+            List of ExpArgs objects to relaunch.
+
+    Raises:
+        ValueError: If the study_dir does not exist.
     """
     study_dir = Path(study_dir)
 
@@ -152,6 +163,16 @@ def _hide_completed(exp_result: bgym.ExpResult, include_errors: bool = True):
 
     This little hack, allows an elegant way to keep the task dependencies for e.g. webarena
     while skipping the tasks that are completed when relaunching.
+
+    Args:
+        exp_result: bgym.ExpResult
+            The experiment result to hide.
+        include_errors: bool
+            If True, include experiments that errored.
+
+    Returns:
+        ExpArgs
+            The ExpArgs object hidden if the experiment is completed.
     """
 
     hide = False

src/agentlab/experiments/reproducibility_util.py

@@ -58,7 +58,7 @@ def _get_git_username(repo: Repo) -> str:
     5. Environment variables (GIT_AUTHOR_NAME and GIT_COMMITTER_NAME)
 
     Args:
-        repo (git.Repo): A GitPython Repo object representing the Git repository.
+        repo (Repo): A GitPython Repo object representing the Git repository.
 
     Returns:
         str: The first non-None username found, or None if no username is found.

src/agentlab/experiments/study.py

@@ -1,4 +1,3 @@
-from abc import ABC, abstractmethod
 import gzip
 import logging
 import pickle
@@ -18,7 +17,6 @@
 from agentlab.experiments.exp_utils import RESULTS_DIR, add_dependencies
 from agentlab.experiments.launch_exp import find_incomplete, non_dummy_count, run_experiments
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -186,7 +184,15 @@ def load_exp_args_list(self):
     def set_reproducibility_info(self, strict_reproducibility=False, comment=None):
         """Gather relevant information that may affect the reproducibility of the experiment
 
-        e.g.: versions of BrowserGym, benchmark, AgentLab..."""
+        e.g.: versions of BrowserGym, benchmark, AgentLab...
+
+        Args:
+            strict_reproducibility: bool
+                If True, all modifications have to be committed before running the experiments.
+                Also, if relaunching a study, it will not be possible if the code has changed.
+            comment: str
+                Extra comment to add to the reproducibility information.
+        """
         agent_names = [a.agent_name for a in self.agent_args]
         info = repro.get_reproducibility_info(
             agent_names,
@@ -252,13 +258,14 @@ def _run(self, n_jobs=1, parallel_backend="joblib", strict_reproducibility=False
         Args:
             n_jobs: int
                 Number of parallel jobs.
-
             parallel_backend: str
                 Parallel backend to use. Either "joblib", "dask" or "sequential".
-
             strict_reproducibility: bool
                 If True, all modifications have to be committed before running the experiments.
                 Also, if relaunching a study, it will not be possible if the code has changed.
+
+        Raises:
+            ValueError: If the exp_args_list is None.
         """
 
         if self.exp_args_list is None:
@@ -276,10 +283,6 @@ def append_to_journal(self, strict_reproducibility=True):
         Args:
             strict_reproducibility: bool
                 If True, incomplete experiments will raise an error.
-
-        Raises:
-            ValueError: If the reproducibility information is not compatible
-                with the report.
         """
         _, summary_df, _ = self.get_results()
         repro.append_to_journal(
@@ -447,9 +450,16 @@ def _agents_on_benchmark(
             If True, the experiments will be run in demo mode.
         logging_level: int
             The logging level for individual jobs.
+        logging_level_stdout: int
+            The logging level for the stdout.
+        ignore_dependencies: bool
+            If True, the dependencies will be ignored and all experiments can be run in parallel.
 
     Returns:
         list[ExpArgs]: The list of experiments to run.
+
+    Raises:
+        ValueError: If multiple agents are run on a benchmark that requires manual reset.
     """
 
     if not isinstance(agents, (list, tuple)):