diff --git a/tests/agents/test_reproducibility_agent.py b/tests/agents/test_reproducibility_agent.py new file mode 100644 index 000000000..f8c18a145 --- /dev/null +++ b/tests/agents/test_reproducibility_agent.py @@ -0,0 +1,124 @@ +"""Unit tests for the trace replay logic of ReproAgent and ReproChatModel.""" + +from dataclasses import dataclass, field + +import pytest +from browsergym.experiments.agent import AgentInfo + +from agentlab.agents.generic_agent import reproducibility_agent +from agentlab.agents.generic_agent.agent_configs import FLAGS_GPT_3_5 +from agentlab.agents.generic_agent.generic_agent import GenericAgent +from agentlab.agents.generic_agent.reproducibility_agent import ReproAgent, ReproChatModel +from agentlab.llm.chat_api import CheatMiniWoBLLMArgs +from agentlab.llm.llm_utils import AIMessage, Discussion, HumanMessage, SystemMessage + + +@dataclass +class StubStepInfo: + agent_info: dict = field(default_factory=dict) + action: str = None + + +@dataclass +class StubExpResult: + step_info: StubStepInfo + summary_info: dict = field(default_factory=lambda: {"err_msg": "some error"}) + + def get_step_info(self, step: int) -> StubStepInfo: + return self.step_info + + +def make_agent(monkeypatch, step_info: StubStepInfo) -> ReproAgent: + monkeypatch.setattr( + reproducibility_agent, "ExpResult", lambda repro_dir: StubExpResult(step_info) + ) + return ReproAgent( + chat_model_args=CheatMiniWoBLLMArgs(), + flags=FLAGS_GPT_3_5, + repro_dir="unused", + ) + + +def patch_parent_get_action(monkeypatch): + """Replay a 2-message discussion through self.chat_llm, like GenericAgent does.""" + + def fake_get_action(self, obs): + messages = Discussion([SystemMessage("system"), HumanMessage("user")]) + response = self.chat_llm(messages) + self.actions.append(response["content"]) + return response["content"], AgentInfo(chat_messages=messages) + + monkeypatch.setattr(GenericAgent, "get_action", fake_get_action) + + +def recorded_discussion(answer: str = None) -> Discussion: + messages = [SystemMessage("system"), HumanMessage("user")] + if answer is not None: + messages.append(AIMessage(answer)) + return Discussion(messages) + + +def test_get_action_without_chat_messages(monkeypatch): + agent = make_agent(monkeypatch, StubStepInfo(agent_info={"chat_messages": None})) + + action, agent_info = agent.get_action(obs={}) + + assert action is None + assert "Agent had no chat messages" in agent_info.markdown_page + assert "some error" in agent_info.markdown_page + + +def test_get_action_rebuilds_missing_assistant_message(monkeypatch): + old_messages = recorded_discussion() + step_info = StubStepInfo(agent_info={"chat_messages": old_messages}, action='click("42")') + agent = make_agent(monkeypatch, step_info) + patch_parent_get_action(monkeypatch) + + action, _ = agent.get_action(obs={}) + + assert len(old_messages) == 3 + assert old_messages[2]["role"] == "assistant" + assert old_messages[2]["content"] == 'click("42")' + assert isinstance(agent.chat_llm, ReproChatModel) + assert action == 'click("42")' + + +def test_get_action_without_recorded_action_keeps_messages(monkeypatch): + old_messages = recorded_discussion() + step_info = StubStepInfo(agent_info={"chat_messages": old_messages}, action=None) + agent = make_agent(monkeypatch, step_info) + patch_parent_get_action(monkeypatch) + + action, _ = agent.get_action(obs={}) + + assert len(old_messages) == 2 + assert action == "None" # fallback of ReproChatModel + + +def test_get_action_replays_recorded_answer(monkeypatch): + old_messages = recorded_discussion('click("7")') + step_info = StubStepInfo(agent_info={"chat_messages": old_messages}, action='click("7")') + agent = make_agent(monkeypatch, step_info) + patch_parent_get_action(monkeypatch) + + action, agent_info = agent.get_action(obs={}) + + assert len(old_messages) == 3 # no reconstruction when the answer was recorded + assert action == 'click("7")' + assert agent_info.stats["lines_added"] == 0 + assert agent_info.stats["lines_removed"] == 0 + + +@pytest.mark.parametrize("n_new_messages", [2, 3]) +def test_repro_chat_model(n_new_messages): + old_messages = recorded_discussion("noop()") + chat_model = ReproChatModel(old_messages, delay=0) + + response = chat_model(Discussion(list(old_messages)[:n_new_messages])) + + if n_new_messages == 2: + assert response is old_messages[2] + assert len(chat_model.new_messages) == 3 + else: # the recorded answer is missing, fall back to a no-op action + assert response["content"] == "None" + assert chat_model.get_stats() == {}