support Literal type, code refactories

Author: comfuture

function_schema/core.py

@@ -3,15 +3,11 @@
 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"],
     format: typing.Annotated[
-        typing.Optional[str], SchemaFormat, "The format of the schema to return"
+        typing.Optional[typing.Literal["openai", "claude"]],
+        "The format of the schema to return",
     ] = "openai",
 ) -> typing.Annotated[dict[str, typing.Any], "The JSON schema for the given function"]:
     """
@@ -66,7 +62,7 @@ def get_function_schema(
     }
     for name, param in params.items():
         param_args = typing.get_args(param.annotation)
-        is_annotated = len(param_args) > 1
+        is_annotated = typing.get_origin(param.annotation) is typing.Annotated
 
         enum_ = None
         default_value = inspect._empty
@@ -84,15 +80,18 @@ def get_function_schema(
             # find enum in param_args tuple
             enum_ = next(
                 (
-                    arg
+                    arg.name
                     for arg in param_args
                     if isinstance(arg, type) and issubclass(arg, enum.Enum)
                 ),
-                None,
+                # use typing.Literal as enum if no enum found
+                typing.get_origin(T) is typing.Literal and typing.get_args(T) or None,
             )
         else:
             T = param.annotation
             description = f"The {name} parameter"
+            if typing.get_origin(T) is typing.Literal:
+                enum_ = typing.get_args(T)
 
         # find default value for param
         if param.default is not inspect._empty:
@@ -104,12 +103,16 @@ def get_function_schema(
         }
 
         if enum_ is not None:
-            schema["properties"][name]["enum"] = [t.name for t in enum_]
+            schema["properties"][name]["enum"] = [t for t in enum_]
 
         if default_value is not inspect._empty:
             schema["properties"][name]["default"] = default_value
 
-        if not isinstance(None, T) and default_value is inspect._empty:
+        if (
+            typing.get_origin(T) is not typing.Literal
+            and not isinstance(None, T)
+            and default_value is inspect._empty
+        ):
             schema["required"].append(name)
 
     parms_key = "input_schema" if format == "claude" else "parameters"
@@ -128,9 +131,15 @@ def guess_type(
 ]:
     """Guesses the JSON schema type for the given python type."""
 
+    # special case
+    if T is typing.Any:
+        return {}
+
+    origin = typing.get_origin(T)
+
     # hacking around typing modules, `typing.Union` and `types.UnitonType`
-    union_types = typing.get_args(T)
-    if len(union_types) > 1:
+    if origin is typing.Union:
+        union_types = [t for t in typing.get_args(T) if t is not type(None)]
         _types = []
         for union_type in union_types:
             _types.append(guess_type(union_type))
@@ -144,6 +153,13 @@ def guess_type(
             return _types[0]
         return _types
 
+    if origin is typing.Literal:
+        return "string"
+    elif origin is list or origin is tuple:
+        return "array"
+    elif origin is dict:
+        return "object"
+
     if not isinstance(T, type):
         return
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "function-schema"
-version = "0.3.0"
+version = "0.3.3"
 requires-python = ">= 3.9"
 description = "A small utility to generate JSON schemas for python functions."
 readme = "README.md"

test/test_guess_type.py

@@ -14,6 +14,22 @@ def test_primitive():
     assert guess_type(bool) == "boolean"
 
 
+def test_typings():
+    """Test typing module types"""
+    assert guess_type(typing.Any) == {}
+    assert guess_type(typing.List) == "array"
+    assert guess_type(typing.Dict) == "object"
+    assert guess_type(typing.Tuple) == "array"
+
+    assert guess_type(typing.List[int]) == "array"
+    assert guess_type(typing.List[str]) == "array"
+    assert guess_type(typing.List[float]) == "array"
+    assert guess_type(typing.List[bool]) == "array"
+
+    assert guess_type(typing.Dict[str, int]) == "object"
+    assert guess_type(typing.Dict[str, str]) == "object"
+
+
 def test_optional():
     """Test optional types"""
     assert guess_type(typing.Optional[int]) == "integer"
@@ -51,12 +67,15 @@ def test_union():
     ]
 
 
+current_version = packaging.version.parse(platform.python_version())
+py_310 = packaging.version.parse("3.10")
+
+
+@pytest.mark.skipif(
+    current_version < py_310, reason="Union type is only available in Python 3.10+"
+)
 def test_union_type():
     """Test union types in Python 3.10+"""
-    current_version = packaging.version.parse(platform.python_version())
-    py_310 = packaging.version.parse("3.10")
-    if current_version < py_310:
-        pytest.skip("Union type is only available in Python 3.10+")
 
     assert guess_type(int | str) == ["integer", "string"]
     assert guess_type(int | float) == "number"
@@ -75,3 +94,9 @@ def test_union_type():
         "number",
         "boolean",
     ]
+
+
+def test_literal_type():
+    """Test literal type"""
+    assert guess_type(typing.Literal["a"]) == "string"
+    assert guess_type(typing.Literal[1]) == "string"

test/test_schema.py

@@ -1,5 +1,5 @@
 import enum
-from typing import Optional, Annotated
+from typing import Annotated, Literal
 from function_schema.core import get_function_schema
 
 
@@ -86,6 +86,7 @@ def func1(
 
 def test_annotated_function_with_enum():
     """Test a function with annotations and enum"""
+
     def func1(
         animal: Annotated[
             str,
@@ -101,6 +102,39 @@ def 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"
+
+
+def test_literal_type():
+    """Test literal type"""
+
+    def func1(animal: Annotated[Literal["Cat", "Dog"], "The animal you want to pet"]):
+        """My function"""
+        ...
+
+    schema = get_function_schema(func1)
+    print(schema)
     assert (
-        schema["parameters"]["properties"]["animal"]["enum"] == ["Cat", "Dog"]
-    ), "parameter animal should have an enum"
+        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"
+
+    def func2(animal: Literal["Cat", "Dog"]):
+        """My function"""
+        ...
+
+    schema = get_function_schema(func2)
+    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"