working gaia bench and gym, with test

Author: ollmer

src/agentlab/benchmarks/abstract_env.py

@@ -1,9 +1,23 @@
 from abc import ABC, abstractmethod
 
-import gym
+import gymnasium as gym
 from pydantic import BaseModel
 
 
+class AbstractBenchmark(BaseModel):
+    name: str
+    env_args_list: list = None
+
+    def get_version(self) -> int:
+        return "1"
+
+    def prepare_backends(self):
+        pass
+
+    def dependency_graph_over_tasks(self) -> dict[str, list[str]]:
+        return {}
+
+
 class AbstractEnvArgs(BaseModel):
     """Easily serialiazable class to store the arguments of an environment"""
 

src/agentlab/benchmarks/gaia.py

@@ -2,7 +2,6 @@
 import shutil
 from typing import Any, Literal
 
-import bgym
 import datasets
 from pydantic import Field
 from tapeagents.core import Observation, StopStep, Thought
@@ -12,24 +11,16 @@
 from tapeagents.tools.media_reader import VideoReader
 from tapeagents.tools.web_search import WebSearch
 
-from agentlab.benchmarks.abstract_env import AbstractEnvArgs
+from agentlab.benchmarks.abstract_env import AbstractBenchmark, AbstractEnvArgs
 from agentlab.benchmarks.multitool_gym import MultiToolGym
 
 
-class GaiaBenchmark(bgym.Benchmark):
-    name = "gaia"
-    split: Literal["test", "validation"]
+class GaiaBenchmark(AbstractBenchmark):
     exp_dir: str
+    name: str = "gaia"
+    split: Literal["test", "validation"]
 
-    high_level_action_set_args = None
-    is_multi_tab = False
-    supports_parallel_seeds = False
-    backends = ["gaia"]
-    env_args_list = None
-    task_metadata = None
-
-    def __post_init__(self):
-        super().__post_init__()
+    def model_post_init(self, __context: Any) -> None:
         self.env_args_list = []
         dataset = datasets.load_dataset("gaia-benchmark/GAIA", "2023_all")[self.split]
         for task in dataset:
@@ -45,7 +36,6 @@ class GaiaGym(MultiToolGym):
 
 class GaiaGymArgs(AbstractEnvArgs):
     task: dict[str, Any]
-    split: Literal["test", "validation"]
     exp_dir: str
     viewport_chars: int = 64000
 

src/agentlab/benchmarks/multitool_gym.py

@@ -1,58 +1,12 @@
-from typing import Any, Literal, Union
-
-from pydantic import Annotated, Field, TypeAdapter
 from tapeagents.core import Action, Observation, Tape
-from tapeagents.environment import ToolCollectionEnvironment
-from tapeagents.tools.base import Multitool, Tool
 
 from agentlab.benchmarks.abstract_env import AbstractEnv
 
 EnvTape = Tape[None, Action | Observation]
 
 
-class FunctionCall(Action):
-    kind: Literal["function_call_action"] = ["function_call_action"]
-    function_name: str
-    args: list[Any] | None
-    kwargs: dict[str, Any] | None
-
-
-class FunctionCallResult(Observation):
-    kind: Literal["function_call_result"] = ["function_call_result"]
-    result: Any
-
-
-class SimpleFunctionCallTool(Tool):
-    action = FunctionCall
-    observation = FunctionCallResult
-    function: callable
-    function_name: str = ""
-
-    def model_post_init(self, __context):
-        function_name = getattr(self.function, "__name__", "")
-        if not function_name and not self.function_name:
-            raise ValueError("Function has no name, function_name must be provided")
-
-    def execute_action(self, action: FunctionCall) -> FunctionCallResult:
-        if not self.function_name == action.function_name:
-            raise ValueError(
-                f"Unexpected function action {action.function_name}, expected {self.function_name}"
-            )
-        result = self.function(*action.args, **action.kwargs)
-        return FunctionCallResult(result=result)
-
-
 class MultiToolGym(AbstractEnv):
-    def __init__(self, tools: list[Tool | Multitool]):
-        self._env = ToolCollectionEnvironment(tools)
-        self._actions = self._env.actions()
-        self._actions_parser: TypeAdapter = TypeAdapter(
-            Annotated[Union[self._actions], Field(discriminator="kind")]
-        )
-        self.reset()
-
     def reset(self):
-        self._tape: EnvTape = EnvTape(steps=[])
         self._env.reset()
 
     def step(self, action: str):
@@ -61,14 +15,12 @@ def step(self, action: str):
         except Exception:
             raise ValueError("Action must be a valid JSON dict")
         assert isinstance(action_step, Action), "{action_step.kind} is not an Action"
-        self._tape += [action_step]
-        self._tape = self._env.react(self._tape)
-        observation_step: Observation = self._tape.steps[-1]
+        observation = self._env.step(action_step)
         reward = self.calculate_reward()
         terminated = False
         truncated = False
-        env_info = {"step_metadata": observation_step.metadata}
-        return observation_step.llm_dict(), reward, terminated, truncated, env_info
+        env_info = {"step_metadata": observation.metadata}
+        return observation, reward, terminated, truncated, env_info
 
     def calculate_reward(self) -> float:
         return 0.0

tests/agents/test_gaia_agent.py

@@ -1,8 +1,34 @@
 from agentlab.agents.tapeagent import TapeAgent, TapeAgentArgs
+from agentlab.benchmarks.gaia import GaiaBenchmark
 
 
 def test_agent_creation():
     args = TapeAgentArgs(agent_name="gaia_agent")
     agent = args.make_agent()
     assert isinstance(agent, TapeAgent)
     assert agent.agent.name == "gaia_agent"
+
+
+def test_gaia_bench():
+    exp_dir = "/tmp/"
+    bench = GaiaBenchmark(exp_dir=exp_dir, split="validation")
+    assert bench.name == "gaia"
+    assert bench.split == "validation"
+    assert bench.exp_dir == exp_dir
+    assert len(bench.env_args_list) == 165
+
+    assert bench.env_args_list[5].exp_dir == "gaia/32102e3e-d12a-4209-9163-7b3a104efe5d"
+    assert bench.env_args_list[5].viewport_chars == 64000
+    task = bench.env_args_list[5].task
+    question = """The attached spreadsheet shows the inventory for a movie and video game rental store in Seattle, Washington. What is the title of the oldest Blu-Ray recorded in this spreadsheet? Return it as appearing in the spreadsheet."""
+    steps = """1. Open the attached file.\n2. Compare the years given in the Blu-Ray section to find the oldest year, 2009.\n3. Find the title of the Blu-Ray disc that corresponds to the year 2009: Time-Parking 2: Parallel Universe."""
+    assert task["task_id"] == "32102e3e-d12a-4209-9163-7b3a104efe5d"
+    assert task["Question"] == question
+    assert task["Level"] == "2"
+    assert task["Final answer"] == "Time-Parking 2: Parallel Universe"
+    assert task["file_name"] == "32102e3e-d12a-4209-9163-7b3a104efe5d.xlsx"
+    assert task["Annotator Metadata"]["Steps"] == steps
+    assert task["Annotator Metadata"]["Number of steps"] == "3"
+    assert task["Annotator Metadata"]["How long did this take?"] == "1 minute"
+    assert task["Annotator Metadata"]["Tools"] == "1. Microsoft Excel"
+    assert task["Annotator Metadata"]["Number of tools"] == "1"