Merge branch 'main' of github.com:comfuture/function-schema

Author: comfuture

README.md

@@ -7,8 +7,8 @@
 This is a small utility to generate JSON schemas for python functions.
 With power of type annotations, it is possible to generate a schema for a function without describing it twice.
 
-At this moment, extracting schema from a function is only useful for [OpenAI API function-call](https://platform.openai.com/docs/guides/gpt/function-calling) feature.
-But it can be used for other purposes for example to generate documentation in the future.
+At this moment, extracting schema from a function is useful for [OpenAI Assistant Toll Calling](https://platform.openai.com/docs/assistants/tools/function-calling), [OpenAI API function-call](https://platform.openai.com/docs/guides/function-calling), and [Anthropic Claude Toll calling](https://docs.anthropic.com/claude/docs/tool-use) feature.
+And it can be used for other purposes for example to generate documentation in the future.
 
 ## Installation
 
@@ -79,20 +79,64 @@ Will output:
 }
 ```
 
+for claude, you should pass 2nd argument as SchemaFormat.claude or `claude`:
+
+```python
+from function_schema import get_function_schema
+
+schema = get_function_schema(get_weather, "claude")
+```
+
+Please refer to the [Claude tool use](https://docs.anthropic.com/claude/docs/tool-use) documentation for more information.
+
 You can use this schema to make a function call in OpenAI API:
 ```python
 import openai
 openai.api_key = "sk-..."
 
-result = openai.ChatCompletion.create(
+# Create an assistant with the function
+assistant = client.beta.assistants.create(
+    instructions="You are a weather bot. Use the provided functions to answer questions.",
+    model="gpt-4-turbo-preview",
+    tools=[{
+        "type": "function",
+        "function": get_function_schema(get_weather),
+    }]
+)
+
+run = client.beta.messages.create(
+    assistant_id=assistant.id,
+    messages=[
+        {"role": "user", "content": "What's the weather like in Seoul?"}
+    ]
+)
+
+# or with chat completion
+
+result = openai.chat.completion.create(
     model="gpt-3.5-turbo",
     messages=[
-      {"role": "user", "content": "What's the weather like in Seoul?"}
-    ],
-    functions=[
-        get_function_schema(get_weather)
+        {"role": "user", "content": "What's the weather like in Seoul?"}
     ],
-    function_call="auto",
+    tools=[get_function_schema(get_weather)],
+    tool_call="auto",
+)
+```
+
+In claude api,
+
+```python
+import anthropic
+
+client = anthropic.Client()
+
+response = client.beta.tools.messages.create(
+    model="claude-3-opus-20240229",
+    max_tokens=4096,
+    tools=[get_function_schema(get_weather, "claude")],
+    messages=[
+        {"role": "user", "content": "What's the weather like in Seoul?"}
+    ]
 )
 ```
 

function_schema/cli.py

@@ -6,7 +6,9 @@
 
 
 def print_usage():
-    print("Usage: function_schema <file_path> <function_name>")
+    print("Usage: function_schema <file_path> <function_name> [format='openai']")
+    print("Example: function_schema ./tests/test_schema.py test_simple_function")
+    print("Example: function_schema ./tests/test_schema.py test_simple_function claude")
 
 
 def main():
@@ -18,14 +20,25 @@ def main():
         spec.loader.exec_module(module)
         members = inspect.getmembers(module)
 
+        try:
+            format = sys.argv[3]
+        except IndexError:
+            format = "openai"
+        if format not in ["openai", "claude"]:
+            print(
+                "Invalid format. Use 'openai' or 'claude'. using 'openai'",
+                file=sys.stderr,
+            )
+
         for name, func in members:
-            if name == sys.argv[2]:
-                print(json.dumps(get_function_schema(func), indent=2))
+            if name >= sys.argv[2]:
+                print(json.dumps(get_function_schema(func, format), indent=2))
                 sys.exit(0)
         print(f"Function {sys.argv[2]} not found in {file_path}")
     except IndexError:
         print_usage()
         sys.exit(1)
 
+
 if __name__ == "__main__":
     main()

function_schema/core.py

@@ -3,8 +3,16 @@
 import inspect
 
 
+class SchemaFormat(str, enum.Enum):
+    openai = "openai"
+    claude = "claude"
+
+
 def get_function_schema(
-    func: typing.Annotated[typing.Callable, "The function to get the schema for"]
+    func: typing.Annotated[typing.Callable, "The function to get the schema for"],
+    format: typing.Annotated[
+        typing.Optional[str], SchemaFormat, "The format of the schema to return"
+    ] = "openai",
 ) -> typing.Annotated[dict[str, typing.Any], "The JSON schema for the given function"]:
     """
     Returns a JSON schema for the given function.
@@ -103,15 +111,18 @@ def get_function_schema(
 
         if not isinstance(None, T) and default_value is inspect._empty:
             schema["required"].append(name)
+
+    parms_key = "input_schema" if format == "claude" else "parameters"
+
     return {
         "name": func.__name__,
         "description": inspect.getdoc(func),
-        "parameters": schema,
+        parms_key: schema,
     }
 
 
 def guess_type(
-    T: typing.Annotated[type, "The type to guess the JSON schema type for"]
+    T: typing.Annotated[type, "The type to guess the JSON schema type for"],
 ) -> typing.Annotated[
     typing.Union[str, list[str]], "str | list of str that representing JSON schema type"
 ]:
@@ -123,11 +134,11 @@ def guess_type(
         _types = []
         for union_type in union_types:
             _types.append(guess_type(union_type))
-        _types = [t for t in _types if t is not None] # exclude None
+        _types = [t for t in _types if t is not None]  # exclude None
 
         # number contains integer in JSON schema
-        if 'number' in _types and 'integer' in _types:
-            _types.remove('integer')
+        if "number" in _types and "integer" in _types:
+            _types.remove("integer")
 
         if len(_types) == 1:
             return _types[0]
@@ -136,7 +147,7 @@ def guess_type(
     if not isinstance(T, type):
         return
 
-    if T.__name__ == 'NoneType':
+    if T.__name__ == "NoneType":
         return
 
     if issubclass(T, str):

test/test_schema_type.py

@@ -0,0 +1,35 @@
+from typing import Annotated
+import enum
+from function_schema.core import get_function_schema
+
+
+def test_function_schema_type():
+    """Test a function schema for claude"""
+
+    def func1(
+        animal: Annotated[
+            str,
+            "The animal you want to pet",
+            enum.Enum("Animal", "Cat Dog"),
+        ],
+    ):
+        """My function"""
+        ...
+
+    schema = get_function_schema(func1)
+    assert (
+        schema["parameters"]["properties"]["animal"]["type"] == "string"
+    ), "parameter animal should be a string"
+    assert schema["parameters"]["properties"]["animal"]["enum"] == [
+        "Cat",
+        "Dog",
+    ], "parameter animal should have an enum"
+
+    schema2 = get_function_schema(func1, format="claude")
+    assert (
+        schema2["input_schema"]["properties"]["animal"]["type"] == "string"
+    ), "parameter animal should be a string"
+    assert schema2["input_schema"]["properties"]["animal"]["enum"] == [
+        "Cat",
+        "Dog",
+    ], "parameter animal should have an enum"