include test for tool calling format

rgbkrk · Jan 16, 2024 · 8e71221 · 8e71221
1 parent 83f80e2
commit 8e71221
Showing 1 changed file with 61 additions and 0 deletions.
diff --git a/tests/test_registry.py b/tests/test_registry.py
@@ -376,3 +376,64 @@ async def test_function_registry_call_edge_cases():
 
     with pytest.raises(UnknownFunctionError):
         await registry.call(None)  # type: ignore
+
+
+# Test the tool calling format
+@pytest.mark.asyncio
+async def test_function_registry_call_tool():
+    registry = FunctionRegistry()
+    registry.register(simple_func, SimpleModel)
+
+    registry.register(simple_func_with_model_arg)
+
+    tools = registry.tools
+
+    assert tools == [{
+        "type": "function",
+        "function": {
+            "name": "simple_func",
+            "description": "A simple test function",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "x": {"type": "integer"},
+                    "y": {"type": "string"},
+                    "z": {"type": "boolean", "default": False},
+                },
+                "required": ["x", "y"],
+            },
+        }
+    }, {
+        "type": "function",
+        "function": {
+            "name": "simple_func_with_model_arg",
+            "description": "A simple test function with a model argument",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "x": {"type": "integer"},
+                    "y": {"type": "string"},
+                    "z": {"default": False, "type": "boolean"},
+                    "model": {"allOf": [{"$ref": "#/$defs/SimpleModel"}], "default": None},
+                },
+                "required": ["x", "y"],
+                "$defs": {
+                    "SimpleModel": {
+                        "title": "SimpleModel",
+                        "type": "object",
+                        "properties": {
+                            "x": {"title": "X", "type": "integer"},
+                            "y": {"title": "Y", "type": "string"},
+                            "z": {
+                                "title": "Z",
+                                "description": "A simple boolean field",
+                                "default": False,
+                                "type": "boolean",
+                            },
+                        },
+                        "required": ["x", "y"],
+                    }
+                },
+            },
+        }
+    }]