From 604666ad997dd2d2afedcac01d4816f7eb6a0dd5 Mon Sep 17 00:00:00 2001 From: ZY Liu <236001964+liuzhengyang699@users.noreply.github.com> Date: Sun, 11 Oct 2026 12:03:12 +0800 Subject: [PATCH] fix(mars): match solution authors across string and integer IDs --- optillm/mars/verifier.py | 14 ++--- tests/test_mars_agent_routing.py | 88 ++++++++++++++++++++++++++++++++ 2 files changed, 96 insertions(+), 6 deletions(-) create mode 100644 tests/test_mars_agent_routing.py diff --git a/optillm/mars/verifier.py b/optillm/mars/verifier.py index 85b1bd17..816960bb 100644 --- a/optillm/mars/verifier.py +++ b/optillm/mars/verifier.py @@ -178,9 +178,11 @@ def _verify_single_solution(self, solution: AgentSolution, request_id: str = Non ] } - def _select_verifier_agent(self, solution_agent_id: int) -> MARSAgent: + def _select_verifier_agent(self, solution_agent_id: str | int) -> MARSAgent: """Select an agent different from the solution creator for verification""" - available_agents = [agent for agent in self.agents if agent.agent_id != solution_agent_id] + # Solutions store string IDs, while the reasoning agents use integers. + solution_agent_id = str(solution_agent_id) + available_agents = [agent for agent in self.agents if str(agent.agent_id) != solution_agent_id] if not available_agents: # If no other agents available, use any agent available_agents = self.agents @@ -188,7 +190,7 @@ def _select_verifier_agent(self, solution_agent_id: int) -> MARSAgent: # Prefer agents with different temperatures for diverse verification if len(available_agents) > 1: # Select agent with most different temperature - solution_agent = next((a for a in self.agents if a.agent_id == solution_agent_id), None) + solution_agent = next((a for a in self.agents if str(a.agent_id) == solution_agent_id), None) if solution_agent: solution_temp = solution_agent.temperature available_agents.sort(key=lambda a: abs(a.temperature - solution_temp), reverse=True) @@ -215,7 +217,7 @@ def iterative_improvement(self, request_id: str = None) -> Dict[str, Any]: if latest_verification['assessment'] in ['INCORRECT', 'INCOMPLETE']: # Find the original agent to improve their solution - original_agent = next((a for a in self.agents if a.agent_id == solution.agent_id), None) + original_agent = next((a for a in self.agents if str(a.agent_id) == str(solution.agent_id)), None) if original_agent: try: @@ -267,7 +269,7 @@ async def iterative_improvement_parallel( if solution.verification_results: latest_verification = solution.verification_results[-1] if latest_verification['assessment'] in ['INCORRECT', 'INCOMPLETE']: - original_agent = next((a for a in self.agents if a.agent_id == solution.agent_id), None) + original_agent = next((a for a in self.agents if str(a.agent_id) == str(solution.agent_id)), None) if original_agent: improvable_solutions.append((solution, original_agent, latest_verification)) @@ -340,4 +342,4 @@ def final_consensus_check(self) -> bool: else: logger.info(f"No consensus: only {len(verified_solutions)} solutions verified (need {consensus_threshold})") - return has_consensus \ No newline at end of file + return has_consensus diff --git a/tests/test_mars_agent_routing.py b/tests/test_mars_agent_routing.py new file mode 100644 index 00000000..04f73a92 --- /dev/null +++ b/tests/test_mars_agent_routing.py @@ -0,0 +1,88 @@ +"""Generated solutions must be verified by peers and refined by their author.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +from optillm.mars.agent import MARSAgent +from optillm.mars.verifier import MARSVerifier +from optillm.mars.workspace import MARSWorkspace + + +def response(text, reasoning_tokens=7): + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content=text))], + usage=SimpleNamespace(total_tokens=20, reasoning_tokens=reasoning_tokens), + ) + + +@pytest.mark.parametrize("parallel", [False, True]) +@pytest.mark.parametrize("author_id,verifier_id", [(0, 2), (1, 2), (2, 0)]) +def test_generated_solution_is_verified_and_improved_by_the_right_agents( + parallel, author_id, verifier_id +): + config = {"max_tokens": 100, "max_verification_attempts": 1} + workspace = MARSWorkspace("Solve 2x + 5 = 15.", config) + feedback = "INCORRECT: the subtraction is wrong. Subtract 5 before dividing." + agents = [MARSAgent(i, Mock(), "test-model", config) for i in range(3)] + for agent in agents: + agent.client.chat.completions.create.return_value = response(feedback) + + author = agents[author_id] + author.client.chat.completions.create.side_effect = [ + response("The answer is 10."), + response("The answer is 5.", reasoning_tokens=11), + ] + solution, _ = author.generate_solution(workspace.problem) + assert solution.agent_id == str(author_id) + workspace.add_solution(solution) + verifier = MARSVerifier(agents, workspace, config) + + if parallel: + verification = asyncio.run(verifier.verify_solutions_parallel()) + else: + verification = verifier.verify_solutions() + + detail = verification["verification_details"][0] + assert detail["verification_results"][0]["verifier_id"] == verifier_id + assert solution.verification_results[-1]["assessment"] == "INCORRECT" + + if parallel: + improvement = asyncio.run(verifier.iterative_improvement_parallel()) + else: + improvement = verifier.iterative_improvement() + + assert improvement == { + "solutions_improved": 1, + "improvement_attempts": 1, + "total_reasoning_tokens": 11, + } + assert solution.solution == "The answer is 5." + assert solution.reasoning_tokens == 18 + assert author.client.chat.completions.create.call_count == 2 + prompt = author.client.chat.completions.create.call_args.kwargs["messages"][1]["content"] + assert "The answer is 10." in prompt + assert feedback in prompt + for agent in agents: + if agent.agent_id not in (author_id, verifier_id): + agent.client.chat.completions.create.assert_not_called() + + +@pytest.mark.parametrize("author_id,verifier_id", [(0, 2), (1, 2), (2, 0)]) +def test_integer_ids_still_choose_a_peer_with_a_different_temperature(author_id, verifier_id): + agents = [MARSAgent(i, Mock(), "test-model", {}) for i in range(3)] + verifier = MARSVerifier(agents, MARSWorkspace("problem", {}), {}) + + assert verifier._select_verifier_agent(author_id) is agents[verifier_id] + + +def test_aggregated_and_single_agent_solutions_keep_the_fallback(): + agent = MARSAgent(0, Mock(), "test-model", {}) + workspace = MARSWorkspace("problem", {}) + verifier = MARSVerifier([agent], workspace, {}) + + assert verifier._select_verifier_agent("0") is agent + assert verifier._select_verifier_agent("agg_123456") is agent + assert MARSVerifier([], workspace, {})._select_verifier_agent("0") is None