Add unittest for parsers.

tests/parser_text.py

@@ -0,0 +1,196 @@
+# -*- coding: utf-8 -*-
+"""Unit test for model response parser."""
+import unittest
+
+from agentscope.models import ModelResponse
+from agentscope.parsers import (
+    MarkdownJsonDictParser,
+    MarkdownJsonObjectParser,
+    MarkdownCodeBlockParser,
+    MultiTaggedContentParser,
+    TaggedContent,
+)
+
+
+class ModelResponseParserTest(unittest.TestCase):
+    """Unit test for model response parser."""
+
+    def setUp(self) -> None:
+        """Init for ExampleTest."""
+        self.res_dict_1 = ModelResponse(
+            text=(
+                "```json\n"
+                '{"speak": "Hello, world!", '
+                '"thought": "xxx", '
+                '"end_discussion": true}\n```'
+            ),
+        )
+        self.instruction_dict_1 = (
+            "You should respond a json object in a json fenced code block "
+            "as follows:\n"
+            "```json\n"
+            '{"speak": "what you speak", '
+            '"thought": "what you thought", '
+            '"end_discussion": true/false}\n'
+            "```"
+        )
+        self.res_dict_2 = ModelResponse(
+            text="[SPEAK]Hello, world![/SPEAK]\n"
+            "[THOUGHT]xxx[/THOUGHT]\n"
+            "[END_DISCUSSION]true[/END_DISCUSSION]",
+        )
+        self.instruction_dict_2 = (
+            "Respond with specific tags as outlined below, and the content "
+            "between [END_DISCUSSION] and [/END_DISCUSSION] MUST be a JSON "
+            "object:\n"
+            "[SPEAK]what you speak[/SPEAK]\n"
+            "[THOUGHT]what you thought[/THOUGHT]\n"
+            "[END_DISCUSSION]true/false[/END_DISCUSSION]"
+        )
+        self.gt_dict = {
+            "speak": "Hello, world!",
+            "thought": "xxx",
+            "end_discussion": True,
+        }
+        self.hint_dict = (
+            '{"speak": "what you speak", '
+            '"thought": "what you thought", '
+            '"end_discussion": true/false}'
+        )
+
+        self.gt_to_memory = {"speak": "Hello, world!", "thought": "xxx"}
+        self.gt_to_content = "Hello, world!"
+        self.gt_to_metadata = {"end_discussion": True}
+
+        self.res_list = ModelResponse(text="""```json\n[1,2,3]\n```""")
+        self.instruction_list = (
+            "You should respond a json object in a json fenced code block as "
+            "follows:\n"
+            "```json\n"
+            "{Your generated list of numbers}\n"
+            "```"
+        )
+        self.gt_list = [1, 2, 3]
+        self.hint_list = "{Your generated list of numbers}"
+
+        self.res_float = ModelResponse(text="""```json\n3.14\n```""")
+        self.instruction_float = (
+            "You should respond a json object in a json fenced code block as "
+            "follows:\n"
+            "```json\n"
+            "{Your generated float number}\n"
+            "```"
+        )
+        self.gt_float = 3.14
+        self.hint_float = "{Your generated float number}"
+
+        self.res_code = ModelResponse(
+            text="""```python\nprint("Hello, world!")\n```""",
+        )
+        self.instruction_code = (
+            "You should generate python code in a python fenced code block as "
+            "follows: \n"
+            "```python\n"
+            "${your_python_code}\n"
+            "```"
+        )
+        self.gt_code = """\nprint("Hello, world!")\n"""
+
+    def test_markdownjsondictparser(self) -> None:
+        """Test for MarkdownJsonDictParser"""
+        parser = MarkdownJsonDictParser(
+            content_hint=self.hint_dict,
+            keys_to_memory=["speak", "thought"],
+            keys_to_content="speak",
+            keys_to_metadata=["end_discussion"],
+        )
+
+        self.assertEqual(parser.format_instruction, self.instruction_dict_1)
+
+        res = parser.parse(self.res_dict_1)
+
+        self.assertDictEqual(res.parsed, self.gt_dict)
+
+        # test filter functions
+        self.assertDictEqual(parser.to_memory(res.parsed), self.gt_to_memory)
+        self.assertEqual(parser.to_content(res.parsed), self.gt_to_content)
+        self.assertDictEqual(
+            parser.to_metadata(res.parsed),
+            self.gt_to_metadata,
+        )
+
+    def test_markdownjsonobjectparser(self) -> None:
+        """Test for MarkdownJsonObjectParser"""
+        # list
+        parser_list = MarkdownJsonObjectParser(content_hint=self.hint_list)
+
+        self.assertEqual(parser_list.format_instruction, self.instruction_list)
+
+        res_list = parser_list.parse(self.res_list)
+        self.assertListEqual(res_list.parsed, self.gt_list)
+
+        # float
+        parser_float = MarkdownJsonObjectParser(content_hint=self.hint_float)
+
+        self.assertEqual(
+            parser_float.format_instruction,
+            self.instruction_float,
+        )
+
+        res_float = parser_float.parse(self.res_float)
+        self.assertEqual(res_float.parsed, self.gt_float)
+
+    def test_markdowncodeblockparser(self) -> None:
+        """Test for MarkdownCodeBlockParser"""
+        parser = MarkdownCodeBlockParser(language_name="python")
+
+        self.assertEqual(parser.format_instruction, self.instruction_code)
+
+        res = parser.parse(self.res_code)
+
+        self.assertEqual(res.parsed, self.gt_code)
+
+    def test_multitaggedcontentparser(self) -> None:
+        """Test for MultiTaggedContentParser"""
+        parser = MultiTaggedContentParser(
+            TaggedContent(
+                "speak",
+                tag_begin="[SPEAK]",
+                content_hint="what you speak",
+                tag_end="[/SPEAK]",
+            ),
+            TaggedContent(
+                "thought",
+                tag_begin="[THOUGHT]",
+                content_hint="what you thought",
+                tag_end="[/THOUGHT]",
+            ),
+            TaggedContent(
+                "end_discussion",
+                tag_begin="[END_DISCUSSION]",
+                content_hint="true/false",
+                tag_end="[/END_DISCUSSION]",
+                parse_json=True,
+            ),
+            keys_to_memory=["speak", "thought"],
+            keys_to_content="speak",
+            keys_to_metadata=["end_discussion"],
+        )
+
+        self.assertEqual(parser.format_instruction, self.instruction_dict_2)
+
+        res = parser.parse(self.res_dict_2)
+
+        self.assertDictEqual(res.parsed, self.gt_dict)
+
+        # test filter functions
+        self.assertDictEqual(parser.to_memory(res.parsed), self.gt_to_memory)
+        self.assertEqual(parser.to_content(res.parsed), self.gt_to_content)
+        self.assertDictEqual(
+            parser.to_metadata(res.parsed),
+            self.gt_to_metadata,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()