From 662b754e4e930b5fac5af16ae2f9e0347ba9b193 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 14 Oct 2026 15:37:22 +0000 Subject: [PATCH] test: feature - unit tests for agent_utils, most_basic, debug, visual and visualwebarena agents Co-Authored-By: Daniella Grimberg --- tests/agents/test_agent_utils.py | 148 ++++++ tests/agents/test_debug_agent.py | 162 +++++++ tests/agents/test_most_basic_agent.py | 166 +++++++ tests/agents/test_visual_agent.py | 443 ++++++++++++++++++ .../agents/test_visualwebarena_agent_unit.py | 322 +++++++++++++ 5 files changed, 1241 insertions(+) create mode 100644 tests/agents/test_agent_utils.py create mode 100644 tests/agents/test_debug_agent.py create mode 100644 tests/agents/test_most_basic_agent.py create mode 100644 tests/agents/test_visual_agent.py create mode 100644 tests/agents/test_visualwebarena_agent_unit.py diff --git a/tests/agents/test_agent_utils.py b/tests/agents/test_agent_utils.py new file mode 100644 index 000000000..0671f0846 --- /dev/null +++ b/tests/agents/test_agent_utils.py @@ -0,0 +1,148 @@ +from math import cos, radians, sin +from unittest.mock import MagicMock + +import numpy as np +import pytest +from PIL import Image, ImageDraw + +from agentlab.agents.agent_utils import ( + draw_arrowhead, + draw_click_indicator, + draw_mouse_pointer, + zoom_webpage, +) + + +def _white(size=(100, 100), mode="RGB"): + color = (255, 255, 255) if mode == "RGB" else (255, 255, 255, 255) + return Image.new(mode, size, color) + + +def _changed_mask(before: Image.Image, after: Image.Image) -> np.ndarray: + a = np.asarray(before.convert("RGBA")).astype(int) + b = np.asarray(after.convert("RGBA")).astype(int) + return np.any(a != b, axis=-1) + + +class TestDrawMousePointer: + @pytest.mark.parametrize("mode", ["RGB", "RGBA"]) + def test_returns_new_rgba_image_of_same_size(self, mode): + img = _white(mode=mode) + out = draw_mouse_pointer(img, 10, 10) + assert out is not img + assert out.size == img.size + assert out.mode == "RGBA" + + def test_does_not_mutate_input(self): + img = _white() + before = np.asarray(img).copy() + draw_mouse_pointer(img, 30, 30) + np.testing.assert_array_equal(np.asarray(img), before) + + def test_pixels_change_only_inside_pointer_box(self): + x, y = 30, 40 + img = _white() + out = draw_mouse_pointer(img, x, y) + mask = _changed_mask(img, out) + ys, xs = np.nonzero(mask) + assert mask.sum() > 0 + # pointer polygon spans [x, x+20] x [y, y+20] + assert xs.min() >= x and xs.max() <= x + 20 + assert ys.min() >= y and ys.max() <= y + 20 + + def test_pointer_is_semi_transparent_black(self): + x, y = 30, 40 + out = draw_mouse_pointer(_white(), x, y) + # a point well inside the arrow polygon + r, g, b, a = out.getpixel((x + 12, y + 10)) + assert a == 255 + assert r == g == b + assert 100 <= r <= 155 # ~50% blend of black over white + + def test_far_away_pixels_untouched(self): + out = draw_mouse_pointer(_white(), 5, 5) + assert out.getpixel((90, 90)) == (255, 255, 255, 255) + + +class TestDrawClickIndicator: + def test_returns_new_rgba_image_of_same_size(self): + img = _white((80, 60)) + out = draw_click_indicator(img, 40, 30) + assert out is not img + assert out.size == (80, 60) + assert out.mode == "RGBA" + + def test_center_gap_is_left_untouched(self): + img = Image.new("RGB", (100, 100), (100, 150, 200)) + out = draw_click_indicator(img, 50, 50) + assert out.getpixel((50, 50)) == (100, 150, 200, 255) + + @pytest.mark.parametrize("dx,dy", [(0, -9), (0, 9), (-9, 0), (9, 0)]) + def test_four_arms_are_black(self, dx, dy): + img = Image.new("RGB", (100, 100), (100, 150, 200)) + out = draw_click_indicator(img, 50, 50) + assert out.getpixel((50 + dx, 50 + dy)) == (0, 0, 0, 255) + + def test_changes_confined_to_cross_extent(self): + img = Image.new("RGB", (100, 100), (100, 150, 200)) + out = draw_click_indicator(img, 50, 50) + ys, xs = np.nonzero(_changed_mask(img, out)) + reach = 4 + 10 + 2 # gap + line_length + outline slack + assert xs.min() >= 50 - reach and xs.max() <= 50 + reach + assert ys.min() >= 50 - reach and ys.max() <= 50 + reach + + def test_does_not_mutate_input(self): + img = _white() + before = np.asarray(img).copy() + draw_click_indicator(img, 50, 50) + np.testing.assert_array_equal(np.asarray(img), before) + + +class TestDrawArrowhead: + def test_draws_two_red_lines_from_end_with_expected_geometry(self): + draw = MagicMock() + draw_arrowhead(draw, (0, 0), (100, 0), arrow_length=15, arrow_angle=30) + assert draw.line.call_count == 2 + (left_pts,), left_kw = draw.line.call_args_list[0] + (right_pts,), right_kw = draw.line.call_args_list[1] + assert left_pts[0] == (100, 0) and right_pts[0] == (100, 0) + exp_x = 100 - 15 * cos(radians(30)) + assert left_pts[1] == pytest.approx((exp_x, 15 * sin(radians(30)))) + assert right_pts[1] == pytest.approx((exp_x, -15 * sin(radians(30)))) + assert left_kw == {"fill": "red", "width": 4} + assert right_kw == {"fill": "red", "width": 4} + + def test_vertical_arrow_head_points_back_toward_start(self): + draw = MagicMock() + draw_arrowhead(draw, (50, 0), (50, 80)) + for call in draw.line.call_args_list: + (_, tip) = call.args[0] + assert tip[1] < 80 # wings trail behind the arrow tip + + def test_draws_red_pixels_on_real_image(self): + img = _white() + draw_arrowhead(ImageDraw.Draw(img), (10, 50), (80, 50)) + arr = np.asarray(img) + red = (arr[..., 0] == 255) & (arr[..., 1] == 0) & (arr[..., 2] == 0) + assert red.any() + assert not red[:, :60].any() # head is near the end point only + + +class TestZoomWebpage: + def test_default_zoom(self): + page = MagicMock() + assert zoom_webpage(page) is page + page.evaluate.assert_called_once_with("document.documentElement.style.zoom='150.0%'") + + @pytest.mark.parametrize("factor,expected", [(0.5, "50.0%"), (2, "200%"), (1.25, "125.0%")]) + def test_custom_zoom(self, factor, expected): + page = MagicMock() + zoom_webpage(page, factor) + page.evaluate.assert_called_once_with(f"document.documentElement.style.zoom='{expected}'") + + @pytest.mark.parametrize("factor", [0, -1, -0.5]) + def test_non_positive_zoom_raises(self, factor): + page = MagicMock() + with pytest.raises(ValueError, match="greater than 0"): + zoom_webpage(page, factor) + page.evaluate.assert_not_called() diff --git a/tests/agents/test_debug_agent.py b/tests/agents/test_debug_agent.py new file mode 100644 index 000000000..22fe1775a --- /dev/null +++ b/tests/agents/test_debug_agent.py @@ -0,0 +1,162 @@ +from dataclasses import dataclass +from unittest.mock import MagicMock + +import bgym +import numpy as np +import pytest +from bgym import HighLevelActionSetArgs + +from agentlab.agents import debug_agent as debug_mod +from agentlab.agents.agent_args import AgentArgs +from agentlab.agents.debug_agent import DEBUG_AGENT, DebugAgent, DebugAgentArgs + + +def _benchmark(name, action_set_args=None): + bench = MagicMock() + bench.name = name + bench.high_level_action_set_args = action_set_args or HighLevelActionSetArgs(subsets=["bid"]) + return bench + + +# ---------------------------------------------------------------- AgentArgs base + + +@dataclass +class _PlainArgs(AgentArgs): + foo: int = 1 + + def make_agent(self): + return None + + +def test_agent_args_default_name_and_set_benchmark_is_noop(): + args = _PlainArgs() + assert args.agent_name == "_PlainArgs" + before = dict(vars(args)) + assert args.set_benchmark(_benchmark("webarena"), demo_mode=True) is None + assert vars(args) == before + + +def test_agent_args_set_reproducibility_mode_not_implemented(): + with pytest.raises(NotImplementedError, match="_PlainArgs"): + _PlainArgs().set_reproducibility_mode() + + +# ---------------------------------------------------------------- DebugAgentArgs + + +def test_post_init_defaults(): + args = DebugAgentArgs() + assert args.agent_name == "debug" + assert args.use_html is False + expected = bgym.DEFAULT_BENCHMARKS["miniwob_tiny_test"]().high_level_action_set_args + assert args.action_set_args == expected + + +def test_module_level_instance(): + assert isinstance(DEBUG_AGENT, DebugAgentArgs) + assert DEBUG_AGENT.agent_name == "debug" + + +@pytest.mark.parametrize( + "name,expected_html", + [("miniwob", True), ("miniwob_tiny_test", True), ("webarena", False), ("workarena_l1", False)], +) +def test_set_benchmark(name, expected_html): + args = DebugAgentArgs() + asa = HighLevelActionSetArgs(subsets=["bid", "nav"]) + args.set_benchmark(_benchmark(name, asa), demo_mode=False) + assert args.use_html is expected_html + assert args.action_set_args is asa + + +def test_set_reproducibility_mode_not_supported(): + with pytest.raises(NotImplementedError, match="DebugAgentArgs"): + DebugAgentArgs().set_reproducibility_mode() + + +@pytest.mark.parametrize("use_html", [True, False]) +def test_make_agent(use_html): + args = DebugAgentArgs() + args.action_set_args = HighLevelActionSetArgs(subsets=["bid"]) + args.use_html = use_html + agent = args.make_agent() + assert isinstance(agent, DebugAgent) + assert agent.use_html is use_html + assert isinstance(agent.action_set, bgym.HighLevelActionSet) + assert "click" in agent.action_set.describe() + + +# ---------------------------------------------------------------- DebugAgent + + +@pytest.fixture +def agent(): + return DebugAgent(HighLevelActionSetArgs(subsets=["bid"])) + + +def test_obs_preprocessor_adds_text_views_without_mutating_input(agent, monkeypatch): + dom = MagicMock(return_value="") + ax = MagicMock(return_value="[ax]") + prune = MagicMock(return_value="") + som = MagicMock(return_value="SOM") + monkeypatch.setattr(debug_mod, "flatten_dom_to_str", dom) + monkeypatch.setattr(debug_mod, "flatten_axtree_to_str", ax) + monkeypatch.setattr(debug_mod, "prune_html", prune) + monkeypatch.setattr(debug_mod, "overlay_som", som) + + screenshot = np.zeros((4, 4, 3), np.uint8) + obs = { + "dom_object": {"d": 1}, + "axtree_object": {"a": 1}, + "extra_element_properties": {"x": {}}, + "screenshot": screenshot, + } + out = agent.obs_preprocessor(obs) + + assert out["dom_txt"] == "" + assert out["axtree_txt"] == "[ax]" + assert out["pruned_html"] == "" + assert out["screenshot_som"] == "SOM" + assert set(obs) == {"dom_object", "axtree_object", "extra_element_properties", "screenshot"} + + expected_kwargs = dict( + extra_properties={"x": {}}, + with_visible=True, + with_clickable=True, + with_center_coords=True, + with_bounding_box_coords=True, + filter_visible_only=False, + filter_with_bid_only=False, + filter_som_only=False, + ) + assert dom.call_args.args == ({"d": 1},) + assert dom.call_args.kwargs == expected_kwargs + assert ax.call_args.args == ({"a": 1},) + assert ax.call_args.kwargs == expected_kwargs + prune.assert_called_once_with("") + assert som.call_args.kwargs == {"extra_properties": {"x": {}}} + # input was deep-copied before processing + assert dom.call_args.args[0] is not obs["dom_object"] + + +@pytest.mark.parametrize( + "use_html,shown", [(True, ""), (False, "RootWebArea 'axtree'")] +) +def test_get_action_reads_from_stdin(monkeypatch, capsys, use_html, shown): + agent = DebugAgent(HighLevelActionSetArgs(subsets=["bid"]), use_html=use_html) + prompts = [] + + def fake_input(prompt): + prompts.append(prompt) + return "click('a1')" + + monkeypatch.setattr("builtins.input", fake_input) + obs = {"pruned_html": "", "axtree_txt": "RootWebArea 'axtree'"} + action, info = agent.get_action(obs) + + assert action == "click('a1')" + assert prompts == [shown + "\n"] + assert info.think == "nope" + assert info.chat_messages == [] + assert info.stats == {} diff --git a/tests/agents/test_most_basic_agent.py b/tests/agents/test_most_basic_agent.py new file mode 100644 index 000000000..96c1b9d90 --- /dev/null +++ b/tests/agents/test_most_basic_agent.py @@ -0,0 +1,166 @@ +from dataclasses import dataclass, field +from unittest.mock import MagicMock + +import bgym +import pytest + +from agentlab.agents.most_basic_agent import most_basic_agent as mba +from agentlab.agents.most_basic_agent.most_basic_agent import ( + MostBasicAgent, + MostBasicAgentArgs, +) +from agentlab.experiments.loop import ExpArgs +from agentlab.llm.base_api import BaseModelArgs +from agentlab.llm.llm_utils import AIMessage, ParseError + + +class FakeChat: + def __init__(self, responses): + self.responses = list(responses) + self.calls = [] + + def __call__(self, messages): + self.calls.append([dict(m) for m in messages]) + return AIMessage(self.responses.pop(0)) + + +@dataclass +class FakeModelArgs(BaseModelArgs): + model_name: str = "fake/model" + responses: list = field(default_factory=list) + + def make_model(self): + self.chat = FakeChat(self.responses) + return self.chat + + +OBS = {"goal": "Click the submit button", "pruned_html": ""} +GOOD = "I should click submit.\n```\nclick('a12')\n```" + + +def _content_text(message) -> str: + content = message["content"] + if isinstance(content, str): + return content + return "".join(part.get("text", "") for part in content) + + +# ---------------------------------------------------------------- args + + +def test_args_defaults(): + args = MostBasicAgentArgs() + assert args.agent_name == "BasicAgent" + assert args.temperature == 0.1 + assert args.use_chain_of_thought is False + assert args.chat_model_args is None + + +def test_make_agent_forwards_fields(): + cma = FakeModelArgs(responses=[GOOD]) + args = MostBasicAgentArgs(temperature=0.7, use_chain_of_thought=True, chat_model_args=cma) + agent = args.make_agent() + assert isinstance(agent, MostBasicAgent) + assert agent.temperature == 0.7 + assert agent.use_chain_of_thought is True + assert agent.chat_model_args is cma + assert agent.chat is cma.chat + assert isinstance(agent.action_set, bgym.HighLevelActionSet) + + +def test_set_reproducibility_mode_zeroes_temperature(): + args = MostBasicAgentArgs(temperature=0.9) + args.set_reproducibility_mode() + assert args.temperature == 0 + + +def test_prepare_and_close_delegate_to_model_args(): + cma = MagicMock() + cma.prepare_server.return_value = "prepared" + cma.close_server.return_value = "closed" + args = MostBasicAgentArgs(chat_model_args=cma) + assert args.prepare() == "prepared" + assert args.close() == "closed" + cma.prepare_server.assert_called_once_with() + cma.close_server.assert_called_once_with() + + +# ---------------------------------------------------------------- get_action + + +def test_get_action_parses_first_code_block(): + cma = FakeModelArgs(responses=["thinking\n```\nclick('a12')\n```\n```\nnoop()\n```"]) + agent = MostBasicAgentArgs(chat_model_args=cma).make_agent() + action, info = agent.get_action(OBS) + + assert action == "click('a12')" + assert info.think.startswith("thinking") + assert "noop()" in info.think + + +def test_get_action_prompt_contents_and_info(): + cma = FakeModelArgs(responses=[GOOD]) + agent = MostBasicAgentArgs(chat_model_args=cma).make_agent() + action, info = agent.get_action(OBS) + + assert action == "click('a12')" + assert info.think == GOOD + sent = cma.chat.calls[0] + assert [m["role"] for m in sent] == ["system", "user"] + assert _content_text(sent[0]) == "You are a web assistant." + human = _content_text(sent[1]) + assert OBS["goal"] in human + assert OBS["pruned_html"] in human + assert "click('a314')" in human # bid example + assert "chain of thoughts" not in human + + # system + user + assistant answer + assert info.stats["prompt_length"] == 3 + assert info.stats["response_length"] == len(GOOD) + # cost tracker stats are merged in + assert {"input_tokens", "output_tokens", "cost"} <= set(info.stats) + assert info.extra_info["chat_model_args"]["model_name"] == "fake/model" + + +def test_chain_of_thought_adds_instruction(): + cma = FakeModelArgs(responses=[GOOD]) + agent = MostBasicAgentArgs(chat_model_args=cma, use_chain_of_thought=True).make_agent() + agent.get_action(OBS) + human = _content_text(cma.chat.calls[0][1]) + assert "chain of thoughts reasoning" in human + + +def test_get_action_retries_until_code_block(): + cma = FakeModelArgs(responses=["no code here", GOOD]) + agent = MostBasicAgentArgs(chat_model_args=cma).make_agent() + action, info = agent.get_action(OBS) + + assert action == "click('a12')" + assert len(cma.chat.calls) == 2 + retry_msgs = cma.chat.calls[1] + assert retry_msgs[-1]["role"] == "user" + assert "No code block found" in _content_text(retry_msgs[-1]) + assert info.stats["prompt_length"] == 5 + + +def test_get_action_raises_after_three_failures(): + cma = FakeModelArgs(responses=["nope"] * 5) + agent = MostBasicAgentArgs(chat_model_args=cma).make_agent() + with pytest.raises(ParseError, match="after 3 retries"): + agent.get_action(OBS) + assert len(cma.chat.calls) == 3 + + +# ---------------------------------------------------------------- module config + + +def test_experiment_config_and_module_constants(): + exps = mba.experiment_config() + assert len(exps) == 2 + assert all(isinstance(e, ExpArgs) for e in exps) + assert [e.agent_args.use_chain_of_thought for e in exps] == [True, False] + assert {e.env_args.task_name for e in exps} == {"miniwob.click-button"} + assert mba.env_args.max_steps == 10 + assert mba.AGENT_4o_MINI.temperature == 0.3 + assert mba.AGENT_4o_MINI.use_chain_of_thought is True + assert mba.AGENT_4o_MINI.chat_model_args.model_name == "gpt-4o-mini-2024-07-18" diff --git a/tests/agents/test_visual_agent.py b/tests/agents/test_visual_agent.py new file mode 100644 index 000000000..d501fc22e --- /dev/null +++ b/tests/agents/test_visual_agent.py @@ -0,0 +1,443 @@ +import logging +from copy import deepcopy +from dataclasses import dataclass, field +from unittest.mock import MagicMock + +import numpy as np +import pytest +from bgym import HighLevelActionSetArgs + +from agentlab.agents import dynamic_prompting as dp +from agentlab.agents.generic_agent import tmlr_config +from agentlab.agents.generic_agent.generic_agent import GenericAgentArgs +from agentlab.agents.visual_agent import agent_configs +from agentlab.agents.visual_agent import visual_agent_prompts as vap +from agentlab.agents.visual_agent.visual_agent import VisualAgent, VisualAgentArgs +from agentlab.agents.visual_agent.visual_agent_prompts import ( + History, + MainPrompt, + Observation, + PromptFlags, + SystemPrompt, + make_instructions, +) +from agentlab.llm.base_api import BaseModelArgs +from agentlab.llm.llm_utils import AIMessage, ParseError + +VALID = "\nI will click the button.\n\n\nmouse_click(10, 20)\n" + + +class FakeChat: + def __init__(self, responses): + self.responses = list(responses) + self.calls = 0 + + def __call__(self, messages): + self.calls += 1 + return AIMessage(self.responses.pop(0)) + + def get_stats(self): + return {"n_calls": self.calls} + + +@dataclass +class FakeModelArgs(BaseModelArgs): + model_name: str = "fake/vision-model" + responses: list = field(default_factory=list) + + def make_model(self): + self.chat = FakeChat(self.responses) + return self.chat + + +def _flags(**overrides) -> PromptFlags: + flags = deepcopy(agent_configs.DEFAULT_PROMPT_FLAGS) + for k, v in overrides.items(): + setattr(flags, k, v) + return flags + + +def _obs(**overrides): + obs = { + "goal_object": [{"type": "text", "text": "Buy the red shoes"}], + "chat_messages": [], + "last_action_error": "", + "open_pages_urls": ["http://a.test", "http://b.test"], + "open_pages_titles": ["Page A", "Page B"], + "active_page_index": 1, + "screenshot": np.zeros((8, 8, 3), np.uint8), + "screenshot_som": np.full((8, 8, 3), 255, np.uint8), + } + obs.update(overrides) + return obs + + +def _text(prompt) -> str: + if isinstance(prompt, str): + return prompt + if isinstance(prompt, list): + return "".join(p.get("text", "") for p in prompt) + content = prompt["content"] + if isinstance(content, str): + return content + return "".join(p.get("text", "") for p in content if p["type"] == "text") + + +def _images(message): + return [p for p in message["content"] if p["type"] == "image_url"] + + +# ---------------------------------------------------------------- prompt pieces + + +def test_prompt_flags_defaults(): + flags = PromptFlags() + assert flags.obs is None and flags.action is None + assert flags.use_thinking is True + assert flags.use_concrete_example is False + assert flags.use_abstract_example is True + assert flags.enable_chat is False + assert flags.extra_instructions is None + + +def test_system_prompt(): + text = SystemPrompt().prompt + assert text.startswith("You are an agent trying to solve a web task") + assert "send messages to the user" in text + + +class TestMakeInstructions: + def test_goal_mode(self): + instr = make_instructions(_obs(), from_chat=False, extra_instructions=None) + assert isinstance(instr, dp.GoalInstructions) + text = _text(instr.prompt) + assert "## Goal:" in text and "Buy the red shoes" in text + assert "Extra instructions" not in text + + def test_goal_mode_with_extra_instructions(self): + instr = make_instructions(_obs(), from_chat=False, extra_instructions="Be quick.") + text = _text(instr.prompt) + assert "## Extra instructions:" in text and "Be quick." in text + + def test_goal_mode_without_chat_messages_key(self, caplog): + obs = _obs() + del obs["chat_messages"] + with caplog.at_level(logging.WARNING): + instr = make_instructions(obs, from_chat=False, extra_instructions=None) + assert isinstance(instr, dp.GoalInstructions) + assert "goal mode" not in caplog.text + + @pytest.mark.parametrize("n_user,warns", [(0, False), (1, False), (2, True), (3, True)]) + def test_warns_on_multiple_user_messages_in_goal_mode(self, caplog, n_user, warns): + msgs = [{"role": "user", "message": "hi", "timestamp": 0}] * n_user + msgs += [{"role": "assistant", "message": "yo", "timestamp": 0}] * 2 + with caplog.at_level(logging.WARNING): + make_instructions(_obs(chat_messages=msgs), from_chat=False, extra_instructions=None) + assert ("enable_chat=True" in caplog.text) is warns + + def test_chat_mode(self): + msgs = [ + {"role": "user", "message": "Find my order", "timestamp": 0}, + {"role": "assistant", "message": "Sure", "timestamp": 1}, + ] + instr = make_instructions(_obs(chat_messages=msgs), True, extra_instructions="Polite.") + assert isinstance(instr, dp.ChatInstructions) + text = instr.prompt + assert "## Chat messages:" in text + assert "[user]" in text and "Find my order" in text + assert "[assistant]" in text and "Sure" in text + assert "Polite." in text + + +class TestHistory: + def test_empty_history(self): + assert History([], []).prompt == "\n" + + def test_steps_are_numbered(self): + text = History(["click('1')", "fill('2', 'x')"], ["t0", "t1"]).prompt + assert "## Step 0\n### Thoughts:\nt0\n### Action:\nclick('1')" in text + assert "## Step 1\n### Thoughts:\nt1\n### Action:\nfill('2', 'x')" in text + assert text.index("Step 0") < text.index("Step 1") + + def test_mismatched_lengths_are_truncated(self): + text = History(["a0", "a1", "a2"], ["t0"]).prompt + assert "Step 0" in text + assert "Step 1" not in text + + +class TestObservation: + def _flags(self, **kw): + flags = deepcopy(agent_configs.DEFAULT_OBS_FLAGS) + for k, v in kw.items(): + setattr(flags, k, v) + return flags + + def test_tabs_shown_when_enabled(self): + text = Observation(_obs(), self._flags(use_tabs=True)).prompt + assert "# Observation of current step:" in text + assert "## Currently open tabs:" in text + assert "Tab 1 (active tab):" in text and "Title: Page B" in text + assert "Tab 0 (active tab)" not in text + + def test_tabs_hidden_when_disabled(self): + text = Observation(_obs(), self._flags(use_tabs=False)).prompt + assert "Currently open tabs" not in text + + @pytest.mark.parametrize( + "error,use_logs,shown", + [("Element not found", True, True), ("Element not found", False, False), ("", True, False)], + ) + def test_error_visibility(self, error, use_logs, shown): + obs = _obs(last_action_error=error) + text = Observation(obs, self._flags(use_error_logs=use_logs, use_tabs=False)).prompt + assert ("Element not found" in text) is shown + + def test_no_screenshot_when_disabled(self): + msg = vap.HumanMessage("hi") + out = Observation(_obs(), self._flags(use_screenshot=False)).add_screenshot(msg) + assert out is msg + assert _text(out) == "hi" + + @pytest.mark.parametrize( + "use_som,key,caption", + [ + (False, "screenshot", "Here is a screenshot of the page:"), + (True, "screenshot_som", "annotated with bounding boxes"), + ], + ) + def test_screenshot_selection(self, monkeypatch, use_som, key, caption): + seen = [] + + def fake_url(img): + seen.append(img) + return "data:image/jpeg;base64,AAA" + + monkeypatch.setattr(vap, "image_to_jpg_base64_url", fake_url) + obs = _obs() + flags = self._flags(use_screenshot=True, use_som=use_som, openai_vision_detail="high") + out = Observation(obs, flags).add_screenshot(vap.HumanMessage("hi")) + assert seen[0] is obs[key] + assert caption in _text(out) + (img,) = _images(out) + assert img["image_url"]["url"] == "data:image/jpeg;base64,AAA" + assert img["image_url"]["detail"] == "high" + + +class TestMainPrompt: + def _prompt(self, flags=None, actions=(), thoughts=(), obs=None): + flags = flags or _flags() + action_set = flags.action.action_set.make_action_set() + return MainPrompt(action_set, obs or _obs(), list(actions), list(thoughts), flags) + + def test_prompt_structure_with_defaults(self): + mp = self._prompt(actions=["mouse_click(1, 2)"], thoughts=["prev thought"]) + msg = mp.prompt + text = _text(msg) + assert "Buy the red shoes" in text + assert "# Observation of current step:" in text + assert "## Step 0" in text and "prev thought" in text + assert "# Action space:" in text and "mouse_click" in text + assert "" in text + assert "# Abstract Example" in text + assert "# Concrete Example" not in text + assert len(_images(msg)) == 1 # DEFAULT flags use the screenshot + + def test_concrete_example_and_no_thinking(self): + flags = _flags(use_concrete_example=True, use_abstract_example=False, use_thinking=False) + text = _text(self._prompt(flags).prompt) + assert "# Concrete Example" in text + assert "# Abstract Example" not in text + assert "" not in text + + def test_chat_mode_uses_chat_instructions(self): + msgs = [{"role": "user", "message": "Hello agent", "timestamp": 0}] + mp = self._prompt(_flags(enable_chat=True), obs=_obs(chat_messages=msgs)) + assert isinstance(mp.instructions, dp.ChatInstructions) + assert "Hello agent" in _text(mp.prompt) + + def test_parse_answer_valid(self): + ans = self._prompt()._parse_answer(VALID) + assert ans["think"] == "I will click the button." + assert ans["action"] == "mouse_click(10, 20)" + + def test_parse_answer_missing_think_is_tolerated(self): + ans = self._prompt()._parse_answer("mouse_click(1, 1)") + assert ans["action"] == "mouse_click(1, 1)" + assert "parse_error" in ans + + @pytest.mark.parametrize( + "answer", + ["hmm", "click('a1')", "not_python("], + ) + def test_parse_answer_invalid_action_raises(self, answer): + with pytest.raises(ParseError): + self._prompt()._parse_answer(answer) + + +# ---------------------------------------------------------------- VisualAgentArgs + + +def test_args_post_init_builds_agent_name(): + args = VisualAgentArgs(chat_model_args=FakeModelArgs(), flags=_flags()) + assert args.agent_name == "VisualAgent-fake_vision-model" + + +def test_args_post_init_tolerates_missing_model_args(): + args = VisualAgentArgs() + assert args.agent_name is None + assert args.max_retry == 4 + + +@pytest.mark.parametrize("multi_tab", [True, False]) +def test_args_set_benchmark_sets_tabs(multi_tab): + args = VisualAgentArgs(chat_model_args=FakeModelArgs(), flags=_flags()) + bench = MagicMock(is_multi_tab=multi_tab) + args.set_benchmark(bench, demo_mode=False) + assert args.flags.obs.use_tabs is multi_tab + + +def test_args_reproducibility_prepare_close(): + cma = MagicMock() + cma.model_name = "m" + cma.prepare_server.return_value = "p" + cma.close_server.return_value = "c" + args = VisualAgentArgs(chat_model_args=cma, flags=_flags()) + args.set_reproducibility_mode() + assert cma.temperature == 0 + assert args.prepare() == "p" + assert args.close() == "c" + + +def test_args_make_agent(): + cma = FakeModelArgs(responses=[VALID]) + args = VisualAgentArgs(chat_model_args=cma, flags=_flags(), max_retry=2) + agent = args.make_agent() + assert isinstance(agent, VisualAgent) + assert agent.max_retry == 2 + assert agent.chat_llm is cma.chat + assert agent.flags is args.flags + assert agent.actions == [] and agent.thoughts == [] and agent.seed is None + + +# ---------------------------------------------------------------- VisualAgent + + +def _agent(responses, max_retry=4, flags=None): + cma = FakeModelArgs(responses=responses) + return VisualAgent(cma, flags or _flags(), max_retry=max_retry), cma + + +def test_obs_preprocessor_delegates(monkeypatch): + seen = {} + + def fake_make(obs_flags): + seen["flags"] = obs_flags + return lambda obs: {**obs, "processed": True} + + monkeypatch.setattr(dp, "make_obs_preprocessor", fake_make) + flags = _flags() + agent, _ = _agent([], flags=flags) + assert seen["flags"] is flags.obs + assert agent.obs_preprocessor({"a": 1}) == {"a": 1, "processed": True} + + +def test_get_action_success(): + agent, cma = _agent([VALID]) + action, info = agent.get_action(_obs()) + + assert action == "mouse_click(10, 20)" + assert info.think == "I will click the button." + assert info.stats["n_retry"] == 0 + assert info.stats["busted_retry"] == 0 + assert info.stats["n_calls"] == 1 + assert {"input_tokens", "output_tokens", "cost"} <= set(info.stats) + assert [m["role"] for m in info.chat_messages] == ["system", "user", "assistant"] + assert info.extra_info["chat_model_args"]["model_name"] == "fake/vision-model" + assert agent.actions == ["mouse_click(10, 20)"] + assert agent.thoughts == ["I will click the button."] + + +def test_get_action_history_flows_into_next_prompt(): + second = VALID.replace("mouse_click(10, 20)", "mouse_click(5, 5)") + agent, _ = _agent([VALID, second]) + agent.get_action(_obs()) + action, info = agent.get_action(_obs()) + assert action == "mouse_click(5, 5)" + user_text = _text(info.chat_messages[1]) + assert "## Step 0" in user_text and "mouse_click(10, 20)" in user_text + assert agent.actions == ["mouse_click(10, 20)", "mouse_click(5, 5)"] + + +def test_get_action_counts_retries(): + agent, cma = _agent(["garbage", "bogus(", VALID]) + action, info = agent.get_action(_obs()) + assert action == "mouse_click(10, 20)" + assert info.stats["n_retry"] == 2 + assert info.stats["busted_retry"] == 0 + assert cma.chat.calls == 3 + + +def test_get_action_busted_retry(): + agent, cma = _agent(["garbage"] * 3, max_retry=3) + action, info = agent.get_action(_obs()) + assert action is None + assert info.think is None + assert info.stats["n_retry"] == 4 + assert info.stats["busted_retry"] == 1 + assert cma.chat.calls == 3 + assert agent.actions == [None] + assert agent.thoughts == [None] + + +def test_reset_clears_state(): + agent, _ = _agent([VALID]) + agent.get_action(_obs()) + agent.reset(seed=7) + assert agent.seed == 7 + assert agent.actions == [] and agent.thoughts == [] + + +# ---------------------------------------------------------------- configs + + +def test_visual_agent_configs(): + assert agent_configs.VISUAL_AGENT_4o.agent_name == "VisualAgent-gpt-4o-2024-05-13" + assert agent_configs.VISUAL_AGENT_CLAUDE_3_5.agent_name.startswith("VisualAgent-") + assert "/" not in agent_configs.VISUAL_AGENT_CLAUDE_3_5.agent_name + flags = agent_configs.DEFAULT_PROMPT_FLAGS + assert flags.obs.use_screenshot is True and flags.obs.use_som is False + assert isinstance(flags.action.action_set, HighLevelActionSetArgs) + assert tuple(flags.action.action_set.subsets) == ("coord",) + + +TMLR_LLM = "openai/gpt-4o-mini-2024-07-18" + + +def test_tmlr_base_agent(): + args = tmlr_config.get_base_agent(TMLR_LLM) + assert isinstance(args, GenericAgentArgs) + assert args.flags is tmlr_config.BASE_FLAGS + assert args.agent_name == "GenericAgent-gpt-4o-mini-2024-07-18" + assert args.flags.obs.use_screenshot is False + + +def test_tmlr_vision_agent_does_not_mutate_base_flags(): + args = tmlr_config.get_vision_agent(TMLR_LLM) + assert args.agent_name == "GenericAgent-gpt-4o-mini-2024-07-18_vision" + assert args.flags.obs.use_screenshot is True + assert args.flags.obs.use_som is False + assert tmlr_config.BASE_FLAGS.obs.use_screenshot is False + + +def test_tmlr_som_agent_does_not_mutate_base_flags(): + args = tmlr_config.get_som_agent(TMLR_LLM) + assert args.agent_name == "GenericAgent-gpt-4o-mini-2024-07-18" + assert args.flags.obs.use_screenshot is True + assert args.flags.obs.use_som is True + assert tmlr_config.BASE_FLAGS.obs.use_som is False + assert tmlr_config.BASE_FLAGS.obs.use_screenshot is False + + +def test_tmlr_unknown_llm_raises(): + with pytest.raises(KeyError): + tmlr_config.get_base_agent("not/a-model") diff --git a/tests/agents/test_visualwebarena_agent_unit.py b/tests/agents/test_visualwebarena_agent_unit.py new file mode 100644 index 000000000..01478506b --- /dev/null +++ b/tests/agents/test_visualwebarena_agent_unit.py @@ -0,0 +1,322 @@ +import base64 +import io +from dataclasses import dataclass, field +from unittest.mock import MagicMock + +import numpy as np +import PIL.Image +import pytest +from bgym import HighLevelActionSetArgs + +from agentlab.agents.visualwebarena import agent as vwa +from agentlab.agents.visualwebarena.agent import ( + VisualWebArenaAgent, + VisualWebArenaAgentArgs, + WebArenaAgent, + image_data_to_uri, + parser, +) +from agentlab.agents.visualwebarena.prompts import TEMPLATES +from agentlab.llm.base_api import BaseModelArgs +from agentlab.llm.llm_utils import ParseError + +GOOD = "Let's think step-by-step. ```click('a7')```" + + +class FakeChat: + def __init__(self, responses): + self.responses = list(responses) + self.calls = [] + + def __call__(self, messages): + self.calls.append([dict(m) for m in messages]) + return {"role": "assistant", "content": self.responses.pop(0)} + + +@dataclass +class FakeModelArgs(BaseModelArgs): + model_name: str = "org/fake-model" + responses: list = field(default_factory=list) + + def make_model(self): + self.chat = FakeChat(self.responses) + return self.chat + + +def _decode(uri: str, fmt: str) -> PIL.Image.Image: + prefix = f"data:image/{fmt};base64," + assert uri.startswith(prefix) + return PIL.Image.open(io.BytesIO(base64.b64decode(uri[len(prefix) :]))) + + +def _png_bytes(mode="RGB", size=(5, 3)): + buf = io.BytesIO() + PIL.Image.new(mode, size).save(buf, format="PNG") + return buf.getvalue() + + +def _action_set(): + return HighLevelActionSetArgs(subsets=["bid"]).make_action_set() + + +def _agent(observation_type="axtree", few_shot=False, responses=()): + chat = FakeChat(responses) + agent = VisualWebArenaAgent( + temperature=0.1, + chat_model=chat, + action_set=_action_set(), + observation_type=observation_type, + with_few_shot_examples=few_shot, + ) + return agent, chat + + +def _obs(**overrides): + obs = { + "active_page_index": np.array([1]), + "open_pages_titles": ["Home", "Shop"], + "axtree_txt": "[a7] button 'Buy'", + "url": "http://shop.test/item", + "screenshot": np.zeros((6, 6, 3), np.uint8), + "extra_element_properties": {}, + "goal_object": [{"type": "text", "text": "Buy the item"}], + } + obs.update(overrides) + return obs + + +def _texts(content): + return [c["text"] for c in content if c["type"] == "text"] + + +def _images(content): + return [c for c in content if c["type"] == "image_url"] + + +# ---------------------------------------------------------------- image_data_to_uri + + +@pytest.mark.parametrize("fmt", ["png", "jpeg"]) +def test_image_data_to_uri_from_ndarray(fmt): + arr = np.zeros((4, 7, 3), np.uint8) + img = _decode(image_data_to_uri(arr, fmt), fmt) + assert img.size == (7, 4) + assert img.format == fmt.upper() + + +def test_image_data_to_uri_from_bytes_defaults_to_png(): + img = _decode(image_data_to_uri(_png_bytes(size=(5, 3))), "png") + assert img.size == (5, 3) + + +@pytest.mark.parametrize("mode", ["RGBA", "LA"]) +def test_image_data_to_uri_drops_alpha(mode): + img = _decode(image_data_to_uri(_png_bytes(mode=mode), "png"), "png") + assert img.mode == "RGB" + + +def test_image_data_to_uri_rgba_array_to_jpeg(): + arr = np.zeros((3, 3, 4), np.uint8) + img = _decode(image_data_to_uri(arr, "jpeg"), "jpeg") + assert img.mode == "RGB" + + +def test_image_data_to_uri_rejects_unknown_format(): + with pytest.raises(AssertionError): + image_data_to_uri(np.zeros((2, 2, 3), np.uint8), "gif") + + +# ---------------------------------------------------------------- parser + + +def test_parser_returns_first_block_and_full_response(): + resp = "thinking\n```\nclick('a7')\n```\nthen\n```\nnoop()\n```" + assert parser(resp) == {"action": "click('a7')", "think": resp} + + +def test_parser_inline_backticks(): + assert parser(GOOD)["action"] == "click('a7')" + + +def test_parser_without_code_raises(): + with pytest.raises(ParseError, match="No code block"): + parser("I would click the button") + + +# ---------------------------------------------------------------- args + + +def test_args_agent_name_format(): + args = VisualWebArenaAgentArgs(chat_model_args=FakeModelArgs()) + assert args.agent_name == "VisualWebArenaAgent-axtree_som-org_fake-model" + assert args.with_few_shot_examples is True + + +def test_webarena_partial_uses_axtree(): + args = WebArenaAgent(chat_model_args=FakeModelArgs()) + assert args.observation_type == "axtree" + assert args.agent_name == "WebArenaAgent-axtree-org_fake-model" + + +def test_module_level_configs(): + assert vwa.WA_AGENT_4O_MINI.agent_name.startswith("WebArenaAgent-axtree-") + assert vwa.VWA_AGENT_4O_MINI.observation_type == "axtree_som" + for args in (vwa.WA_AGENT_SONNET, vwa.VWA_AGENT_SONNET, vwa.VWA_AGENT_4O): + assert "/" not in args.agent_name + + +def test_args_set_benchmark_deepcopies_action_set(): + asa = HighLevelActionSetArgs(subsets=["bid", "nav"]) + bench = MagicMock(high_level_action_set_args=asa) + args = VisualWebArenaAgentArgs(chat_model_args=FakeModelArgs()) + args.set_benchmark(bench, demo_mode=False) + assert args.action_set_args == asa + assert args.action_set_args is not asa + + +def test_args_reproducibility_prepare_close(): + cma = MagicMock() + cma.model_name = "x" + cma.prepare_server.return_value = "p" + cma.close_server.return_value = "c" + args = VisualWebArenaAgentArgs(chat_model_args=cma, temperature=0.8) + args.set_reproducibility_mode() + assert args.temperature == 0.0 + assert args.prepare() == "p" + assert args.close() == "c" + + +def test_args_make_agent(): + cma = FakeModelArgs() + args = VisualWebArenaAgentArgs( + chat_model_args=cma, + action_set_args=HighLevelActionSetArgs(subsets=["bid"]), + observation_type="axtree", + with_few_shot_examples=False, + temperature=0.3, + ) + agent = args.make_agent() + assert isinstance(agent, VisualWebArenaAgent) + assert agent.chat_model is cma.chat + assert agent.temperature == 0.3 + assert agent.observation_type == "axtree" + assert agent.few_shot_messages == [] + + +# ---------------------------------------------------------------- agent init + + +@pytest.mark.parametrize("obs_type", ["axtree", "axtree_screenshot", "axtree_som"]) +def test_intro_message_contains_action_space(obs_type): + agent, _ = _agent(obs_type) + (intro,) = agent.intro_messages + assert intro["type"] == "text" + assert "click(bid" in intro["text"] + assert "{action_space_description}" not in intro["text"] + assert agent.action_history == ["None"] + assert agent.few_shot_messages == [] + + +def test_text_only_few_shot_examples(): + agent, _ = _agent("axtree", few_shot=True) + n = len(TEMPLATES["axtree"]["examples"]) + assert len(agent.few_shot_messages) == n + assert agent.few_shot_messages[0]["text"].startswith(f"Example 1/{n}:") + assert "ACTION: " in agent.few_shot_messages[0]["text"] + assert _images(agent.few_shot_messages) == [] + + +@pytest.mark.parametrize("obs_type", ["axtree_screenshot", "axtree_som"]) +def test_multimodal_few_shot_examples_embed_screenshots(obs_type): + agent, _ = _agent(obs_type, few_shot=True) + examples = TEMPLATES[obs_type]["examples"] + n_img = sum(len(e) == 3 for e in examples) + n_txt = sum(len(e) == 2 for e in examples) + assert n_img > 0 + imgs = _images(agent.few_shot_messages) + assert len(imgs) == n_img + assert all(i["image_url"]["url"].startswith("data:image/png;base64,") for i in imgs) + assert len(agent.few_shot_messages) == 4 * n_img + n_txt + + +def test_bad_few_shot_example_format_raises(monkeypatch): + monkeypatch.setitem( + TEMPLATES, "axtree", {"intro": "{action_space_description}", "examples": [("only-one",)]} + ) + with pytest.raises(ValueError, match="Unexpected format"): + _agent("axtree", few_shot=True) + + +# ---------------------------------------------------------------- get_action + + +def test_get_action_axtree_builds_messages(): + agent, chat = _agent("axtree", responses=[GOOD]) + action, info = agent.get_action(_obs()) + + assert action == "click('a7')" + assert info.think == GOOD + system, user, answer = info.chat_messages + assert answer["content"] == GOOD + assert system["role"] == "system" and system["content"] == agent.intro_messages + assert user["role"] == "user" + texts = _texts(user["content"]) + obs_text = texts[0] + assert obs_text.startswith("OBSERVATION:") + assert "Tab 0: Home | Tab 1 (current): Shop" in obs_text + assert "[a7] button 'Buy'" in obs_text + assert "URL: http://shop.test/item" in obs_text + assert "PREVIOUS ACTION: None" in obs_text + assert texts[-2].startswith("OBJECTIVE:") + assert user["content"][-1] == {"type": "text", "text": "Buy the item"} + assert _images(user["content"]) == [] + assert agent.action_history == ["None", "click('a7')"] + assert {"input_tokens", "output_tokens", "cost"} <= set(info.stats) + + +def test_get_action_tracks_previous_action(): + agent, chat = _agent("axtree", responses=[GOOD, "```scroll(0, 100)```"]) + agent.get_action(_obs()) + action, info = agent.get_action(_obs()) + assert action == "scroll(0, 100)" + assert "PREVIOUS ACTION: click('a7')" in _texts(info.chat_messages[1]["content"])[0] + + +def test_get_action_screenshot_without_som(monkeypatch): + som = MagicMock() + monkeypatch.setattr(vwa, "overlay_som", som) + agent, _ = _agent("axtree_screenshot", responses=[GOOD]) + _, info = agent.get_action(_obs()) + content = info.chat_messages[1]["content"] + assert len(_images(content)) == 1 + assert "SCREENSHOT:\n" in _texts(content) + som.assert_not_called() + + +def test_get_action_screenshot_with_som(monkeypatch): + marked = np.full((6, 6, 3), 255, np.uint8) + som = MagicMock(return_value=marked) + monkeypatch.setattr(vwa, "overlay_som", som) + agent, _ = _agent("axtree_som", responses=[GOOD]) + obs = _obs(extra_element_properties={"a7": {"som": True}}) + _, info = agent.get_action(obs) + som.assert_called_once_with(obs["screenshot"], {"a7": {"som": True}}) + (img,) = _images(info.chat_messages[1]["content"]) + assert img["image_url"]["url"] == image_data_to_uri(marked) + + +def test_get_action_includes_few_shots_before_observation(): + agent, _ = _agent("axtree", few_shot=True, responses=[GOOD]) + _, info = agent.get_action(_obs()) + content = info.chat_messages[1]["content"] + n = len(agent.few_shot_messages) + assert content[:n] == agent.few_shot_messages + assert content[n]["text"].startswith("OBSERVATION:") + + +def test_get_action_retries_then_raises(): + agent, chat = _agent("axtree", responses=["no code"] * 3) + with pytest.raises(ParseError): + agent.get_action(_obs()) + assert len(chat.calls) == 3 + assert agent.action_history == ["None"]