381 lines
12 KiB
Python
381 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
|
""" Unit test for service toolkit. """
|
|
import json
|
|
import unittest
|
|
from typing import Literal
|
|
|
|
from agentscope.models import ModelWrapperBase, ModelResponse
|
|
from agentscope.parsers import MultiTaggedContentParser, TaggedContent
|
|
from agentscope.service import (
|
|
bing_search,
|
|
execute_python_code,
|
|
retrieve_from_list,
|
|
query_mysql,
|
|
summarization,
|
|
)
|
|
from agentscope.service import ServiceToolkit
|
|
|
|
|
|
class ServiceToolkitTest(unittest.TestCase):
|
|
"""
|
|
Unit test for service toolkit.
|
|
"""
|
|
|
|
def setUp(self) -> None:
|
|
"""Init for ExampleTest."""
|
|
self.json_schema_bing_search1 = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "bing_search",
|
|
"description": (
|
|
"Search question in Bing Search API and "
|
|
"return the searching results"
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"num_results": {
|
|
"type": "number",
|
|
"description": (
|
|
"The number of search " "results to return."
|
|
),
|
|
"default": 10,
|
|
},
|
|
"question": {
|
|
"type": "string",
|
|
"description": "The search query string.",
|
|
},
|
|
},
|
|
"required": [
|
|
"question",
|
|
],
|
|
},
|
|
},
|
|
}
|
|
|
|
self.json_schema_bing_search2 = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "bing_search",
|
|
"description": (
|
|
"Search question in Bing Search API and "
|
|
"return the searching results"
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"question": {
|
|
"type": "string",
|
|
"description": "The search query string.",
|
|
},
|
|
},
|
|
"required": ["question"],
|
|
},
|
|
},
|
|
}
|
|
|
|
self.json_schema_func = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "func",
|
|
"description": None,
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"c": {"default": "test"},
|
|
"d": {
|
|
"type": "typing.Literal",
|
|
"enum": [1, "abc", "d"],
|
|
"default": 1,
|
|
},
|
|
"b": {},
|
|
"a": {"type": "string"},
|
|
},
|
|
"required": ["a", "b"],
|
|
},
|
|
},
|
|
}
|
|
|
|
self.json_schema_execute_python_code = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "execute_python_code",
|
|
"description": "Execute a piece of python code.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"code": {
|
|
"type": "string",
|
|
"description": (
|
|
"The Python code to be " "executed."
|
|
),
|
|
},
|
|
},
|
|
"required": ["code"],
|
|
},
|
|
},
|
|
}
|
|
|
|
self.json_schema_retrieve_from_list = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "retrieve_from_list",
|
|
"description": "Retrieve data in a list.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {
|
|
"type": "Any",
|
|
"description": "A message to be retrieved.",
|
|
},
|
|
},
|
|
"required": [
|
|
"query",
|
|
],
|
|
},
|
|
},
|
|
}
|
|
|
|
self.json_schema_query_mysql = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "query_mysql",
|
|
"description": "Execute query within MySQL database.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"query": {
|
|
"type": "string",
|
|
"description": "SQL query to execute.",
|
|
},
|
|
},
|
|
"required": [
|
|
"query",
|
|
],
|
|
},
|
|
},
|
|
}
|
|
|
|
self.json_schema_summarization = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "summarization",
|
|
"description": "Summarize the input text.",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"text": {
|
|
"type": "string",
|
|
"description": (
|
|
"Text to be summarized by " "the model."
|
|
),
|
|
},
|
|
},
|
|
"required": [
|
|
"text",
|
|
],
|
|
},
|
|
},
|
|
}
|
|
|
|
def test_bing_search(self) -> None:
|
|
"""Test bing_search."""
|
|
# api_key is specified by developer, while question and num_results
|
|
# are specified by model
|
|
_, doc_dict = ServiceToolkit.get(bing_search, api_key="xxx")
|
|
print(json.dumps(doc_dict, indent=4))
|
|
self.assertDictEqual(
|
|
doc_dict,
|
|
self.json_schema_bing_search1,
|
|
)
|
|
|
|
# Set num_results by developer rather than model
|
|
_, doc_dict = ServiceToolkit.get(
|
|
bing_search,
|
|
num_results=3,
|
|
api_key="xxx",
|
|
)
|
|
|
|
self.assertDictEqual(
|
|
doc_dict,
|
|
self.json_schema_bing_search2,
|
|
)
|
|
|
|
def test_enum(self) -> None:
|
|
"""Test enum in service toolkit."""
|
|
|
|
def func( # type: ignore
|
|
a: str,
|
|
b,
|
|
c="test",
|
|
d: Literal[1, "abc", "d"] = 1,
|
|
) -> None:
|
|
print(a, b, c, d)
|
|
|
|
_, doc_dict = ServiceToolkit.get(func)
|
|
|
|
self.assertDictEqual(
|
|
doc_dict,
|
|
self.json_schema_func,
|
|
)
|
|
|
|
def test_exec_python_code(self) -> None:
|
|
"""Test execute_python_code in service toolkit."""
|
|
_, doc_dict = ServiceToolkit.get(
|
|
execute_python_code,
|
|
timeout=300,
|
|
use_docker=True,
|
|
maximum_memory_bytes=None,
|
|
)
|
|
|
|
self.assertDictEqual(
|
|
doc_dict,
|
|
self.json_schema_execute_python_code,
|
|
)
|
|
|
|
def test_retrieval(self) -> None:
|
|
"""Test retrieval in service toolkit."""
|
|
_, doc_dict = ServiceToolkit.get(
|
|
retrieve_from_list,
|
|
knowledge=[1, 2, 3],
|
|
score_func=lambda x, y: 1.0,
|
|
top_k=10,
|
|
embedding_model=10,
|
|
preserve_order=True,
|
|
)
|
|
|
|
self.assertDictEqual(
|
|
doc_dict,
|
|
self.json_schema_retrieve_from_list,
|
|
)
|
|
|
|
def test_sql_query(self) -> None:
|
|
"""Test sql_query in service toolkit."""
|
|
_, doc_dict = ServiceToolkit.get(
|
|
query_mysql,
|
|
database="test",
|
|
host="localhost",
|
|
user="root",
|
|
password="xxx",
|
|
port=3306,
|
|
allow_change_data=False,
|
|
maxcount_results=None,
|
|
)
|
|
|
|
self.assertDictEqual(
|
|
doc_dict,
|
|
self.json_schema_query_mysql,
|
|
)
|
|
|
|
def test_summary(self) -> None:
|
|
"""Test summarization in service toolkit."""
|
|
_, doc_dict = ServiceToolkit.get(
|
|
summarization,
|
|
model=ModelWrapperBase("abc"),
|
|
system_prompt="",
|
|
summarization_prompt="",
|
|
max_return_token=-1,
|
|
token_limit_prompt="",
|
|
)
|
|
|
|
print(json.dumps(doc_dict, indent=4))
|
|
|
|
self.assertDictEqual(
|
|
doc_dict,
|
|
self.json_schema_summarization,
|
|
)
|
|
|
|
def test_service_toolkit(self) -> None:
|
|
"""Test the object of ServiceToolkit."""
|
|
service_toolkit = ServiceToolkit()
|
|
|
|
service_toolkit.add(bing_search, api_key="xxx", num_results=3)
|
|
service_toolkit.add(
|
|
execute_python_code,
|
|
timeout=300,
|
|
use_docker=True,
|
|
maximum_memory_bytes=None,
|
|
)
|
|
|
|
self.assertEqual(
|
|
service_toolkit.tools_instruction,
|
|
"""## Tool Functions:
|
|
The following tool functions are available in the format of
|
|
```
|
|
{index}. {function name}: {function description}
|
|
{argument1 name} ({argument type}): {argument description}
|
|
{argument2 name} ({argument type}): {argument description}
|
|
...
|
|
```
|
|
|
|
1. bing_search: Search question in Bing Search API and return the searching results
|
|
question (string): The search query string.
|
|
2. execute_python_code: Execute a piece of python code.
|
|
code (string): The Python code to be executed.
|
|
""", # noqa
|
|
)
|
|
self.assertDictEqual(
|
|
service_toolkit.json_schemas,
|
|
{
|
|
"bing_search": self.json_schema_bing_search2,
|
|
"execute_python_code": self.json_schema_execute_python_code,
|
|
},
|
|
)
|
|
|
|
def test_multi_tagged_content(self) -> None:
|
|
"""Test multi tagged content"""
|
|
|
|
parser = MultiTaggedContentParser(
|
|
TaggedContent(
|
|
"thought",
|
|
"[THOUGHT]",
|
|
"what you think",
|
|
"[/THOUGHT]",
|
|
),
|
|
TaggedContent("speak", "[SPEAK]", "what you speak", "[/SPEAK]"),
|
|
TaggedContent(
|
|
"function",
|
|
"[FUNCTION]",
|
|
'{"function name": "xxx", "args": {"arg name": "xxx"}}',
|
|
"[/FUNCTION]",
|
|
parse_json=True,
|
|
),
|
|
)
|
|
|
|
target_format_instruction = (
|
|
"Respond with specific tags as outlined below, and the content "
|
|
"between [FUNCTION] and [/FUNCTION] MUST be a JSON object:\n"
|
|
"[THOUGHT]what you think[/THOUGHT]\n"
|
|
"[SPEAK]what you speak[/SPEAK]\n"
|
|
'[FUNCTION]{"function name": "xxx", "args": {"arg name": "xxx"}}'
|
|
"[/FUNCTION]"
|
|
)
|
|
|
|
self.assertEqual(parser.format_instruction, target_format_instruction)
|
|
|
|
response = ModelResponse(
|
|
text=(
|
|
"This is a test\n[THOUGHT]I think this is a good idea"
|
|
"[/THOUGHT]\n[SPEAK]I am speaking now[/SPEAK]\n[FUNCTION]"
|
|
'{"function name": "bing_search", "args": {"query": "news of '
|
|
'today"}}[/FUNCTION]'
|
|
),
|
|
)
|
|
res = parser.parse(response)
|
|
|
|
self.assertDictEqual(
|
|
res.parsed,
|
|
{
|
|
"thought": "I think this is a good idea",
|
|
"speak": "I am speaking now",
|
|
"function": {
|
|
"function name": "bing_search",
|
|
"args": {"query": "news of today"},
|
|
},
|
|
},
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|