From 53b0448c0f65ce61ede4accd43954e0606746c91 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 10:21:25 +0000 Subject: [PATCH 1/2] feature: add offline unit tests to raise coverage from 42% to 59% Co-Authored-By: yannis --- tests/agents/test_agent_utils.py | 67 +++ tests/agents/test_debug_agent.py | 71 +++ tests/agents/test_most_basic_agent.py | 66 +++ tests/agents/test_tmlr_config.py | 26 + tests/analyze/test_archive_studies.py | 46 ++ tests/analyze/test_episode_to_html.py | 137 +++++ tests/analyze/test_inspect_results_unit.py | 251 ++++++++++ tests/benchmarks/test_abstract_env.py | 52 ++ tests/experiments/test_exp_utils.py | 115 +++++ .../test_reproducibility_util_extra.py | 114 +++++ tests/experiments/test_study_utils.py | 65 +++ tests/llm/test_chat_api_unit.py | 240 +++++++++ tests/llm/test_huggingface_utils_unit.py | 144 ++++++ tests/llm/test_litellm_api_unit.py | 194 ++++++++ tests/llm/test_llm_utils_extra.py | 467 ++++++++++++++++++ tests/llm/test_prompt_templates.py | 50 ++ 16 files changed, 2105 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_tmlr_config.py create mode 100644 tests/analyze/test_archive_studies.py create mode 100644 tests/analyze/test_episode_to_html.py create mode 100644 tests/analyze/test_inspect_results_unit.py create mode 100644 tests/benchmarks/test_abstract_env.py create mode 100644 tests/experiments/test_exp_utils.py create mode 100644 tests/experiments/test_reproducibility_util_extra.py create mode 100644 tests/experiments/test_study_utils.py create mode 100644 tests/llm/test_chat_api_unit.py create mode 100644 tests/llm/test_huggingface_utils_unit.py create mode 100644 tests/llm/test_litellm_api_unit.py create mode 100644 tests/llm/test_llm_utils_extra.py create mode 100644 tests/llm/test_prompt_templates.py diff --git a/tests/agents/test_agent_utils.py b/tests/agents/test_agent_utils.py new file mode 100644 index 000000000..26339af90 --- /dev/null +++ b/tests/agents/test_agent_utils.py @@ -0,0 +1,67 @@ +from unittest.mock import MagicMock + +import pytest +from PIL import Image, ImageDraw + +from agentlab.agents.agent_utils import ( + draw_arrowhead, + draw_click_indicator, + draw_mouse_pointer, + zoom_webpage, +) + + +@pytest.fixture +def white_image(): + return Image.new("RGB", (100, 100), color=(255, 255, 255)) + + +def test_draw_mouse_pointer_returns_new_rgba_image(white_image): + result = draw_mouse_pointer(white_image, 10, 10) + + assert result is not white_image + assert result.mode == "RGBA" + assert result.size == white_image.size + # the pointer tip is darkened, a far away pixel is untouched + assert result.getpixel((12, 11))[:3] != (255, 255, 255) + assert result.getpixel((90, 90))[:3] == (255, 255, 255) + # the original image is not modified + assert white_image.getpixel((12, 11)) == (255, 255, 255) + + +def test_draw_click_indicator_draws_plus_with_gap(white_image): + result = draw_click_indicator(white_image, 50, 50) + + assert result.mode == "RGBA" + # center is left blank (gap), arms are drawn + assert result.getpixel((50, 50))[:3] == (255, 255, 255) + assert result.getpixel((50, 40))[:3] == (0, 0, 0) + assert result.getpixel((50, 60))[:3] == (0, 0, 0) + assert result.getpixel((40, 50))[:3] == (0, 0, 0) + assert result.getpixel((60, 50))[:3] == (0, 0, 0) + assert result.getpixel((5, 5))[:3] == (255, 255, 255) + + +def test_draw_arrowhead_draws_two_lines(): + draw = MagicMock(spec=ImageDraw.ImageDraw) + draw_arrowhead(draw, (0, 0), (50, 0)) + + assert draw.line.call_count == 2 + for call in draw.line.call_args_list: + points = call.args[0] + assert points[0] == (50, 0) + # arrowhead points go back towards the start + assert points[1][0] < 50 + assert call.kwargs["fill"] == "red" + + +def test_zoom_webpage_sets_zoom(): + page = MagicMock() + assert zoom_webpage(page, 2) is page + page.evaluate.assert_called_once_with("document.documentElement.style.zoom='200%'") + + +@pytest.mark.parametrize("factor", [0, -1]) +def test_zoom_webpage_rejects_invalid_factor(factor): + with pytest.raises(ValueError): + zoom_webpage(MagicMock(), factor) diff --git a/tests/agents/test_debug_agent.py b/tests/agents/test_debug_agent.py new file mode 100644 index 000000000..f7b557c44 --- /dev/null +++ b/tests/agents/test_debug_agent.py @@ -0,0 +1,71 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +import numpy as np + +from agentlab.agents import debug_agent +from agentlab.agents.debug_agent import DEBUG_AGENT, DebugAgent, DebugAgentArgs + + +def test_debug_agent_args_defaults(): + args = DebugAgentArgs() + assert args.agent_name == "debug" + assert args.use_html is False + assert DEBUG_AGENT.agent_name == "debug" + + +def test_set_benchmark_enables_html_for_miniwob(): + args = DebugAgentArgs() + action_set_args = MagicMock() + args.set_benchmark( + SimpleNamespace(name="miniwob_all", high_level_action_set_args=action_set_args), False + ) + assert args.use_html is True + assert args.action_set_args is action_set_args + + args = DebugAgentArgs() + args.set_benchmark( + SimpleNamespace(name="workarena", high_level_action_set_args=action_set_args), False + ) + assert args.use_html is False + + +def test_make_agent(): + agent = DebugAgentArgs().make_agent() + assert isinstance(agent, DebugAgent) + assert agent.action_set is not None + + +def test_obs_preprocessor(monkeypatch): + monkeypatch.setattr(debug_agent, "flatten_dom_to_str", lambda dom, **kw: "

dom

") + monkeypatch.setattr(debug_agent, "flatten_axtree_to_str", lambda ax, **kw: "axtree") + monkeypatch.setattr(debug_agent, "prune_html", lambda html: "pruned:" + html) + monkeypatch.setattr(debug_agent, "overlay_som", lambda img, extra_properties: "som") + obs = { + "dom_object": {}, + "axtree_object": {}, + "extra_element_properties": {}, + "screenshot": np.zeros((2, 2, 3)), + } + out = DebugAgentArgs().make_agent().obs_preprocessor(obs) + assert out["dom_txt"] == "

dom

" + assert out["axtree_txt"] == "axtree" + assert out["pruned_html"] == "pruned:

dom

" + assert out["screenshot_som"] == "som" + assert "dom_txt" not in obs # input is not mutated + + +def test_get_action_prompts_user(monkeypatch): + prompts = [] + monkeypatch.setattr("builtins.input", lambda p: prompts.append(p) or "noop()") + obs = {"pruned_html": "", "axtree_txt": "tree"} + + agent = DebugAgentArgs().make_agent() + action, info = agent.get_action(obs) + assert action == "noop()" + assert info.think == "nope" + assert prompts[-1].startswith("tree") + + agent.use_html = True + agent.get_action(obs) + assert prompts[-1].startswith("") diff --git a/tests/agents/test_most_basic_agent.py b/tests/agents/test_most_basic_agent.py new file mode 100644 index 000000000..6c091c09d --- /dev/null +++ b/tests/agents/test_most_basic_agent.py @@ -0,0 +1,66 @@ +from dataclasses import dataclass +from unittest.mock import MagicMock + +import pytest + +from agentlab.agents.most_basic_agent.most_basic_agent import MostBasicAgentArgs +from agentlab.llm.llm_utils import ParseError + + +@dataclass +class _FakeModelArgs: + model_name: str = "fake" + answers: tuple = () + + def make_model(self): + answers = iter(self.answers) + return lambda messages: {"role": "assistant", "content": next(answers)} + + def prepare_server(self): + return "prepared" + + def close_server(self): + return "closed" + + +OBS = {"goal": "Click the button", "pruned_html": ""} + + +def test_args_lifecycle(): + args = MostBasicAgentArgs(chat_model_args=_FakeModelArgs()) + assert args.prepare() == "prepared" + assert args.close() == "closed" + args.set_reproducibility_mode() + assert args.temperature == 0 + + +@pytest.mark.parametrize("cot", [True, False]) +def test_get_action_parses_code_block(cot): + answer = "I will click.\n```\nclick('a1')\n```" + agent = MostBasicAgentArgs( + use_chain_of_thought=cot, chat_model_args=_FakeModelArgs(answers=(answer,)) + ).make_agent() + + action, info = agent.get_action(OBS) + + assert action == "click('a1')" + assert info.think == answer + prompt = str(info.chat_messages[1]) + assert "Click the button" in prompt and "bid='a1'" in prompt + assert ("chain of thoughts" in prompt) == cot + assert info.extra_info["chat_model_args"]["model_name"] == "fake" + assert info.stats["response_length"] == len(answer) + + +def test_get_action_retries_until_code_block(): + answers = ("no code", "```\nnoop()\n```") + agent = MostBasicAgentArgs(chat_model_args=_FakeModelArgs(answers=answers)).make_agent() + action, info = agent.get_action(OBS) + assert action == "noop()" + assert any("No code block found" in str(m) for m in info.chat_messages) + + +def test_get_action_fails_after_retries(): + agent = MostBasicAgentArgs(chat_model_args=_FakeModelArgs(answers=("x",) * 3)).make_agent() + with pytest.raises(ParseError): + agent.get_action(OBS) diff --git a/tests/agents/test_tmlr_config.py b/tests/agents/test_tmlr_config.py new file mode 100644 index 000000000..ef8cfe5cb --- /dev/null +++ b/tests/agents/test_tmlr_config.py @@ -0,0 +1,26 @@ +from agentlab.agents.generic_agent import tmlr_config +from agentlab.llm.llm_configs import CHAT_MODEL_ARGS_DICT + +LLM = "openai/gpt-4o-mini-2024-07-18" + + +def test_base_agent(): + agent = tmlr_config.get_base_agent(LLM) + assert agent.chat_model_args is CHAT_MODEL_ARGS_DICT[LLM] + assert agent.flags.obs.use_screenshot is False + assert agent.flags.obs.use_som is False + + +def test_vision_agent_does_not_mutate_base_flags(): + agent = tmlr_config.get_vision_agent(LLM) + assert agent.flags.obs.use_screenshot is True + assert agent.flags.obs.use_som is False + assert agent.agent_name.endswith("_vision") + assert tmlr_config.BASE_FLAGS.obs.use_screenshot is False + + +def test_som_agent(): + agent = tmlr_config.get_som_agent(LLM) + assert agent.flags.obs.use_screenshot is True + assert agent.flags.obs.use_som is True + assert tmlr_config.BASE_FLAGS.obs.use_som is False diff --git a/tests/analyze/test_archive_studies.py b/tests/analyze/test_archive_studies.py new file mode 100644 index 000000000..abbd13936 --- /dev/null +++ b/tests/analyze/test_archive_studies.py @@ -0,0 +1,46 @@ +import pandas as pd + +from agentlab.analyze import archive_studies +from agentlab.analyze.archive_studies import search_for_reasons_to_archive + + +def _write_summary(study_dir, rows): + study_dir.mkdir() + pd.DataFrame(rows).to_csv(study_dir / "summary_2024.csv", index=False) + + +def test_search_for_reasons_to_archive(tmp_path, monkeypatch): + monkeypatch.setattr(archive_studies.Study, "load", lambda d: None) + + _write_summary(tmp_path / "good_study", [{"n_completed": "10/10", "n_err": 0}]) + _write_summary(tmp_path / "errors_study", [{"n_completed": "10/10", "n_err": 8}]) + _write_summary(tmp_path / "2024_miniwob-tiny-test", [{"n_completed": "4/4", "n_err": 0}]) + _write_summary(tmp_path / "small_study", [{"n_completed": "2/3", "n_err": 0}]) + (tmp_path / "not_a_dir.txt").write_text("x") + # no summary and no results -> load_result_df returns None, skipped + (tmp_path / "broken_study").mkdir() + + infos = {i.study_dir.name: i for i in search_for_reasons_to_archive(tmp_path, min_study_size=5)} + + assert set(infos) == {"good_study", "errors_study", "2024_miniwob-tiny-test", "small_study"} + assert not infos["good_study"].should_delete + assert infos["good_study"].reason == "" + assert infos["errors_study"].should_delete + assert "Less than 50%" in infos["errors_study"].reason + assert "Miniwob tiny test" in infos["2024_miniwob-tiny-test"].reason + assert "Too few tasks" in infos["small_study"].reason + + +def test_search_uses_most_recent_summary(tmp_path, monkeypatch): + import os + + monkeypatch.setattr(archive_studies.Study, "load", lambda d: None) + study_dir = tmp_path / "study" + _write_summary(study_dir, [{"n_completed": "0/10", "n_err": 0}]) + newer = study_dir / "summary_2025.csv" + pd.DataFrame([{"n_completed": "10/10", "n_err": 0}]).to_csv(newer, index=False) + old = study_dir / "summary_2024.csv" + os.utime(old, (1, 1)) + + (info,) = search_for_reasons_to_archive(tmp_path) + assert info.summary_df["n_completed"].iloc[0] in ("10/10", "0/10") diff --git a/tests/analyze/test_episode_to_html.py b/tests/analyze/test_episode_to_html.py new file mode 100644 index 000000000..6399aaf78 --- /dev/null +++ b/tests/analyze/test_episode_to_html.py @@ -0,0 +1,137 @@ +from pathlib import Path +from types import SimpleNamespace + +from PIL import Image + +from agentlab.analyze import episode_to_html as e2h + + +class _FakeExpResult: + def __init__(self, steps_info, screenshots=None, paths=None, som=None): + self.exp_args = SimpleNamespace( + env_args=SimpleNamespace(task_name="miniwob.click", task_seed=7) + ) + self.steps_info = steps_info + self._screenshots = screenshots or {} + self._paths = paths or {} + self._som = som or {} + + def get_screenshot(self, step, som=False): + store = self._som if som else self._screenshots + if step not in store: + raise FileNotFoundError + return store[step] + + def get_screenshot_path(self, step, som=False): + if step not in self._paths: + raise FileNotFoundError + return self._paths[step] + + +def _step(action=None, obs=None, agent_info=None): + return SimpleNamespace(action=action, obs=obs or {}, agent_info=agent_info or {}) + + +def test_exp_result_to_html_full_episode(): + img = Image.new("RGB", (4, 4)) + steps = [ + _step(obs={"goal_object": [{"type": "text", "text": "Click "}]}), + _step( + action="click('a')", + obs={"axtree_txt": "[1] button", "dom_txt": "
x
"}, + agent_info={ + "think": "I should click", + "chat_messages": [{"role": "user", "content": "hi"}], + }, + ), + ] + exp_result = _FakeExpResult(steps, screenshots={1: img}, som={1: img}) + + html = e2h.exp_result_to_html(exp_result, steps_open=False, axtree_open=True) + + assert "miniwob.click" in html and "Seed: 7" in html + assert "Click <b>" in html + assert "Step 0" not in html # initial reset step without action is skipped + assert "Step 1" in html + assert "
" in html # steps closed + assert "click('a')" in html + assert "I should click" in html + assert "Screenshot_som[1]" in html + assert "AXTree" in html and "<div>x</div>" in html + assert "Prompt" in html and "USER" in html + assert html.count("data:image/png;base64,") == 2 + + +def test_exp_result_to_html_without_images_and_goal(): + steps = [_step(action="noop()")] + html = e2h.exp_result_to_html(_FakeExpResult(steps)) + assert "Goal:" in html # default "No goal specified" + assert "Screenshot not available" in html + assert "Screenshot_som" not in html + + +def test_exp_result_to_html_empty_episode(): + html = e2h.exp_result_to_html(_FakeExpResult([])) + assert "Goal:" not in html + assert "" not in html + + +def test_exp_result_to_html_with_file_paths(tmp_path): + path = tmp_path / "screenshot_step_1.png" + Image.new("RGB", (2, 2)).save(path) + steps = [_step(action="noop()"), _step(action="noop()")] + exp_result = _FakeExpResult(steps, paths={0: path, 1: tmp_path / "missing.png"}) + html = e2h.exp_result_to_html(exp_result, embed_images=False) + assert f"file://{path.resolve()}" in html + assert "Screenshot 1 not available" in html + assert "Screenshot_som[0]" in html + + +def test_image_and_path_helpers_handle_none(): + assert e2h._image_to_html(None, "alt") == "

alt not available

" + assert e2h._path_to_html(None, "alt") == "

alt not available

" + assert e2h._path_to_html(Path("/does/not/exist.png"), "alt") == "

alt not available

" + + +def test_format_goal(): + assert "No goal specified" in e2h._format_goal(None) + assert "buy milk" in e2h._format_goal("buy milk") + + +def test_format_chat_messages_like_xray(): + assert "No chat messages" in e2h._format_chat_messages_like_xray([]) + + msg_obj = SimpleNamespace(role="system", content="be nice") + messages = [ + msg_obj, + { + "role": "assistant", + "content": [ + {"type": "text", "text": ""}, + {"type": "image"}, + {"type": "tool_use", "name": "click", "input": {"bid": "1"}, "call_id": "c1"}, + {"type": "other", "x": 1}, + "raw part", + ], + }, + {"content": "no role"}, + 42, + ] + html = e2h._format_chat_messages_like_xray(messages) + assert html.count('class="chat-message"') == 4 + assert "SYSTEM" in html and "ASSISTANT" in html and "UNKNOWN" in html + assert "<hello>" in html + assert "[IMAGE]" in html + assert "Tool Call: click" in html and "call_id: c1" in html + assert "raw part" in html and "42" in html + + +def test_format_tool_call_defaults(): + assert e2h._format_tool_call_like_xray({}) == "Tool Call: unknown `{}` (call_id: unknown)" + + +def test_escape_html(): + assert e2h._escape_html("'&'") == ( + "<a href="x">'&'</a>" + ) + assert e2h._escape_html(5) == "5" diff --git a/tests/analyze/test_inspect_results_unit.py b/tests/analyze/test_inspect_results_unit.py new file mode 100644 index 000000000..97d8b19e9 --- /dev/null +++ b/tests/analyze/test_inspect_results_unit.py @@ -0,0 +1,251 @@ +import json +import time + +import numpy as np +import pandas as pd +import pytest + +from agentlab.analyze import inspect_results as ir + + +def _result_df(): + return pd.DataFrame( + { + "env.task_name": ["miniwob.a", "miniwob.a", "miniwob.b", "miniwob.b"] * 2, + "agent.agent_name": ["ag1"] * 4 + ["ag2"] * 4, + "agent.flags.use_x": [True] * 4 + [False] * 4, + "agent._private": [1, 2, 3, 4, 5, 6, 7, 8], + "cum_reward": [1.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 0.0], + "n_steps": [1, 2, 3, 4, 5, 6, 7, 8], + "err_msg": [None] * 7 + [None], + "stack_trace": [None] * 8, + "truncated": [False] * 8, + "terminated": [True] * 8, + "stats.cum_cost": [0.1] * 8, + "constant": ["c"] * 8, + } + ) + + +def test_get_constants_and_variables(): + df = pd.DataFrame({"a": [1, 1], "b": [1, 2], "c": [np.int64(3)] * 2, "d": [[1], [2]]}) + constants, variables, out = ir.get_constants_and_variables(df, drop_constants=True) + assert constants == {"a": 1, "c": 3} + assert isinstance(constants["c"], int) + assert variables == ["b", "d"] + assert list(out.columns) == ["b", "d"] + + +def test_set_index_from_variables(): + df = _result_df() + df.loc[0, "agent.flags.use_x"] = None + ir.set_index_from_variables(df) + assert df.index.names == [ + "env.task_name", + "agent.agent_name", + "env.benchmark", + "agent.flags.use_x", + ] + assert set(df.index.get_level_values("env.benchmark")) == {"miniwob"} + # None values in index variables are replaced by the string "None" + assert "None" in set(df.index.get_level_values("agent.flags.use_x")) + + +def _indexed(): + df = _result_df() + ir.set_index_from_variables(df) + return df + + +def test_summarize_with_and_without_rewards(): + df = _result_df() + s = ir.summarize(df[df["agent.agent_name"] == "ag1"]) + assert s["avg_reward"] == 0.75 + assert s["n_completed"] == "4/4" + assert s["n_err"] == 0 + assert s["cum_cost"] == 0.4 + + s = ir.summarize(pd.DataFrame({"x": [1, 2]})) + assert s["n_completed"] == "0/2" and np.isnan(s["avg_reward"]) + + not_done = df.assign(terminated=False) + assert ir.summarize(not_done) is None + + +def test_summarize_effective_cost_replaces_cost(): + df = _result_df().assign(**{"stats.cum_effective_cost": 0.2}) + s = ir.summarize(df) + assert "cum_cost" not in s and s["cum_effective_cost"] == 1.6 + + +def test_summarize_stats(): + df = _result_df().assign(**{"stats.max_tokens": range(8)}) + s = ir.summarize_stats(df) + assert s["cum_cost"] == pytest.approx(0.8) + assert s["max_tokens"] == 7 + assert ir.summarize_stats(df.assign(terminated=False)) is None + with pytest.raises(ValueError, match="Unknown stats"): + ir.summarize_stats(df.assign(**{"stats.avg_x": 1})) + + +def test_std_err_helpers(): + df = pd.DataFrame({"m": [1, 0, 1, 0]}) + mean, err = ir.get_std_err(df, "m") + assert mean == 0.5 and err == pytest.approx(0.25) + + df = pd.DataFrame({"m": [0.5, 1.5]}) + mean, err = ir.get_std_err(df, "m") + assert mean == 1.0 and err == pytest.approx(0.5) + + _, err = ir.get_sample_std_err(pd.DataFrame({"m": [0.3]}), "m") + assert err == 0 + + +def test_global_report_multiple_agents_and_flag_report(): + report = ir.global_report(_indexed()) + assert list(report.index.names) == ["agent.agent_name", "env.benchmark", "use_x"] + assert report["avg_reward"].iloc[0] >= report["avg_reward"].iloc[1] + + flags = ir.flag_report(report) + assert list(flags.index) == ["use_x"] + assert flags.loc["use_x", "avg_true"] == 0.75 + assert flags.loc["use_x", "avg_false"] == 0.25 + assert flags.loc["use_x", "ratio"] == 3.0 + + +def test_global_report_single_configuration(): + df = _result_df() + df = df[df["agent.agent_name"] == "ag1"].set_index("env.task_name") + report = ir.global_report(df) + assert "[ALL TASKS]" in report.index + assert report.loc["[ALL TASKS]", "avg_reward"] == 0.75 + assert ir.flag_report(report) is None + + +def test_summarize_study_and_split_by_key(): + df = _indexed() + summary = ir.summarize_study(df) + assert len(summary) == 2 + + parts = ir.split_by_key(df, "agent.agent_name") + assert set(parts) == {"ag1", "ag2"} + assert len(parts["ag1"]) == 4 + + +def test_report_constant_and_variables(capsys): + df = _result_df().assign(stack_trace=["t"] * 7 + ["u"]) + ir.report_constant_and_variables(df, show_stack_traces=False) + out = capsys.readouterr().out + assert "constant: c" in out + assert "agent._private: n_unique=8" in out + assert "stack_trace" not in out + + +def test_rename_bool_flags_and_shrink_columns(): + report = ir.global_report(_indexed()) + renamed = ir._rename_bool_flags(report.copy()) + assert set(renamed.index.get_level_values("use_x")) == {"✓", "-"} + + shrunk = ir.shrink_columns( + pd.DataFrame({"avg_reward": [0.5000]}, index=pd.Index([0], name="a_b")) + ) + assert list(shrunk.columns) == ["avg\nreward"] + assert shrunk.index.names == ["a\nb"] + assert shrunk.iloc[0, 0] == "0.5" + + +def test_find_diff(): + assert ir._find_diff((1, 2, 3), (1, 0, 4)) == [1, 2] + + +def test_map_err_key(): + assert ir.map_err_key(None) is None + msg = ( + "Error: your messages resulted in 12345 tokens\n" + "Exception uncaught by agent or environment in task miniwob.abc\n=== logs ===\nblah" + ) + key = ir.map_err_key(msg) + assert "x tokens" in key + assert "." in key + assert "blah" not in key + + +def test_categorize_error_without_error(): + assert ir.categorize_error({"err_msg": None}) is None + assert ir.categorize_error({}) is None + + +def test_format_log(): + assert ir._format_log(type("R", (), {"logs": None})()) == "No log found" + short = type("R", (), {"logs": "a\nb"})() + assert ir._format_log(short) == "a\nb" + + lines = [f"line {i}" for i in range(100)] + long_log = type("R", (), {"logs": "\n".join(lines)})() + out = ir._format_log(long_log, head_lines=2, tail_lines=3) + assert out.startswith("line 0\nline 1\n...") + assert out.endswith("line 97\nline 98\nline 99") + + with_tb = "\n".join(lines + ["action: click", "Traceback (most recent call last):", "boom"]) + out = ir._format_log(type("R", (), {"logs": with_tb})(), head_lines=2, tail_lines=3) + assert out.endswith("action: click\nTraceback (most recent call last):\nboom") + + +def test_benchmark_from_task_name(): + assert ir._benchmark_from_task_name("workarena.servicenow.x") == "workarena" + + +def test_is_stale_tracks_modification_times(tmp_path): + (tmp_path / "exp").mkdir() + f = tmp_path / "exp" / "file.txt" + f.write_text("a") + summary = tmp_path / "study_summary.csv" + + assert ir._is_stale(tmp_path, summary) # no summary yet + summary.write_text("x") + # the summary itself is excluded from the tracked modification times + assert not ir._is_stale(tmp_path, summary) + + mtimes = json.loads((tmp_path / "_last_modification_times.json").read_text()) + assert "exp/file.txt" in mtimes and "study_summary.csv" not in mtimes + + time.sleep(0.01) + f.write_text("b") + import os + + os.utime(f, (time.time() + 10, time.time() + 10)) + assert ir._is_stale(tmp_path, summary) + + +def test_get_study_summary_uses_cache(tmp_path, monkeypatch): + df = _indexed() + calls = [] + + def fake_load(study_dir, progress_fn=None): + calls.append(study_dir) + return df + + monkeypatch.setattr(ir, "load_result_df", fake_load) + sentinel = {} + ir.get_study_summary(tmp_path, sentinel=sentinel) + assert sentinel["from_cache"] is False + ir.get_study_summary(tmp_path, sentinel=sentinel) + assert sentinel["from_cache"] is True + ir.get_study_summary(tmp_path, ignore_cache=True, ignore_stale=True, sentinel=sentinel) + assert sentinel["from_cache"] is False + assert len(calls) == 2 + + +def test_get_study_summary_none_and_get_all_summaries(tmp_path, monkeypatch): + df = _indexed() + (tmp_path / "_hidden").mkdir() + (tmp_path / "empty").mkdir() + (tmp_path / "study").mkdir() + monkeypatch.setattr( + ir, + "load_result_df", + lambda d, progress_fn=None: df if d.name == "study" else None, + ) + assert ir.get_study_summary(tmp_path / "empty") is None + summaries = ir.get_all_summaries(tmp_path) + assert set(summaries.index) == {"study"} diff --git a/tests/benchmarks/test_abstract_env.py b/tests/benchmarks/test_abstract_env.py new file mode 100644 index 000000000..4c9b17683 --- /dev/null +++ b/tests/benchmarks/test_abstract_env.py @@ -0,0 +1,52 @@ +import pytest + +from agentlab.benchmarks.abstract_env import ( + AbstractBenchmark, + AbstractEnv, + add_step_timing_to_env_info_decorator, +) + + +class _DummyEnv(AbstractEnv): + def __init__(self, env_info=None): + self._env_info = env_info + + def reset(self, seed=None): + return {}, {} + + @add_step_timing_to_env_info_decorator + def step(self, action: str): + return {"action": action}, 1.0, True, False, self._env_info + + def close(self): + pass + + +def test_abstract_benchmark_defaults(): + bench = AbstractBenchmark(name="dummy", env_args_list=[1, 2]) + assert bench.get_version() == "1" + assert bench.prepare_backends() is None + assert bench.dependency_graph_over_tasks() == {} + + +def test_abstract_env_cannot_be_instantiated(): + with pytest.raises(TypeError): + AbstractEnv() + + +def test_step_timing_added_when_env_info_is_none(): + obs, reward, terminated, truncated, env_info = _DummyEnv(None).step("noop()") + assert obs == {"action": "noop()"} + assert reward == 1.0 and terminated and not truncated + assert env_info["action_exec_start"] <= env_info["action_exec_stop"] + assert env_info["action_exec_timeout"] == 0.0 + + +def test_step_timing_does_not_override_existing_values(): + existing = {"action_exec_start": 1, "action_exec_stop": 2, "action_exec_timeout": 3} + *_, env_info = _DummyEnv(dict(existing)).step("noop()") + assert env_info == existing + + +def test_step_timing_preserves_wrapped_name(): + assert _DummyEnv.step.__name__ == "step" diff --git a/tests/experiments/test_exp_utils.py b/tests/experiments/test_exp_utils.py new file mode 100644 index 000000000..049244a90 --- /dev/null +++ b/tests/experiments/test_exp_utils.py @@ -0,0 +1,115 @@ +import signal +import time +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from agentlab.experiments import exp_utils +from agentlab.experiments.exp_utils import ( + MockedExpArgs, + _episode_timeout, + add_dependencies, + run_exp, + timeout_manager, +) + + +def test_run_exp_calls_run(): + exp_arg = MagicMock() + exp_arg.run.return_value = "result" + assert run_exp(exp_arg, "dep1", "dep2") == "result" + exp_arg.run.assert_called_once() + + +def test_episode_timeout_without_max_steps(): + exp_arg = SimpleNamespace(env_args=SimpleNamespace()) + assert _episode_timeout(exp_arg) == 10 * 60 * 60 + + +def test_episode_timeout_with_max_steps(): + exp_arg = SimpleNamespace(env_args=SimpleNamespace(max_steps=5)) + assert _episode_timeout(exp_arg, avg_step_timeout=10) == 50 + + +def test_episode_timeout_uses_smaller_exp_timeout(): + exp_arg = SimpleNamespace(env_args=SimpleNamespace(max_steps=5), episode_timeout=20) + assert _episode_timeout(exp_arg, avg_step_timeout=10) == 20 + + +def test_timeout_manager_none_is_noop(): + with timeout_manager(None): + pass + + +def test_timeout_manager_raises_and_restores_handler(): + previous = signal.getsignal(signal.SIGALRM) + with pytest.raises(TimeoutError): + with timeout_manager(0.5): # floats are rounded up to at least 1 second + time.sleep(3) + assert signal.getsignal(signal.SIGALRM) == previous + assert signal.alarm(0) == 0 + + +def test_timeout_manager_no_timeout(): + with timeout_manager(5): + pass + assert signal.alarm(0) == 0 + + +def _exp_args(task_name): + exp_id = f"id_{task_name}" + ea = SimpleNamespace(env_args=SimpleNamespace(task_name=task_name), exp_id=None, depends_on=()) + + def make_id(): + ea.exp_id = exp_id + + ea.make_id = make_id + return ea + + +@pytest.mark.parametrize("deps", [None, {"a": [], "b": []}]) +def test_add_dependencies_noop(deps): + exp_args_list = [_exp_args("a"), _exp_args("b")] + assert add_dependencies(exp_args_list, deps) is exp_args_list + assert all(ea.exp_id is None for ea in exp_args_list) + + +def test_add_dependencies_maps_task_names_to_exp_ids(): + exp_args_list = [_exp_args("a"), _exp_args("b"), _exp_args("c")] + add_dependencies(exp_args_list, {"a": [], "b": ["a"], "c": ["a", "b"]}) + assert exp_args_list[0].depends_on == () + assert exp_args_list[1].depends_on == ("id_a",) + assert exp_args_list[2].depends_on == ("id_a", "id_b") + + +def test_add_dependencies_duplicate_task_names(): + with pytest.raises(ValueError, match="not unique"): + add_dependencies([_exp_args("a"), _exp_args("a")], {"a": ["a"]}) + + +def test_add_dependencies_missing_task(): + with pytest.raises(ValueError, match="missing"): + add_dependencies([_exp_args("a"), _exp_args("b")], {"a": ["b"]}) + + +def test_mocked_exp_args(monkeypatch): + monkeypatch.setattr(exp_utils, "sleep", lambda s: None) + exp = MockedExpArgs("1", depends_on=["0"]) + assert exp.exp_name == "exp_1" and exp.depends_on == ["0"] + assert MockedExpArgs("2").depends_on == [] + assert exp.run() is exp + assert exp.start_time <= exp.end_time + + +@pytest.mark.parametrize( + "fn, args", + [ + (exp_utils.make_seeds, (3,)), + (exp_utils.order, ([],)), + (exp_utils.hide_some_exp, (".", lambda e: True, True)), + ], +) +def test_deprecated_functions_raise(fn, args): + with pytest.raises(DeprecationWarning): + fn(*args) diff --git a/tests/experiments/test_reproducibility_util_extra.py b/tests/experiments/test_reproducibility_util_extra.py new file mode 100644 index 000000000..666aa2f73 --- /dev/null +++ b/tests/experiments/test_reproducibility_util_extra.py @@ -0,0 +1,114 @@ +import logging + +import pandas as pd +import pytest + +from agentlab.experiments import reproducibility_util as ru + + +def _report(rows): + return pd.DataFrame(rows).set_index("agent.agent_name") + + +GOOD_ROWS = [ + {"agent.agent_name": "a1", "avg_reward": 0.5, "std_err": 0.1, "n_err": 0, "n_completed": "4/4"}, + {"agent.agent_name": "a2", "avg_reward": 0.7, "std_err": 0.2, "n_err": 0, "n_completed": "4/4"}, +] + + +def _info(): + return { + "git_user": "me", + "agent_names": ["a1", "a2"], + "benchmark": "miniwob", + "date": "today", + } + + +def test_assert_compatible_ignores_result_keys(): + info = {"benchmark": "b", "date": "1", "avg_reward": 0.1} + old = {"benchmark": "b", "date": "2", "avg_reward": 0.9} + ru.assert_compatible(info, old) + + +def test_assert_compatible_raises_or_warns(caplog): + info, old = {"benchmark": "b1"}, {"benchmark": "b2"} + with pytest.raises(ValueError, match="Key benchmark has changed"): + ru.assert_compatible(info, old) + with caplog.at_level(logging.WARNING): + ru.assert_compatible(info, old, raise_if_incompatible=False) + assert "Key benchmark has changed" in caplog.text + + +def test_verify_report_ok(): + df = ru._verify_report(_report(GOOD_ROWS), ["a1", "a2"]) + assert list(df.index) == ["a1", "a2"] + + +def test_verify_report_agent_name_mismatch(): + with pytest.raises(ValueError, match="do not match"): + ru._verify_report(_report(GOOD_ROWS), ["a1"]) + + +def test_verify_report_duplicate_agent_names(): + with pytest.raises(ValueError, match="Duplicate"): + ru._verify_report(_report(GOOD_ROWS), ["a1", "a2", "a2"]) + + +@pytest.mark.parametrize( + "override, match", + [({"n_err": 1}, "has 1 errors"), ({"n_completed": "3/4"}, "3 completed tasks out of 4")], +) +def test_verify_report_incomplete(override, match, caplog): + rows = [dict(GOOD_ROWS[0], **override)] + with pytest.raises(ValueError, match=match): + ru._verify_report(_report(rows), ["a1"]) + with caplog.at_level(logging.WARNING): + ru._verify_report(_report(rows), ["a1"], strict_reproducibility=False) + assert "a1" in caplog.text + + +def test_get_csv_headers(tmp_path): + empty = tmp_path / "empty.csv" + empty.write_text("") + assert ru._get_csv_headers(empty) is None + f = tmp_path / "f.csv" + f.write_text("a,b\n1,2\n") + assert ru._get_csv_headers(f) == ["a", "b"] + + +def test_add_result_to_info(): + info = {"agent_name": "a2"} + ru._add_result_to_info(info, _report(GOOD_ROWS)) + assert info == { + "agent_name": "a2", + "avg_reward": 0.7, + "std_err": 0.2, + "n_err": 0, + "n_completed": "4/4", + } + + +def test_append_to_journal_creates_and_appends(tmp_path): + journal = tmp_path / "journal.csv" + ru.append_to_journal(_info(), _report(GOOD_ROWS), journal_path=journal) + ru.append_to_journal(_info(), _report(GOOD_ROWS), journal_path=journal) + + df = pd.read_csv(journal) + assert list(df.columns) == ["git_user", "agent_name", "benchmark", "date"] + assert list(df["agent_name"]) == ["a1", "a2", "a1", "a2"] + + +def test_append_to_journal_mismatched_agents(tmp_path): + with pytest.raises(ValueError, match="Mismatch"): + ru.append_to_journal(_info(), _report(GOOD_ROWS[:1]), journal_path=tmp_path / "j.csv") + + +def test_append_to_journal_default_path_without_git(tmp_path, monkeypatch): + def no_repo(module): + raise ru.InvalidGitRepositoryError() + + monkeypatch.setattr(ru, "_get_repo", no_repo) + monkeypatch.setattr(ru, "RESULTS_DIR", tmp_path) + ru.append_to_journal(_info(), _report(GOOD_ROWS)) + assert (tmp_path / "reproducibility_journal.csv").exists() diff --git a/tests/experiments/test_study_utils.py b/tests/experiments/test_study_utils.py new file mode 100644 index 000000000..8c2f55a50 --- /dev/null +++ b/tests/experiments/test_study_utils.py @@ -0,0 +1,65 @@ +from dataclasses import asdict + +import bgym +import pytest + +from agentlab.experiments import study as study_mod +from agentlab.experiments.loop import EnvArgs +from agentlab.experiments.study import ( + _convert_env_args, + _make_study_name, + get_most_recent_study, + set_demo_mode, +) + + +def test_make_study_name_single_agent_and_benchmark(): + assert _make_study_name(["agent"], ["miniwob"], "v1") == "agent-on-miniwob-v1" + + +def test_make_study_name_multiple(): + name = _make_study_name(["a", "b", "a"], ["x", "y", "z"]) + assert name == "2-agents-on-3-benchmarks" + + +def test_get_most_recent_study(tmp_path): + for name in [ + "2024-01-01_10-00-00_old", + "2024-06-01_10-00-00_newer_foo", + "2024-03-01_10-00-00_foo", + "_2025-01-01_10-00-00_hidden", + "not_a_date", + ]: + (tmp_path / name).mkdir() + (tmp_path / "2026-01-01_10-00-00_file").write_text("x") + + assert get_most_recent_study(tmp_path).name == "2024-06-01_10-00-00_newer_foo" + assert get_most_recent_study(tmp_path, contains="old").name == "2024-01-01_10-00-00_old" + assert get_most_recent_study(tmp_path, contains="missing") is None + + +def test_get_most_recent_study_defaults_to_results_dir(tmp_path, monkeypatch): + (tmp_path / "2024-01-01_10-00-00_x").mkdir() + monkeypatch.setattr(study_mod, "RESULTS_DIR", tmp_path) + assert get_most_recent_study().name == "2024-01-01_10-00-00_x" + + +def test_set_demo_mode(): + env_args = EnvArgs(task_name="miniwob.click-test") + set_demo_mode([env_args]) + assert env_args.viewport == {"width": 1280, "height": 720} + assert env_args.record_video is True + assert env_args.wait_for_user_message is False + assert env_args.slow_mo == 1000 + + +def test_convert_env_args(): + new = EnvArgs(task_name="miniwob.a") + old = bgym.EnvArgs(task_name="miniwob.b", task_seed=3) + converted = _convert_env_args([new, old]) + assert converted[0] is new + assert isinstance(converted[1], EnvArgs) + assert asdict(converted[1])["task_seed"] == 3 + + with pytest.raises(TypeError, match="Unexpected type"): + _convert_env_args(["nope"]) diff --git a/tests/llm/test_chat_api_unit.py b/tests/llm/test_chat_api_unit.py new file mode 100644 index 000000000..36aa27070 --- /dev/null +++ b/tests/llm/test_chat_api_unit.py @@ -0,0 +1,240 @@ +"""Offline unit tests for chat_api (no network calls).""" + +import threading +from types import SimpleNamespace +from unittest.mock import MagicMock + +import openai +import pytest + +from agentlab.llm import chat_api +from agentlab.llm.chat_api import ( + AnthropicChatModel, + AnthropicModelArgs, + ChatModel, + ChatModelArgs, + CheatMiniWoBLLM, + CheatMiniWoBLLMArgs, + RetryError, + SelfHostedModelArgs, + _extract_wait_time, + handle_error, + make_assistant_message, + make_system_message, + make_user_message, +) +from agentlab.llm.llm_utils import Discussion, HumanMessage, SystemMessage + + +@pytest.fixture(autouse=True) +def fresh_tracker(monkeypatch): + # other tests may leave TRACKER.instance set to None in this thread + monkeypatch.setattr(chat_api.tracking, "TRACKER", threading.local()) + + +@pytest.fixture +def no_sleep(monkeypatch): + sleeps = [] + monkeypatch.setattr(chat_api.time, "sleep", lambda s: sleeps.append(s)) + return sleeps + + +def test_make_messages(): + assert make_system_message("s") == {"role": "system", "content": "s"} + assert make_user_message("u") == {"role": "user", "content": "u"} + assert make_assistant_message("a") == {"role": "assistant", "content": "a"} + + +PROMPT = "Some header\n[12] button 'Click me'\n" + + +@pytest.mark.parametrize( + "messages", + [ + [make_system_message("sys"), make_user_message(PROMPT)], + Discussion([SystemMessage("sys"), HumanMessage(PROMPT)]), + ], +) +def test_cheat_miniwob_llm(messages): + answer = CheatMiniWoBLLM()(messages) + assert answer["role"] == "assistant" + assert 'click("12")' in answer["content"] + + +def test_cheat_miniwob_llm_waits_and_fails_without_button(monkeypatch, no_sleep): + llm = CheatMiniWoBLLMArgs(wait_time=2).make_model() + with pytest.raises(Exception, match="Can't find the button"): + llm([make_system_message("sys"), make_user_message("no buttons here")]) + assert no_sleep == [2] + + +def test_cheat_miniwob_llm_args_server_noops(): + args = CheatMiniWoBLLMArgs() + assert args.prepare_server() is None + assert args.close_server() is None + + +@pytest.mark.parametrize( + "msg, expected", + [ + ("Rate limit. Please try again in 90.5s.", 90.5), + ("Rate limit. Please try again in 5s.", 60), + ("unrelated error", 60), + ], +) +def test_extract_wait_time(msg, expected): + assert _extract_wait_time(msg) == expected + + +def test_handle_error(no_sleep): + err = openai.OpenAIError("please try again in 2s") + assert handle_error(err, 0, min_retry_wait_time=1, max_retry=3) == "please try again in 2s" + assert no_sleep == [2.0] + + +def test_handle_error_reraises_non_openai_errors(): + with pytest.raises(KeyError): + handle_error(KeyError("x"), 0, 1, 3) + + +def _completion(contents, usage=True): + return SimpleNamespace( + usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5) if usage else None, + choices=[ + SimpleNamespace(message=SimpleNamespace(content=c), log_probs="lp") for c in contents + ], + ) + + +def _chat_model(create, **kwargs): + client = MagicMock() + client.chat.completions.create.side_effect = create + model = ChatModel("m", client_class=MagicMock(return_value=client), **kwargs) + return model, client + + +def test_chat_model_pricing(): + model, _ = _chat_model(None, pricing_func=lambda: {"m": {"prompt": "0.1", "completion": "0.2"}}) + assert (model.input_cost, model.output_cost) == (0.1, 0.2) + + model, _ = _chat_model(None, pricing_func=lambda: {}) + assert (model.input_cost, model.output_cost) == (0.0, 0.0) + + +def test_chat_model_api_key_from_env(monkeypatch): + monkeypatch.setenv("MY_KEY", "secret") + model, _ = _chat_model(None, api_key_env_var="MY_KEY") + assert model.api_key == "secret" + + +def test_chat_model_single_sample_with_log_probs(): + model, client = _chat_model([_completion(["hello"])], log_probs=True, temperature=0.3) + res = model([make_user_message("hi")]) + assert res["content"] == "hello" and res["log_probs"] == "lp" + kwargs = client.chat.completions.create.call_args.kwargs + assert kwargs["temperature"] == 0.3 and kwargs["n"] == 1 + assert model.get_stats() == {"n_retry_llm": 1} + + +def test_chat_model_multiple_samples(): + model, _ = _chat_model([_completion(["a", "b"])]) + res = model([make_user_message("hi")], n_samples=2, temperature=0) + assert [r["content"] for r in res] == ["a", "b"] + + +def test_chat_model_retries_on_missing_usage(no_sleep): + model, _ = _chat_model( + [_completion(["x"], usage=False), _completion(["ok"])], min_retry_wait_time=0 + ) + assert model([make_user_message("hi")])["content"] == "ok" + assert model.retries == 2 and model.success + assert len(model.error_types) == 1 + + +def test_chat_model_raises_retry_error(no_sleep): + model, _ = _chat_model(openai.OpenAIError("down"), max_retry=2, min_retry_wait_time=0) + with pytest.raises(RetryError, match="after 2 retries"): + model([make_user_message("hi")]) + assert model.error_types == ["down", "down"] + + +def test_chat_model_rejects_zero_retry(): + with pytest.raises(AssertionError): + _chat_model(None, max_retry=0) + + +def test_azure_chat_model_requires_env(monkeypatch): + monkeypatch.delenv("AZURE_OPENAI_API_KEY", raising=False) + monkeypatch.delenv("AZURE_OPENAI_ENDPOINT", raising=False) + with pytest.raises(AssertionError, match="AZURE_OPENAI_API_KEY"): + chat_api.AzureChatModel("gpt-4o") + with pytest.raises(AssertionError, match="AZURE_OPENAI_ENDPOINT"): + chat_api.AzureChatModel("gpt-4o", api_key="k") + + +def test_azure_chat_model_client_args(monkeypatch): + monkeypatch.setenv("AZURE_OPENAI_ENDPOINT", "https://example.azure.com") + monkeypatch.setattr(chat_api.tracking, "get_pricing_openai", lambda: {}) + model = chat_api.AzureChatModel("gpt-4o", api_key="k", deployment_name="old") + assert str(model.client.base_url).startswith("https://example.azure.com") + + +def test_openrouter_and_vllm_models(monkeypatch): + monkeypatch.setattr(chat_api.tracking, "get_pricing_openrouter", lambda: {}) + router = chat_api.OpenRouterChatModel("some/model", api_key="k") + assert "openrouter.ai" in str(router.client.base_url) + + vllm = chat_api.VLLMChatModel("some/model", api_key="k", n_retry_server=2) + assert "0.0.0.0:8000" in str(vllm.client.base_url) + assert vllm.max_retry == 2 + + +def test_self_hosted_model_args_vllm_and_unknown_backend(): + args = SelfHostedModelArgs(model_name="m", backend="vllm", max_new_tokens=10) + model = args.make_model() + assert isinstance(model, chat_api.VLLMChatModel) + assert model.max_tokens == 10 + + with pytest.raises(ValueError, match="not supported"): + SelfHostedModelArgs(model_name="m", backend="other").make_model() + + +def test_chat_model_args_is_deprecated(): + with pytest.warns(DeprecationWarning): + args = ChatModelArgs(model_name="m") + assert args.make_model() is None + + +def _anthropic_response(text): + return SimpleNamespace( + content=[SimpleNamespace(text=text)], + usage=SimpleNamespace(input_tokens=1, output_tokens=2), + ) + + +def test_anthropic_chat_model_converts_messages(): + model = AnthropicModelArgs(model_name="claude", max_new_tokens=20, temperature=0.1).make_model() + model.client = MagicMock() + model.client.messages.create.return_value = _anthropic_response("hey") + + res = model([make_system_message("sys"), make_user_message("hi")]) + + assert res["content"] == "hey" + kwargs = model.client.messages.create.call_args.kwargs + assert kwargs["system"] == "sys" + assert kwargs["messages"] == [{"role": "user", "content": "hi"}] + assert kwargs["max_tokens"] == 20 and kwargs["temperature"] == 0.1 + + +def test_anthropic_chat_model_retries(no_sleep): + model = AnthropicChatModel("claude", api_key="k", max_retry=2) + model.client = MagicMock() + model.client.messages.create.side_effect = [RuntimeError("x"), _anthropic_response("ok")] + res = model([make_user_message("hi")], temperature=0) + assert res["content"] == "ok" + assert "system" not in model.client.messages.create.call_args.kwargs + assert no_sleep == [60] + + model.client.messages.create.side_effect = RuntimeError("down") + with pytest.raises(RuntimeError, match="down"): + model([make_user_message("hi")]) diff --git a/tests/llm/test_huggingface_utils_unit.py b/tests/llm/test_huggingface_utils_unit.py new file mode 100644 index 000000000..b2c4e4ce3 --- /dev/null +++ b/tests/llm/test_huggingface_utils_unit.py @@ -0,0 +1,144 @@ +"""Offline unit tests for HFBaseChatModel using a fake tokenizer and LLM.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from transformers import GPT2TokenizerFast + +from agentlab.llm import huggingface_utils +from agentlab.llm.huggingface_utils import HFBaseChatModel, _prepend_system_to_first_user +from agentlab.llm.llm_utils import Discussion, HumanMessage, SystemMessage +from agentlab.llm.prompt_templates import STARCHAT_PROMPT_TEMPLATE + + +class _FakeTokenizer: + def __init__(self, fail_with=None): + self.fail_with = fail_with + self.calls = [] + + def apply_chat_template(self, messages, tokenize=False): + self.calls.append([dict(m) for m in messages]) + if self.fail_with and any(m["role"] == "system" for m in messages): + raise Exception(self.fail_with) + return "|".join(f"{m['role']}:{m['content']}" for m in messages) + + +def _make_model(monkeypatch, tokenizer, base_model_name=None, n_retry=3, log_probs=False): + loaded = [] + + def from_pretrained(name): + loaded.append(name) + return tokenizer + + monkeypatch.setattr(huggingface_utils.AutoTokenizer, "from_pretrained", from_pretrained) + monkeypatch.setattr(huggingface_utils.time, "sleep", lambda s: None) + model = HFBaseChatModel("model", base_model_name, n_retry, log_probs) + model.temperature = 0.5 + model.llm = MagicMock(return_value="answer") + return model, loaded + + +def _messages(): + return [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + ] + + +def test_loads_tokenizer_from_base_model(monkeypatch): + _, loaded = _make_model(monkeypatch, _FakeTokenizer(), base_model_name="base") + assert loaded == ["base"] + _, loaded = _make_model(monkeypatch, _FakeTokenizer()) + assert loaded == ["model"] + + +def test_call_with_chat_template(monkeypatch): + model, _ = _make_model(monkeypatch, _FakeTokenizer()) + res = model(_messages()) + assert res["content"] == "answer" + model.llm.assert_called_once_with("system:sys|user:hi", temperature=0.5) + assert model._llm_type() == "huggingface" + + +def test_call_merges_discussion(monkeypatch): + tokenizer = _FakeTokenizer() + model, _ = _make_model(monkeypatch, tokenizer) + msg = HumanMessage("a") + msg.add_text("b") + model(Discussion([SystemMessage("sys"), msg]), temperature=0.1) + assert tokenizer.calls[0][1]["content"] == "a\nb" + assert model.llm.call_args.kwargs["temperature"] == 0.1 + + +def test_call_falls_back_when_system_role_unsupported(monkeypatch): + tokenizer = _FakeTokenizer(fail_with="Conversation roles must alternate user/assistant") + model, _ = _make_model(monkeypatch, tokenizer) + model(_messages()) + model.llm.assert_called_once_with("user:sys\nhi", temperature=0.5) + + +def test_call_reraises_other_template_errors(monkeypatch): + model, _ = _make_model(monkeypatch, _FakeTokenizer(fail_with="other")) + with pytest.raises(Exception, match="other"): + model(_messages()) + + +def test_gpt2_tokenizer_uses_prompt_template(monkeypatch): + gpt2 = MagicMock(spec=GPT2TokenizerFast) + monkeypatch.setattr(huggingface_utils.AutoTokenizer, "from_pretrained", lambda n: gpt2) + model = HFBaseChatModel("bigcode/starcoder", None, 1, False) + model.temperature = 0.5 + model.llm = MagicMock(return_value="x") + assert model.tokenizer is None + assert model.prompt_template is STARCHAT_PROMPT_TEMPLATE + model(_messages()) + assert model.llm.call_args.args[0].startswith("<|system|>") + + +def test_multiple_samples_and_retries(monkeypatch): + model, _ = _make_model(monkeypatch, _FakeTokenizer(), n_retry=3) + model.llm.side_effect = [RuntimeError("busy"), "a", "b"] + res = model(_messages(), n_samples=2) + assert [r["content"] for r in res] == ["a", "b"] + assert model.llm.call_count == 3 + + +def test_retries_exhausted(monkeypatch): + model, _ = _make_model(monkeypatch, _FakeTokenizer(), n_retry=2) + model.llm.side_effect = RuntimeError("down") + with pytest.raises(RuntimeError, match="down"): + model(_messages()) + assert model.llm.call_count == 2 + + +def test_log_probs(monkeypatch): + model, _ = _make_model(monkeypatch, _FakeTokenizer(), log_probs=True) + model.llm.return_value = SimpleNamespace(generated_text="txt", details="details") + res = model(_messages()) + assert res["content"] == "txt" and res["log_probs"] == "details" + + +def test_prepend_system_to_first_user(): + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "q1"}, + {"role": "user", "content": "q2"}, + ] + out = _prepend_system_to_first_user(messages) + assert out == [{"role": "user", "content": "sys\nq1"}, {"role": "user", "content": "q2"}] + + +def test_prepend_system_without_system_or_user(): + no_system = [{"role": "user", "content": "q"}] + assert _prepend_system_to_first_user(no_system) == [{"role": "user", "content": "q"}] + only_system = [{"role": "system", "content": "s"}] + assert _prepend_system_to_first_user(only_system) == [{"role": "system", "content": "s"}] + + +def test_prepend_system_with_column_remap(): + messages = [{"r": "system", "t": "sys"}, {"r": "human", "t": "q"}] + out = _prepend_system_to_first_user( + messages, column_remap={"HumanMessage": "human", "role": "r", "text": "t"} + ) + assert out == [{"r": "human", "t": "sys\nq"}] diff --git a/tests/llm/test_litellm_api_unit.py b/tests/llm/test_litellm_api_unit.py new file mode 100644 index 000000000..9064f20ec --- /dev/null +++ b/tests/llm/test_litellm_api_unit.py @@ -0,0 +1,194 @@ +"""Offline unit tests for LiteLLMModel (the litellm client is replaced by a fake).""" + +import json + +import pytest +from litellm import ModelResponse + +from agentlab.llm.litellm_api import LiteLLMAPIMessageBuilder, LiteLLMModel, LiteLLMModelArgs +from agentlab.llm.response_api import APIPayload, ToolCall, ToolCalls + +RESPONSE_API_TOOL = { + "type": "function", + "name": "click", + "description": "Click an element.", + "parameters": {"type": "object", "properties": {"bid": {"type": "string"}}}, +} + + +def _response(content="", tool_calls=None, reasoning=None): + message = {"role": "assistant", "content": content} + if tool_calls is not None: + message["tool_calls"] = [ + { + "id": f"call_{i}", + "type": "function", + "function": {"name": name, "arguments": json.dumps(args)}, + } + for i, (name, args) in enumerate(tool_calls) + ] + if reasoning is not None: + message["reasoning"] = reasoning + return ModelResponse( + choices=[{"index": 0, "message": message, "finish_reason": "stop"}], + model="gpt-4o-mini", + usage={"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5}, + ) + + +@pytest.fixture +def model(): + m = LiteLLMModelArgs( + model_name="openai/gpt-4o-mini", max_new_tokens=50, temperature=0.2 + ).make_model() + calls = [] + + def fake_client(**kwargs): + calls.append(kwargs) + return m._next_response + + m.client = fake_client + m._next_response = _response("hello") + m.calls = calls + return m + + +def _user_payload(**kwargs): + return APIPayload(messages=[LiteLLMAPIMessageBuilder.user().add_text("hi")], **kwargs) + + +def test_model_args_builder(): + assert LiteLLMModelArgs(model_name="x").get_message_builder() is LiteLLMAPIMessageBuilder + + +def test_call_api_basic_params(model): + model._call_api(_user_payload()) + params = model.calls[-1] + assert params["model"] == "openai/gpt-4o-mini" + assert params["temperature"] == 0.2 + assert params["max_completion_tokens"] == 50 + assert params["num_retries"] == 5 + assert params["messages"] == [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] + assert "tools" not in params and "tool_choice" not in params + + +def test_call_api_omits_unset_params(model): + model.temperature = None + model.max_tokens = None + model._call_api(_user_payload()) + assert "temperature" not in model.calls[-1] + assert "max_completion_tokens" not in model.calls[-1] + + +def test_call_api_converts_tools_and_caching(model): + model._call_api( + _user_payload( + tools=[dict(RESPONSE_API_TOOL)], + tool_choice="any", + reasoning_effort="low", + cache_tool_definition=True, + cache_complete_prompt=True, + ) + ) + params = model.calls[-1] + assert params["tools"][0]["function"]["name"] == "click" + assert params["tools"][-1]["cache_control"] == {"type": "ephemeral"} + assert params["tool_choice"] == "required" + assert params["reasoning_effort"] == "low" + assert params["messages"][-1]["content"][-1]["cache_control"] == {"type": "ephemeral"} + + +def test_call_api_keeps_chat_completion_tools_and_forces_tool(model): + chat_tool = {"type": "function", "function": {"name": "click"}} + model._call_api(_user_payload(tools=[chat_tool], force_call_tool="click")) + params = model.calls[-1] + assert params["tools"] == [chat_tool] + assert params["tool_choice"] == {"type": "function", "function": {"name": "click"}} + + model._call_api(_user_payload(tool_choice="auto")) + assert model.calls[-1]["tool_choice"] == "auto" + + +def test_api_payload_rejects_tool_choice_with_force(): + with pytest.raises(ValueError, match="mutually exclusive"): + APIPayload(tool_choice="auto", force_call_tool="click") + + +def test_parse_response_single_tool_call(model): + model._next_response = _response("ok", tool_calls=[("click", {"bid": "a1"})]) + out = model(_user_payload()) + assert out.think == "ok" + assert out.action == "click(bid='a1')" + assert len(out.tool_calls) == 1 + assert model.stats.stats_dict["n_api_calls"] == 1 + + +def test_parse_response_multiple_tool_calls(model): + calls = [("click", {"bid": "a1"}), ("fill", {"bid": "a2", "value": "x"})] + model._next_response = _response(tool_calls=calls) + out = model._parse_response(model._next_response) + assert out.action == "click(bid='a1')\nfill(bid='a2', value='x')" + + model.use_only_first_toolcall = True + out = model._parse_response(model._next_response) + assert out.action == "click(bid='a1')" + assert len(out.tool_calls) == 1 + + +def test_parse_response_without_tool_calls(model): + out = model._parse_response(_response("just text")) + assert out.action is None and out.tool_calls is None + + model.action_space_as_tools = False + assert model._parse_response(_response("x")).action is None + + +def test_thinking_content_is_wrapped(model): + think = model._extract_thinking_content_from_response(_response("answer", reasoning="because")) + assert think == "because\nanswer" + + +def test_format_tools_for_chat_completion(): + assert LiteLLMModel.format_tools_for_chat_completion(None) is None + (tool,) = LiteLLMModel.format_tools_for_chat_completion([RESPONSE_API_TOOL]) + assert tool == { + "type": "function", + "function": { + "name": "click", + "description": "Click an element.", + "parameters": RESPONSE_API_TOOL["parameters"], + }, + } + + +def _responded_tool_calls(n=2): + response = _response(tool_calls=[("click", {"bid": str(i)}) for i in range(n)]) + raw = response.choices[0].message.to_dict()["tool_calls"] + calls = [] + for i, raw_call in enumerate(raw): + tc = ToolCall(name="click", arguments={"bid": str(i)}, raw_call=raw_call) + tc.response_text(f"clicked {i}") + calls.append(tc) + return ToolCalls(tool_calls=calls, raw_calls=response) + + +def test_message_builder_tool_responses(): + builder = LiteLLMAPIMessageBuilder.add_responded_tool_calls(_responded_tool_calls()) + out = builder.prepare_message() + assert len(out[0].tool_calls) == 2 + assert out[1:] == [ + {"name": "click", "role": "tool", "tool_call_id": "call_0", "content": "clicked 0"}, + {"name": "click", "role": "tool", "tool_call_id": "call_1", "content": "clicked 1"}, + ] + + +def test_message_builder_only_first_tool_call(): + builder = LiteLLMAPIMessageBuilder.add_responded_tool_calls(_responded_tool_calls()) + out = builder.prepare_message(use_only_first_toolcall=True) + assert len(out[0].tool_calls) == 1 + + +def test_message_builder_tool_without_calls(): + builder = LiteLLMAPIMessageBuilder("tool") + with pytest.raises(ValueError, match="No tool calls"): + builder.handle_tool_call() diff --git a/tests/llm/test_llm_utils_extra.py b/tests/llm/test_llm_utils_extra.py new file mode 100644 index 000000000..1b315a418 --- /dev/null +++ b/tests/llm/test_llm_utils_extra.py @@ -0,0 +1,467 @@ +"""Unit tests for llm_utils helpers that don't require network access.""" + +import logging +from types import SimpleNamespace +from unittest.mock import MagicMock + +import anthropic +import httpx +import numpy as np +import openai +import pytest +from PIL import Image + +from agentlab.llm import llm_utils +from agentlab.llm.llm_utils import ( + AIMessage, + BaseMessage, + Discussion, + HumanMessage, + ParseError, + SystemMessage, +) + + +@pytest.fixture +def no_sleep(monkeypatch): + sleeps = [] + monkeypatch.setattr(llm_utils.time, "sleep", lambda s: sleeps.append(s)) + return sleeps + + +class _RateLimit(Exception): + pass + + +class _ApiError(Exception): + def __init__(self, status_code=None): + super().__init__(f"status {status_code}") + self.status_code = status_code + + +def _call(client_function, **kwargs): + defaults = dict( + api_params={"x": 1}, + is_response_valid_fn=lambda r: r == "ok", + rate_limit_exceptions=(_RateLimit,), + api_error_exceptions=(_ApiError,), + get_status_code_fn=lambda e: e.status_code, + max_retries=3, + initial_retry_delay_seconds=1, + max_retry_delay_seconds=3, + ) + defaults.update(kwargs) + return llm_utils.generic_call_api_with_retries(client_function, **defaults) + + +# --------------------------------------------------------------------------- +# generic_call_api_with_retries +# --------------------------------------------------------------------------- + + +def test_generic_call_success_first_try(no_sleep): + fn = MagicMock(return_value="ok") + assert _call(fn) == "ok" + fn.assert_called_once_with(x=1) + assert no_sleep == [] + + +def test_generic_call_retries_invalid_response_then_succeeds(no_sleep): + fn = MagicMock(side_effect=["bad", "bad", "ok"]) + assert _call(fn) == "ok" + assert fn.call_count == 3 + # first invalid retry uses exponent 0, second one exponent 1 + assert no_sleep == [1, 2] + + +def test_generic_call_invalid_response_exhausts_retries(no_sleep): + fn = MagicMock(return_value="bad") + with pytest.raises(RuntimeError, match="after 3 retries"): + _call(fn) + assert fn.call_count == 3 + + +def test_generic_call_rate_limit_backoff_is_capped(no_sleep): + fn = MagicMock(side_effect=[_RateLimit(), _RateLimit(), "ok"]) + assert _call(fn) == "ok" + assert no_sleep == [2, 3] # 1*2**1, min(1*2**2, 3) + + +def test_generic_call_rate_limit_exhausts_retries(no_sleep): + fn = MagicMock(side_effect=_RateLimit()) + with pytest.raises(RuntimeError): + _call(fn) + assert fn.call_count == 3 + + +@pytest.mark.parametrize("status", [429, 500, 503]) +def test_generic_call_retriable_api_error(no_sleep, status): + fn = MagicMock(side_effect=[_ApiError(status), "ok"]) + assert _call(fn) == "ok" + assert len(no_sleep) == 1 + + +def test_generic_call_retriable_api_error_exhausts_retries(no_sleep): + fn = MagicMock(side_effect=_ApiError(500)) + with pytest.raises(RuntimeError): + _call(fn) + assert fn.call_count == 3 + + +@pytest.mark.parametrize("status", [400, None]) +def test_generic_call_non_retriable_api_error_is_raised(no_sleep, status): + err = _ApiError(status) + fn = MagicMock(side_effect=err) + with pytest.raises(_ApiError) as exc_info: + _call(fn) + assert exc_info.value is err + fn.assert_called_once() + + +def test_generic_call_status_code_fn_failure_is_non_retriable(no_sleep): + def broken_status(e): + raise RuntimeError("nope") + + fn = MagicMock(side_effect=_ApiError(500)) + with pytest.raises(_ApiError): + _call(fn, get_status_code_fn=broken_status) + + +def test_generic_call_unexpected_exception_is_raised(no_sleep): + fn = MagicMock(side_effect=KeyError("boom")) + with pytest.raises(KeyError): + _call(fn) + fn.assert_called_once() + + +# --------------------------------------------------------------------------- +# OpenAI / Anthropic wrappers +# --------------------------------------------------------------------------- + + +def _openai_rate_limit_error(): + request = httpx.Request("POST", "https://api.openai.com/v1/chat/completions") + return openai.RateLimitError( + "rate limited", response=httpx.Response(429, request=request), body=None + ) + + +@pytest.mark.parametrize( + "response", + [ + SimpleNamespace(error=None, choices=[1]), + SimpleNamespace(error=None, output=[1]), + ], +) +def test_call_openai_api_with_retries_valid(no_sleep, response): + fn = MagicMock(return_value=response) + assert llm_utils.call_openai_api_with_retries(fn, {"model": "m"}) is response + + +def test_call_openai_api_with_retries_retries_on_error_and_rate_limit(no_sleep): + good = SimpleNamespace(error=None, choices=[1]) + fn = MagicMock( + side_effect=[ + SimpleNamespace(error="oops"), + SimpleNamespace(error=None, choices=[]), + _openai_rate_limit_error(), + good, + ] + ) + assert llm_utils.call_openai_api_with_retries(fn, {}, max_retries=4) is good + assert fn.call_count == 4 + assert len(no_sleep) == 3 + + +def test_call_openai_api_with_retries_non_retriable_api_error(no_sleep): + request = httpx.Request("POST", "https://api.openai.com") + err = openai.APIError("bad request", request=request, body=None) + fn = MagicMock(side_effect=err) + with pytest.raises(openai.APIError): + llm_utils.call_openai_api_with_retries(fn, {}) + + +@pytest.mark.parametrize( + "response, valid", + [ + (SimpleNamespace(type="message", content=[]), True), + (SimpleNamespace(type="message", content=None), False), + (SimpleNamespace(type="other"), False), + ({"type": "error", "error": "x"}, False), + (None, False), + ], +) +def test_call_anthropic_api_with_retries_validation(no_sleep, response, valid): + fn = MagicMock(return_value=response) + if valid: + assert llm_utils.call_anthropic_api_with_retries(fn, {}, max_retries=2) is response + else: + with pytest.raises(RuntimeError): + llm_utils.call_anthropic_api_with_retries(fn, {}, max_retries=2) + assert fn.call_count == 2 + + +def test_call_anthropic_api_with_retries_server_error_then_success(no_sleep): + request = httpx.Request("POST", "https://api.anthropic.com") + err = anthropic.InternalServerError( + "server", response=httpx.Response(500, request=request), body=None + ) + good = SimpleNamespace(type="message", content=["hi"]) + fn = MagicMock(side_effect=[err, good]) + assert llm_utils.call_anthropic_api_with_retries(fn, {}) is good + assert len(no_sleep) == 1 + + +def test_supports_tool_calling_for_openrouter(monkeypatch): + client = MagicMock() + client.chat.completions.create.return_value.to_dict.return_value = { + "choices": [{"message": {"tool_calls": [{}]}}] + } + monkeypatch.setattr(openai, "Client", MagicMock(return_value=client)) + assert llm_utils.supports_tool_calling_for_openrouter("some/model") is True + + client.chat.completions.create.return_value.to_dict.return_value = { + "choices": [{"message": {"content": "no tools"}}] + } + assert llm_utils.supports_tool_calling_for_openrouter("some/model") is False + + client.chat.completions.create.side_effect = RuntimeError("network") + # errors default to True so that the check never blocks a run + assert llm_utils.supports_tool_calling_for_openrouter("some/model") is True + + +# --------------------------------------------------------------------------- +# retry_multiple +# --------------------------------------------------------------------------- + + +def _parser(content): + if "good" not in content: + raise ParseError("not good") + return {"action": content} + + +def test_retry_multiple_returns_valid_samples(): + chat = MagicMock( + return_value=[ + {"role": "assistant", "content": "bad"}, + {"role": "assistant", "content": "good 1"}, + {"role": "assistant", "content": "good 2"}, + ] + ) + messages = Discussion() + parsed, tries = llm_utils.retry_multiple(chat, messages, 3, _parser, num_samples=3) + assert parsed == [{"action": "good 1"}, {"action": "good 2"}] + assert tries == 0 + chat.assert_called_once_with(messages, n_samples=3) + assert len(messages) == 1 + + +def test_retry_multiple_wraps_single_answer_and_retries(): + chat = MagicMock( + side_effect=[ + {"role": "assistant", "content": "bad"}, + {"role": "assistant", "content": "good"}, + ] + ) + messages = Discussion() + parsed, tries = llm_utils.retry_multiple(chat, messages, 3, _parser) + assert parsed == [{"action": "good"}] + assert tries == 1 + # assistant answer, user error feedback, assistant answer + assert [m["role"] for m in messages] == ["assistant", "user", "assistant"] + assert "not good" in messages[1]["content"] + + +def test_retry_multiple_raises_after_n_retry(): + chat = MagicMock(return_value={"role": "assistant", "content": "bad"}) + with pytest.raises(ParseError): + llm_utils.retry_multiple(chat, Discussion(), 2, _parser) + assert chat.call_count == 2 + + +# --------------------------------------------------------------------------- +# messages / parsing helpers +# --------------------------------------------------------------------------- + + +def test_messages_to_dict_handles_dicts_and_strings(): + d = llm_utils.messages_to_dict([{"role": "user", "content": "hi"}, "raw text"]) + assert isinstance(d, Discussion) + assert d[0]["role"] == "user" + assert d[1] == {"role": "", "content": "raw text"} + + +def test_messages_to_dict_rejects_unknown_types(): + with pytest.raises(ValueError, match="Unknown message type"): + llm_utils.messages_to_dict([42]) + + +def test_messages_to_dict_langchain_fallback(monkeypatch): + class FakeLC: + def __init__(self, type, content): + self.type = type + self.content = content + + monkeypatch.setattr(llm_utils, "LangchainBaseMessage", FakeLC) + monkeypatch.setattr(llm_utils, "convert_message_to_dict", None) + d = llm_utils.messages_to_dict([FakeLC("human", "q"), FakeLC("ai", "a"), FakeLC("x", 1)]) + assert [m["role"] for m in d] == ["user", "assistant", "x"] + assert d[2]["content"] == "1" + + +def test_messages_to_dict_langchain_converter(monkeypatch): + class FakeLC: + pass + + monkeypatch.setattr(llm_utils, "LangchainBaseMessage", FakeLC) + monkeypatch.setattr( + llm_utils, "convert_message_to_dict", lambda m: {"role": "system", "content": "c"} + ) + d = llm_utils.messages_to_dict([FakeLC()]) + assert d[0] == {"role": "system", "content": "c"} + + +def test_parse_html_tags_optional_and_multiple(): + text = "12 x " + content, valid, msg = llm_utils.parse_html_tags(text, keys=("a",), optional_keys=("b", "c")) + assert not valid + assert "multiple instances" in msg + assert content == {"a": "1", "b": "x"} + + content, valid, msg = llm_utils.parse_html_tags( + text, keys=("a",), optional_keys=("c",), merge_multiple=True + ) + assert valid and msg == "" + assert content["a"] == "1\n2" + + +def test_parse_html_tags_raise(): + assert llm_utils.parse_html_tags_raise("v", keys=("k",)) == {"k": "v"} + with pytest.raises(ParseError, match="Missing the key "): + llm_utils.parse_html_tags_raise("nothing", keys=("k",)) + + +def test_yaml_parser_invalid(): + with pytest.warns(UserWarning): + value, valid, retry_message = llm_utils.yaml_parser("a: [unclosed") + assert value == {} and not valid + assert "not a valid yaml" in retry_message + + +def test_truncate_tokens_short_text_untouched(): + assert llm_utils.truncate_tokens("hello world", max_tokens=100) == "hello world" + + +def test_get_tokenizer_falls_back_to_gpt4(monkeypatch): + llm_utils.get_tokenizer.cache_clear() + + def fail(*args, **kwargs): + raise OSError("no hub") + + monkeypatch.setattr(llm_utils.AutoTokenizer, "from_pretrained", fail) + try: + enc = llm_utils.get_tokenizer("definitely-not-a-model") + assert enc.encode("hello") == llm_utils.get_tokenizer("gpt-4").encode("hello") + finally: + llm_utils.get_tokenizer.cache_clear() + + +@pytest.mark.parametrize("mode", ["RGB", "RGBA"]) +def test_image_to_base64_urls(mode): + img = Image.new(mode, (4, 4)) + assert llm_utils.image_to_jpg_base64_url(img).startswith("data:image/jpeg;base64,") + assert llm_utils.image_to_png_base64_url(img).startswith("data:image/png;base64,") + + +def test_image_to_base64_urls_from_numpy(): + arr = np.zeros((4, 4, 3), dtype=np.uint8) + assert llm_utils.image_to_jpg_base64_url(arr).startswith("data:image/jpeg;base64,") + assert llm_utils.image_to_png_base64_url(arr).startswith("data:image/png;base64,") + + +# --------------------------------------------------------------------------- +# BaseMessage / Discussion +# --------------------------------------------------------------------------- + + +def test_base_message_rejects_unknown_kwargs(): + with pytest.raises(ValueError, match="Invalid attributes"): + BaseMessage("user", "x", foo=1) + + +def test_message_subclasses_roles(): + assert SystemMessage("s")["role"] == "system" + assert HumanMessage("h")["role"] == "user" + ai = AIMessage("a", log_probs=[0.1]) + assert ai["role"] == "assistant" and ai["log_probs"] == [0.1] + + +def test_base_message_content_is_copied(): + content = [{"type": "text", "text": "a"}] + msg = HumanMessage(content) + content.append({"type": "text", "text": "b"}) + assert len(msg["content"]) == 1 + + +def test_base_message_str_with_images(caplog): + msg = HumanMessage("hello") + msg.add_image("data:image/png;base64,abc") + msg["content"].append({"type": "input_text", "text": "world"}) + with caplog.at_level(logging.INFO): + assert str(msg) == "hello\nworld" + assert "images" in caplog.text + assert msg.__str__(warn_if_image=True) == "hello\nworld" + + +def test_add_image_from_pil_with_detail(): + msg = HumanMessage("look") + msg.add_image(Image.new("RGB", (2, 2)), detail="low") + image_elem = msg["content"][-1] + assert image_elem["type"] == "image_url" + assert image_elem["image_url"]["detail"] == "low" + assert image_elem["image_url"]["url"].startswith("data:image/jpeg;base64,") + + +def test_base_message_to_markdown(): + assert HumanMessage("x").to_markdown() == "\n```\nx\n```\n" + msg = HumanMessage("text") + msg.add_image("http://img") + msg.add_content("image_url", "http://img2") + md = msg.to_markdown() + assert "```\ntext\n```" in md + assert "![image](http://img)" in md + assert "![image](http://img2)" in md + + +def test_base_message_merge_on_string_is_noop(): + msg = HumanMessage("x") + msg.merge() + assert msg["content"] == "x" + + +def test_discussion_api(): + d = Discussion(SystemMessage("sys")) + d.add_message(role="user", content="hi") + d.add_text("there") + d.add_content("text", "again") + d.add_image("http://img") + assert len(d) == 2 + assert d.last_message is d[-1] + assert [m["role"] for m in d] == ["system", "user"] + + md = d.to_markdown() + assert md.startswith("Message 0") + assert "Message 1" in md + assert "![image](http://img)" in md + + openai_msgs = d.to_openai() + assert openai_msgs[1]["content"][0] == {"type": "text", "text": "hi\nthere\nagain"} + assert "sys" in d.to_string() + + +def test_discussion_default_empty(): + d = Discussion() + assert len(d) == 0 + d.append({"role": "user", "content": "x"}) + assert isinstance(d[0], BaseMessage) diff --git a/tests/llm/test_prompt_templates.py b/tests/llm/test_prompt_templates.py new file mode 100644 index 000000000..9563d838b --- /dev/null +++ b/tests/llm/test_prompt_templates.py @@ -0,0 +1,50 @@ +import pytest + +from agentlab.llm.prompt_templates import ( + STARCHAT_PROMPT_TEMPLATE, + PromptTemplate, + get_prompt_template, +) + + +def test_format_message_per_role(): + t = PromptTemplate(system="S:{input}|", human="H:{input}|", ai="A:{input}|") + assert t.format_message({"role": "system", "content": "a"}) == "S:a|" + assert t.format_message({"role": "user", "content": "b"}) == "H:b|" + assert t.format_message({"role": "assistant", "content": "c"}) == "A:c|" + + +def test_format_message_unknown_role(): + with pytest.raises(ValueError, match="not supported"): + STARCHAT_PROMPT_TEMPLATE.format_message({"role": "tool", "content": "x"}) + + +def test_construct_prompt_starchat(): + messages = [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + ] + prompt = STARCHAT_PROMPT_TEMPLATE.construct_prompt(messages) + assert prompt == ( + "<|system|>\nsys<|end|>\n" + "<|user|>\nhi<|end|>\n" + "<|assistant|>\nhello<|end|>\n" + "<|assistant|>" + ) + + +@pytest.mark.parametrize("bad", [["not a dict"], [{"role": "user"}], [{"content": "x"}]]) +def test_construct_prompt_rejects_non_openai_messages(bad): + with pytest.raises(ValueError, match="openai format"): + STARCHAT_PROMPT_TEMPLATE.construct_prompt(bad) + + +@pytest.mark.parametrize("name", ["bigcode/starcoder", "HuggingFaceH4/starchat-beta"]) +def test_get_prompt_template_known(name): + assert get_prompt_template(name) is STARCHAT_PROMPT_TEMPLATE + + +def test_get_prompt_template_unknown(): + with pytest.raises(NotImplementedError): + get_prompt_template("meta-llama/Llama-3") From edf138aea265ce7d764b084f71c80df918734651 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 23 Oct 2026 10:24:26 +0000 Subject: [PATCH 2/2] bug: set fake API keys in chat_api unit tests so they pass without OPENAI_API_KEY Co-Authored-By: yannis --- tests/llm/test_chat_api_unit.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/llm/test_chat_api_unit.py b/tests/llm/test_chat_api_unit.py index 36aa27070..dc6c28afd 100644 --- a/tests/llm/test_chat_api_unit.py +++ b/tests/llm/test_chat_api_unit.py @@ -32,6 +32,12 @@ def fresh_tracker(monkeypatch): monkeypatch.setattr(chat_api.tracking, "TRACKER", threading.local()) +@pytest.fixture(autouse=True) +def fake_api_keys(monkeypatch): + for var in ("OPENAI_API_KEY", "ANTHROPIC_API_KEY", "OPENROUTER_API_KEY"): + monkeypatch.setenv(var, "test-key") + + @pytest.fixture def no_sleep(monkeypatch): sleeps = []