fix tool convert bug and add more tests

This commit is contained in:
yzlin 2024-04-11 16:55:52 +08:00
parent 391eda6a35
commit 5376a86929
2 changed files with 22 additions and 3 deletions

View file

@ -48,6 +48,14 @@ class DummyClass:
pass
class DummySubClass(DummyClass):
"""sub class docstring"""
def sub_method(self, df: pd.DataFrame):
"""sub method"""
pass
def dummy_fn(
df: pd.DataFrame,
s: str,
@ -117,6 +125,18 @@ def test_convert_code_to_tool_schema_class():
assert schema == expected
def test_convert_code_to_tool_schema_subclass():
schema = convert_code_to_tool_schema(DummySubClass)
assert "sub_method" in schema["methods"] # sub class method should be included
assert "fit" in schema["methods"] # parent class method should be included
def test_convert_code_to_tool_schema_include():
schema = convert_code_to_tool_schema(DummyClass, include=["fit"])
assert "fit" in schema["methods"]
assert "transform" not in schema["methods"]
def test_convert_code_to_tool_schema_function():
expected = {
"type": "function",