mirror of https://github.com/aliasrobotics/cai.git
51 lines
1.5 KiB
Python
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"],
|
|
},
|
|
},
|
|
}
|