fix(mars): route verification and refinement using matching agent IDs

optillm/mars/verifier.py

@@ -178,17 +178,19 @@ 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
 
         # 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
+        return has_consensus

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