| 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"']) |
|
|
