Cover VisualWebArena agent step/parse path with a non-pricy test

tests/agents/test_visualwebarena_agent.py

@@ -1,13 +1,104 @@
 import logging
+import re
 import tempfile
+from dataclasses import dataclass
+from pathlib import Path
 
+import bgym
 import pytest
 
 from agentlab.agents.visualwebarena.agent import VisualWebArenaAgentArgs
+from agentlab.analyze import inspect_results
+from agentlab.experiments import launch_exp
 from agentlab.experiments.loop import EnvArgs, ExpArgs
 from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT
 
 
+def _messages_to_text(messages) -> str:
+    """Flatten the multi-modal messages of the VisualWebArena agent into plain text."""
+    chunks = []
+    for message in messages:
+        content = message.get("content", "")
+        if isinstance(content, str):
+            chunks.append(content)
+        else:
+            chunks.extend(part["text"] for part in content if part.get("type") == "text")
+    return "\n".join(chunks)
+
+
+@dataclass
+class CheatVisualWebArenaLLM:
+    """For unit-testing purposes only. It only works with the miniwob.click-test task."""
+
+    def __call__(self, messages) -> dict:
+        prompt = _messages_to_text(messages)
+        # only look at the observation, few-shot examples also contain buttons
+        prompt = prompt.rsplit("OBSERVATION:", 1)[-1]
+        match = re.search(r"^\s*\[(\d+)\].*button", prompt, re.MULTILINE | re.IGNORECASE)
+        if match is None:
+            raise Exception("Can't find the button's bid")
+
+        action = f'click("{match.group(1)}")'
+        answer = f"""\
+Let's think step-by-step. The objective is to click the button. In summary, the next action I \
+will perform is ```{action}```"""
+        return dict(role="assistant", content=answer)
+
+    def get_stats(self):
+        return {}
+
+
+@dataclass
+class CheatVisualWebArenaLLMArgs:
+    model_name: str = "test/cheat_visualwebarena_click_test"
+    max_total_tokens: int = 10240
+    max_input_tokens: int = 8000
+    max_new_tokens: int = 128
+
+    def make_model(self):
+        return CheatVisualWebArenaLLM()
+
+    def prepare_server(self):
+        pass
+
+    def close_server(self):
+        pass
+
+
+@pytest.mark.parametrize("observation_type", ["axtree", "axtree_som"])
+def test_visualwebarena_agent_on_miniwob(observation_type):
+    agent_args = VisualWebArenaAgentArgs(
+        temperature=0.0,
+        chat_model_args=CheatVisualWebArenaLLMArgs(),
+        observation_type=observation_type,
+    )
+    agent_args.set_benchmark(bgym.DEFAULT_BENCHMARKS["miniwob_tiny_test"](), demo_mode=False)
+
+    exp_args = ExpArgs(
+        agent_args=agent_args,
+        env_args=EnvArgs(task_name="miniwob.click-test", task_seed=42, max_steps=5, headless=True),
+    )
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        launch_exp.run_experiments(
+            1, [exp_args], Path(tmp_dir) / "visualwebarena_agent_test", parallel_backend="joblib"
+        )
+        result_record = inspect_results.load_result_df(tmp_dir, progress_fn=None)
+
+        target = {
+            "n_steps": 1,
+            "cum_reward": 1.0,
+            "terminated": True,
+            "truncated": False,
+            "err_msg": None,
+            "stack_trace": None,
+        }
+
+        for key, target_val in target.items():
+            assert key in result_record
+            assert result_record[key].iloc[0] == target_val
+
+
 @pytest.mark.pricy
 def test_agent():
     with tempfile.TemporaryDirectory() as exp_dir: