From 7354d05b7ae7580c9a2e583cd3a0547c2cf49dd5 Mon Sep 17 00:00:00 2001 From: Whit Date: Sat, 9 May 2026 16:58:19 +0800 Subject: [PATCH] fix research agent tool loop --- contest_trade/agents/research_agent.py | 16 +++-- contest_trade/tools/tool_utils.py | 3 +- tests/test_research_agent_loop.py | 86 ++++++++++++++++++++++++++ 3 files changed, 99 insertions(+), 6 deletions(-) create mode 100644 tests/test_research_agent_loop.py diff --git a/contest_trade/agents/research_agent.py b/contest_trade/agents/research_agent.py index 3863960..dda3e69 100644 --- a/contest_trade/agents/research_agent.py +++ b/contest_trade/agents/research_agent.py @@ -217,7 +217,7 @@ async def _plan(self, state: ResearchAgentState) -> ResearchAgentState: async def _tool_selection(self, state: ResearchAgentState) -> ResearchAgentState: """选择工具""" - if not self.react: + if not self.react or state["tool_call_count"] >= self.config.max_react_step: state["selected_tool"] = {"tool_name": "final_report"} return state @@ -259,10 +259,11 @@ async def _enough_information(self, state: ResearchAgentState) -> str: return "enough_information" selected_tool = state["selected_tool"] + if state["tool_call_count"] >= self.config.max_react_step: + return "enough_information" if "error" in selected_tool: return "not_enough_information" - if selected_tool["tool_name"] == "final_report" or \ - state["tool_call_count"] >= self.config.max_react_step: + if selected_tool["tool_name"] == "final_report": return "enough_information" except Exception as e: logger.error(f"Error in enough_information: {e}") @@ -273,6 +274,14 @@ async def _call_tool(self, state: ResearchAgentState) -> ResearchAgentState: """调用工具""" selected_tool = state["selected_tool"] try: + if "error" in selected_tool: + state["tool_call_count"] += 1 + state["tool_call_context"] += json.dumps({ + "tool_called": selected_tool, + "tool_result": selected_tool, + }, ensure_ascii=False) + "\n" + return state + print('Begin to call tool: ', selected_tool) tool_name = selected_tool["tool_name"] tool_args = selected_tool["properties"] @@ -457,4 +466,3 @@ async def run_with_monitoring(self, input: ResearchAgentInput) -> ResearchAgentO ) agent_output = asyncio.run(agent.run_with_monitoring(agent_input)) print(agent_output.to_dict()) - diff --git a/contest_trade/tools/tool_utils.py b/contest_trade/tools/tool_utils.py index cac4f13..5ed1031 100644 --- a/contest_trade/tools/tool_utils.py +++ b/contest_trade/tools/tool_utils.py @@ -63,7 +63,7 @@ def register_from_module_path(self, module_path: str): module = importlib.import_module(module_name) func = getattr(module, func_name) - if not callable(func): + if not callable(func) and not hasattr(func, 'invoke'): raise ValueError(f"{module_path} is not callable") return self.register_function(func) @@ -341,4 +341,3 @@ async def print_string(input_string: str): # call tool result = asyncio.run(registry.call_tool("print_string", {"input_string": "Hello World!"})) print(result) - diff --git a/tests/test_research_agent_loop.py b/tests/test_research_agent_loop.py new file mode 100644 index 0000000..4cf6e74 --- /dev/null +++ b/tests/test_research_agent_loop.py @@ -0,0 +1,86 @@ +import asyncio +import json +import shutil +import sys +import unittest +from pathlib import Path + + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +CONTEST_TRADE_ROOT = PROJECT_ROOT / "contest_trade" +if str(CONTEST_TRADE_ROOT) not in sys.path: + sys.path.insert(0, str(CONTEST_TRADE_ROOT)) + +from agents.research_agent import ResearchAgent, ResearchAgentConfig +from tools.tool_utils import ToolManager, ToolManagerConfig + + +class ResearchAgentLoopTests(unittest.TestCase): + def setUp(self): + self.agent = ResearchAgent( + ResearchAgentConfig(agent_name="unit_loop_test", belief="test belief") + ) + + def tearDown(self): + shutil.rmtree(self.agent.signal_dir, ignore_errors=True) + + def _state(self, tool_call_count, selected_tool): + return { + "trigger_time": "2026-05-09 16:26:25", + "task": "test task", + "belief": "test belief", + "background_information": "", + "plan_result": "", + "tool_call_context": "", + "selected_tool": selected_tool, + "tool_call_count": tool_call_count, + "step_count": 0, + "final_result": "", + "final_result_thinking": "", + "result": None, + } + + def test_tool_selection_uses_final_report_after_max_steps(self): + state = self._state(self.agent.config.max_react_step, {}) + + asyncio.run(self.agent._tool_selection(state)) + + self.assertEqual({"tool_name": "final_report"}, state["selected_tool"]) + + def test_max_steps_takes_precedence_over_tool_selection_error(self): + state = self._state( + self.agent.config.max_react_step, + {"error": "Call tool Failed", "error_msg": "bad tool selection"}, + ) + + result = asyncio.run(self.agent._enough_information(state)) + + self.assertEqual("enough_information", result) + + def test_tool_selection_error_is_recorded_as_a_tool_step(self): + state = self._state( + 0, + {"error": "Call tool Failed", "error_msg": "bad tool selection"}, + ) + + asyncio.run(self.agent._call_tool(state)) + recorded_context = json.loads(state["tool_call_context"]) + + self.assertEqual(1, state["tool_call_count"]) + self.assertEqual(state["selected_tool"], recorded_context["tool_result"]) + + +class ToolManagerRegistrationTests(unittest.TestCase): + def test_registers_langchain_structured_tools(self): + manager = ToolManager( + ToolManagerConfig(["tools.final_report.final_report"]) + ) + + registered_tool = manager.get_tool("final_report") + + self.assertIsNotNone(registered_tool) + self.assertTrue(hasattr(registered_tool, "invoke")) + + +if __name__ == "__main__": + unittest.main()