init
This commit is contained in:
CakeCN
380
AlgoriAgent/tests/service_toolkit_test.py
Normal file
380
AlgoriAgent/tests/service_toolkit_test.py
Normal file
@@ -0,0 +1,380 @@
|
||||
# -*- 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()
|
||||
Reference in New Issue
Block a user