|
from unittest import TestCase, mock |
|
|
|
from lagent.actions import ActionExecutor |
|
from lagent.actions.llm_qa import LLMQA |
|
from lagent.actions.serper_search import SerperSearch |
|
from lagent.agents.rewoo import ReWOO, ReWOOProtocol |
|
from lagent.schema import ActionReturn, ActionStatusCode |
|
|
|
|
|
class TestReWOO(TestCase): |
|
|
|
@mock.patch.object(SerperSearch, 'run') |
|
@mock.patch.object(LLMQA, 'run') |
|
@mock.patch.object(ReWOOProtocol, 'parse_worker') |
|
def test_normal_chat(self, mock_parse_worker_func, mock_qa_func, |
|
mock_search_func): |
|
mock_model = mock.Mock() |
|
mock_model.generate_from_template.return_value = 'LLM response' |
|
|
|
mock_parse_worker_func.return_value = (['Thought1', 'Thought2' |
|
], ['LLMQA', 'SerperSearch'], |
|
['abc', 'abc']) |
|
|
|
search_return = ActionReturn(args=None) |
|
search_return.state = ActionStatusCode.SUCCESS |
|
search_return.result = dict(text='search_return') |
|
mock_search_func.return_value = search_return |
|
|
|
qa_return = ActionReturn(args=None) |
|
qa_return.state = ActionStatusCode.SUCCESS |
|
qa_return.result = dict(text='qa_return') |
|
mock_qa_func.return_value = qa_return |
|
|
|
chatbot = ReWOO( |
|
llm=mock_model, |
|
action_executor=ActionExecutor(actions=[ |
|
LLMQA(mock_model), |
|
SerperSearch(api_key=''), |
|
])) |
|
agent_return = chatbot.chat('abc') |
|
self.assertEqual(agent_return.response, 'LLM response') |
|
|
|
def test_parse_worker(self): |
|
prompt = ReWOOProtocol() |
|
message = """ |
|
Plan: a. |
|
#E1 = tool1["a"] |
|
#E2 = tool2["b"] |
|
""" |
|
try: |
|
thoughts, actions, actions_input = prompt.parse_worker(message) |
|
except Exception as e: |
|
self.assertEqual( |
|
'Each Plan should only correspond to only ONE action', str(e)) |
|
else: |
|
self.assertFalse( |
|
True, 'it should raise exception when the format is incorrect') |
|
|
|
message = """ |
|
Plan: a. |
|
#E1 = tool1("a") |
|
Plan: b. |
|
#E2 = tool2["b"] |
|
""" |
|
try: |
|
thoughts, actions, actions_input = prompt.parse_worker(message) |
|
except Exception as e: |
|
self.assertIsInstance(e, BaseException) |
|
else: |
|
self.assertFalse( |
|
True, 'it should raise exception when the format is incorrect') |
|
|
|
message = """ |
|
Plan: a. |
|
#E1 = tool1["a"] |
|
Plan: b. |
|
#E2 = tool2["b"] |
|
""" |
|
try: |
|
thoughts, actions, actions_input = prompt.parse_worker(message) |
|
except Exception: |
|
self.assertFalse( |
|
True, |
|
'it should not raise exception when the format is correct') |
|
self.assertEqual(thoughts, ['a.', 'b.']) |
|
self.assertEqual(actions, ['tool1', 'tool2']) |
|
self.assertEqual(actions_input, ['"a"', '"b"']) |
|
|
