289 lines
9.9 KiB
Python
289 lines
9.9 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Unit test for model response parser."""
|
|
import unittest
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
from agentscope.models import ModelResponse
|
|
from agentscope.parsers import (
|
|
MarkdownJsonDictParser,
|
|
MarkdownJsonObjectParser,
|
|
MarkdownCodeBlockParser,
|
|
MultiTaggedContentParser,
|
|
TaggedContent,
|
|
)
|
|
from agentscope.parsers.parser_base import DictFilterMixin
|
|
|
|
|
|
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 = (
|
|
"Respond a JSON dictionary in a markdown's 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.instruction_dict_3 = (
|
|
"Respond a JSON dictionary in a markdown's fenced code block as "
|
|
"follows:\n"
|
|
"```json\n"
|
|
"{a_JSON_dictionary}\n"
|
|
"```\n"
|
|
"The generated JSON dictionary MUST follow this schema: \n"
|
|
"{'properties': {'speak': {'description': 'what you speak', "
|
|
"'title': 'Speak', 'type': 'string'}, 'thought': {'description': "
|
|
"'what you thought', 'title': 'Thought', 'type': 'string'}, "
|
|
"'end_discussion': {'description': 'whether the discussion "
|
|
"reached an agreement or not', 'title': 'End Discussion', "
|
|
"'type': 'boolean'}}, 'required': ['speak', 'thought', "
|
|
"'end_discussion'], 'title': 'Schema', 'type': 'object'}"
|
|
)
|
|
|
|
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.instruction_code_with_hint = (
|
|
"You should generate python code in a python fenced code block as "
|
|
"follows: \n"
|
|
"```python\n"
|
|
"abc\n"
|
|
"```"
|
|
)
|
|
self.gt_code = """\nprint("Hello, world!")\n"""
|
|
|
|
def test_markdownjsondictparser_with_schema(self) -> None:
|
|
"""Test for MarkdownJsonDictParser with schema"""
|
|
|
|
class Schema(BaseModel): # pylint: disable=missing-class-docstring
|
|
speak: str = Field(description="what you speak")
|
|
thought: str = Field(description="what you thought")
|
|
end_discussion: bool = Field(
|
|
description="whether the discussion reached an agreement or "
|
|
"not",
|
|
)
|
|
|
|
parser = MarkdownJsonDictParser(
|
|
content_hint=Schema,
|
|
keys_to_memory=["speak", "thought"],
|
|
keys_to_content="speak",
|
|
keys_to_metadata=["end_discussion"],
|
|
)
|
|
|
|
self.assertEqual(parser.format_instruction, self.instruction_dict_3)
|
|
|
|
res = parser.parse(self.res_dict_1)
|
|
|
|
self.assertDictEqual(res.parsed, self.gt_dict)
|
|
|
|
res = parser.parse(
|
|
ModelResponse(
|
|
text="""```json
|
|
{
|
|
"speak" : "Hello, world!",
|
|
"thought" : "xxx",
|
|
"end_discussion" : "true"
|
|
}
|
|
```""",
|
|
),
|
|
)
|
|
|
|
self.assertDictEqual(res.parsed, self.gt_dict)
|
|
|
|
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_markdowncodeblockparser_with_hint(self) -> None:
|
|
"""Test for MarkdownCodeBlockParser"""
|
|
parser = MarkdownCodeBlockParser(
|
|
language_name="python",
|
|
content_hint="abc",
|
|
)
|
|
|
|
self.assertEqual(
|
|
parser.format_instruction,
|
|
self.instruction_code_with_hint,
|
|
)
|
|
|
|
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,
|
|
)
|
|
|
|
def test_DictFilterMixin_default_value(self) -> None:
|
|
"""Test the default value of the DictFilterMixin class"""
|
|
mixin = DictFilterMixin(
|
|
keys_to_memory=True,
|
|
keys_to_content=True,
|
|
keys_to_metadata=False,
|
|
)
|
|
|
|
self.assertDictEqual(mixin.to_memory(self.gt_dict), self.gt_dict)
|
|
self.assertDictEqual(mixin.to_content(self.gt_dict), self.gt_dict)
|
|
self.assertEqual(mixin.to_metadata(self.gt_dict), None)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|