erniebot-agent/tests/unit_tests/tools/test_schema.py

# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations

import unittest
from typing import List, Optional

from erniebot_agent.tools.base import RemoteToolkit
from erniebot_agent.tools.schema import (
    ToolParameterView,
    get_typing_list_type,
    is_optional_type,
    json_type,
)
from openapi_spec_validator.readers import read_from_filename
from pydantic import Field


class TestToolSchema(unittest.TestCase):
    openapi_file = "./tests/fixtures/openapi.yaml"

    def test_plugin_schema(self):
        schema = RemoteToolkit.from_openapi_file(self.openapi_file)

        self.assertEqual(schema.info.title, "单词本")
        self.assertEqual(schema.servers[0].url, "http://127.0.0.1:8081")

    def test_load_and_save(self):
        """function_call requires empty fields, eg: items: {}, but yaml file doesn't
        contain any empty field
        """
        spec_dict = read_from_filename(self.openapi_file)
        schema = RemoteToolkit.from_openapi_file(self.openapi_file)
        saved_spec_dict = schema.to_openapi_dict()
        self.assertEqual(spec_dict[0], saved_spec_dict)

    def test_function_call_schemas(self):
        toolkit = RemoteToolkit.from_openapi_file(self.openapi_file)
        function_call_schemas = [tool.function_call_schema() for tool in toolkit.get_tools()]
        self.assertEqual(len(function_call_schemas), 4)

        self.assertEqual(function_call_schemas[0]["name"], "getWordbook")
        self.assertEqual(function_call_schemas[0]["responses"]["required"], ["wordbook"])
        self.assertEqual(function_call_schemas[0]["responses"]["properties"]["wordbook"]["type"], "array")
        self.assertEqual(function_call_schemas[3]["name"], "deleteWord")

    def test_get_typing_list_type(self):
        result = get_typing_list_type(List[int])
        self.assertEqual(result, "integer")

        result = get_typing_list_type(List[str])
        self.assertEqual(result, "string")

        result = get_typing_list_type(int)
        self.assertEqual(result, None)

        result = get_typing_list_type(dict)
        self.assertEqual(result, None)

        result = get_typing_list_type(List[ToolParameterView])
        self.assertEqual(result, "object")

    def test_json_type(self):
        result = json_type(List[int])
        self.assertEqual(result, "array")

        result = json_type(int)
        self.assertEqual(result, "integer")

        result = json_type(float)
        self.assertEqual(result, "number")

        result = json_type(ToolParameterView)
        self.assertEqual(result, "object")

    def test_list_tool_parameter_view(self):
        class SearchResponseDocument(ToolParameterView):
            document: str = Field(description="和query相关的规章片段")
            filename: str = Field(description="规章名称")
            page_num: int = Field(description="规章页数")

        class SearchToolOutputView(ToolParameterView):
            documents: List[SearchResponseDocument] = Field(description="检索结果，内容为住房和城乡建设部规章中和query相关的规章片段")

        open_api_dict = SearchToolOutputView.to_openapi_dict()

        expected_schema = {
            "type": "object",
            "required": ["documents"],
            "properties": {
                "documents": {
                    "type": "array",
                    "description": "检索结果，内容为住房和城乡建设部规章中和query相关的规章片段",
                    "items": {
                        "type": "object",
                        "required": ["document", "filename", "page_num"],
                        "properties": {
                            "document": {"type": "string", "description": "和query相关的规章片段"},
                            "filename": {"type": "string", "description": "规章名称"},
                            "page_num": {"type": "integer", "description": "规章页数"},
                        },
                    },
                }
            },
        }
        self.assertDictEqual(open_api_dict, expected_schema)

    def test_optional_tool_parameter_view(self):
        class SearchResponseDocument(ToolParameterView):
            document: str = Field(description="和query相关的规章片段")
            filename: str = Field(description="规章名称")
            page_num: int = Field(description="规章页数")

        class SearchToolOutputView(ToolParameterView):
            name: str = Field(description="测试名称")
            documents: Optional[SearchResponseDocument] = Field(
                description="检索结果，内容为住房和城乡建设部规章中和query相关的规章片段", default_factory=list
            )

        openapi_dict = SearchToolOutputView.to_openapi_dict()
        expected_openapi_dict = {
            "type": "object",
            "properties": {
                "name": {
                    "type": "string",
                    "description": "测试名称",
                },
                "documents": {
                    "type": "object",
                    "description": "检索结果，内容为住房和城乡建设部规章中和query相关的规章片段",
                    "required": ["document", "filename", "page_num"],
                    "properties": {
                        "document": {
                            "type": "string",
                            "description": "和query相关的规章片段",
                        },
                        "filename": {
                            "type": "string",
                            "description": "规章名称",
                        },
                        "page_num": {
                            "type": "integer",
                            "description": "规章页数",
                        },
                    },
                },
            },
            "required": ["name"],
        }
        self.assertDictEqual(openapi_dict, expected_openapi_dict)

    def test_enum_value(self):
        # TODO(wj-Mcat): to support enum[int/str/float]
        pass

    def test_is_optional_type(self):
        self.assertFalse(is_optional_type(List[int]))
        self.assertFalse(is_optional_type(int))
        self.assertFalse(is_optional_type(str))

        self.assertTrue(is_optional_type(Optional[int]))
        self.assertTrue(is_optional_type(Optional[str]))
        self.assertTrue(is_optional_type(Optional[ToolParameterView]))

    def test_load_examples(self):
        toolkit = RemoteToolkit.from_openapi_file("./tests/fixtures/openapi.yaml")
        toolkit.examples = toolkit.load_examples_yaml("./tests/fixtures/examples.yaml")
        self.assertEqual(len(toolkit.examples), 12)

        # add_word examples
        examples = toolkit.get_tool("getWordbook").examples

        self.assertEqual(len(examples), 2)
        self.assertEqual(examples[0].content, "展示单词列表")
        self.assertEqual(examples[1].function_call["name"], "getWordbook")
        self.assertEqual(examples[1].function_call["thoughts"], "这是一个展示单词本的需求")