cai/tests/test_util.py

51 lines
1.5 KiB
Python

from cai.util import function_to_json
def test_basic_function():
def basic_function(arg1, arg2):
return arg1 + arg2
result = function_to_json(basic_function)
assert result == {
"type": "function",
"function": {
"name": "basic_function",
"description": "",
"parameters": {
"type": "object",
"properties": {
"arg1": {"type": "string"},
"arg2": {"type": "string"},
},
"required": ["arg1", "arg2"],
},
},
}
def test_complex_function():
def complex_function_with_types_and_descriptions(
arg1: int, arg2: str, arg3: float = 3.14, arg4: bool = False
):
"""This is a complex function with a docstring."""
pass
result = function_to_json(complex_function_with_types_and_descriptions)
assert result == {
"type": "function",
"function": {
"name": "complex_function_with_types_and_descriptions",
"description": "This is a complex function with a docstring.",
"parameters": {
"type": "object",
"properties": {
"arg1": {"type": "integer"},
"arg2": {"type": "string"},
"arg3": {"type": "number"},
"arg4": {"type": "boolean"},
},
"required": ["arg1", "arg2"],
},
},
}