Finalize multiagent tests

This commit is contained in:
Aymeric 2025-01-06 19:33:50 +01:00
parent ff8e20d93f
commit e8617b3037
1 changed files with 14 additions and 14 deletions

View File

@ -369,21 +369,21 @@ class AgentTests(unittest.TestCase):
): ):
if len(messages) < 3: if len(messages) < 3:
return """ return """
Thought: Let's call our search agent. Thought: Let's call our search agent.
Code: Code:
```py ```py
result = search_agent("Who is the current US president?") result = search_agent("Who is the current US president?")
```<end_code> ```<end_code>
""" """
else: else:
assert "Report on the current US president" in str(messages) assert "Report on the current US president" in str(messages)
return """ return """
Thought: Let's return the report. Thought: Let's return the report.
Code: Code:
```py ```py
final_answer("Final report.") final_answer("Final report.")
```<end_code> ```<end_code>
""" """
def get_tool_call( def get_tool_call(
self, messages, available_tools, stop_sequences=None, grammar=None self, messages, available_tools, stop_sequences=None, grammar=None
@ -398,7 +398,7 @@ class AgentTests(unittest.TestCase):
assert "Report on the current US president" in str(messages) assert "Report on the current US president" in str(messages)
return ( return (
"final_answer", "final_answer",
{"prompt": "Final report."}, "Final report.",
"call_0", "call_0",
) )
manager_model = FakeModelMultiagentsManagerAgent() manager_model = FakeModelMultiagentsManagerAgent()
@ -409,7 +409,7 @@ class AgentTests(unittest.TestCase):
): ):
return ( return (
"final_answer", "final_answer",
{"prompt": "Report on the current US president"}, {"report": "Report on the current US president"},
"call_0", "call_0",
) )
managed_model = FakeModelMultiagentsManagedAgent() managed_model = FakeModelMultiagentsManagedAgent()