Merge pull request #8 from comfuture/enhance-doc-meta

Enhance `Doc` in `Annotation` and implement its tests

Author: comfuture

README.md

@@ -20,13 +20,14 @@ pip install function-schema
 
 ```python
 from typing import Annotated, Optional
+from function_schema import Doc
 import enum
 
 def get_weather(
-    city: Annotated[str, "The city to get the weather for"],
+    city: Annotated[str, Doc("The city to get the weather for")],
     unit: Annotated[
         Optional[str],
-        "The unit to return the temperature in",
+        Doc("The unit to return the temperature in"),
         enum.Enum("Unit", "celcius fahrenheit")
     ] = "celcius",
 ) -> str:
@@ -37,10 +38,7 @@ def get_weather(
 Function description is taken from the docstring.
 Type hinting with `typing.Annotated` for annotate additional information about the parameters and return type.
 
-- type can be `typing.Union`, `typing.Optional`. (`T | None` for python 3.10+)
-- string value of `Annotated` is used as a description
-- enum value of `Annotated` is used as an enum schema
-
+Then you can generate a schema for this function:
 ```python
 import json
 from function_schema import get_function_schema
@@ -89,20 +87,55 @@ 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.
 
-### Literal types can be used as Enum
+You can use any type hinting supported by python for the first argument of `Annotated`. including:
+`typing.Literal`, `typing.Optional`, `typing.Union`, and `T | None` for python 3.10+.  
+`Doc` class or plain string in `Annotated` is used for describe the parameter.
+`Doc` metadata is the [PEP propose](https://peps.python.org/pep-0727/) for standardizing the metadata in type hints.
+currently, implemented in `typing-extensions` module. Also `function_schema.Doc` is provided for compatibility.
+
+Enumeratable candidates can be defined with `enum.Enum` in the argument of `Annotated`.
+
+```python
+import enum
+
+class AnimalType(enum.Enum):
+    dog = enum.auto()
+    cat = enum.auto()
+
+def get_animal(
+    animal: Annotated[str, Doc("The animal to get"), AnimalType],
+) -> str:
+    """Returns the animal."""
+    return f"Animal is {animal.value}"
+```
+In this example, each name of `AnimalType` enums(`dog`, `cat`) is used as an enum schema.
+In shorthand, you can use `typing.Literal` as the type will do the same thing.
+
+```python
+def get_animal(
+    animal: Annotated[Literal["dog", "cat"], Doc("The animal to get")],
+) -> str:
+    """Returns the animal."""
+    return f"Animal is {animal}"
+```
+
+
+### Plain String in Annotated
+
+The string value of `Annotated` is used as a description for convenience.
 
 ```python
 def get_weather(
-    city: Annotated[str, "The city to get the weather for"],
-    unit: Annotated[
-        Optional[Literal["celcius", "fahrenheit"]], # <- Literal type represents Enum
-        "The unit to return the temperature in",
-    ] = "celcius",
+    city: Annotated[str, "The city to get the weather for"], # <- string value of Annotated is used as a description
+    unit: Annotated[Optional[str], "The unit to return the temperature in"] = "celcius",
 ) -> str:
     """Returns the weather for the given city."""
     return f"Weather for {city} is 20°C"
 ```
-The schema will be generated as the same as the previous example.
+
+But this would create a predefined meaning for any plain string inside of `Annotated`,
+and any tool that was using plain strings in them for any other purpose, which is currently allowed, would now be invalid.
+Please refer to the [PEP 0727, Plain String in Annotated](https://peps.python.org/pep-0727/#plain-string-in-annotated) for more information.
 
 ### Usage with OpenAI API
 
@@ -163,7 +196,7 @@ response = client.beta.tools.messages.create(
 ### CLI usage
 
 ```sh
-function_schema mymodule.py my_function
+function_schema mymodule.py my_function | jq
 ```
 
 ## License

function_schema/__init__.py

@@ -3,7 +3,7 @@
 """
 from .core import get_function_schema
 
-__version__ = "0.4.0"
+__version__ = "0.4.1"
 __all__ = (
     "get_function_schema",
     "__version__",

function_schema/core.py

@@ -1,16 +1,16 @@
 import enum
-import typing
 import inspect
 import platform
 import packaging.version
+from typing import Annotated, Optional, Union, Callable, Literal, Any, get_args, get_origin
 
 current_version = packaging.version.parse(platform.python_version())
 py_310 = packaging.version.parse("3.10")
 
 if current_version >= py_310:
     from types import UnionType
 else:
-    UnionType = typing.Union  # type: ignore
+    UnionType = Union  # type: ignore
 
 try:
     from typing import Doc
@@ -22,29 +22,23 @@ class Doc:
             def __init__(self, documentation: str, /):
                 self.documentation = documentation
 
-__all__ = ("get_function_schema", "guess_type", "Doc")
+__all__ = ("get_function_schema", "guess_type", "Doc", "Annotated")
 
-def is_doc_meta(obj):
+
+def is_doc_meta(obj: Annotated[Any, Doc("The object to be checked.")]) -> Annotated[bool, Doc("True if the object is a documentation object, False otherwise.")]:
     """
     Check if the given object is a documentation object.
-    Parameters:
-        obj (object): The object to be checked.
-    Returns:
-        bool: True if the object is a documentation object, False otherwise.
 
     Example:
     >>> is_doc_meta(Doc("This is a documentation object"))
     True
     """
     return getattr(obj, '__class__') == Doc and hasattr(obj, 'documentation')
 
-def unwrap_doc(obj: typing.Union[Doc, str]):
+
+def unwrap_doc(obj: Annotated[Union[Doc, str], Doc("The object to get the documentation string from.")]) -> Annotated[str, Doc("The documentation string.")]:
     """
     Get the documentation string from the given object.
-    Parameters:
-        obj (Doc | str): The object to get the documentation string from.
-    Returns:
-        str: The documentation string.
 
     Example:
     >>> unwrap_doc(Doc("This is a documentation object"))
@@ -58,12 +52,12 @@ def unwrap_doc(obj: typing.Union[Doc, str]):
 
 
 def get_function_schema(
-    func: typing.Annotated[typing.Callable, "The function to get the schema for"],
-    format: typing.Annotated[
-        typing.Optional[typing.Literal["openai", "claude"]],
-        "The format of the schema to return",
+    func: Annotated[Callable, Doc("The function to get the schema for")],
+    format: Annotated[
+        Optional[Literal["openai", "claude"]],
+        Doc("The format of the schema to return"),
     ] = "openai",
-) -> typing.Annotated[dict[str, typing.Any], "The JSON schema for the given function"]:
+) -> Annotated[dict[str, Any], Doc("The JSON schema for the given function")]:
     """
     Returns a JSON schema for the given function.
 
@@ -76,10 +70,10 @@ def get_function_schema(
     >>> from typing import Annotated, Optional
     >>> import enum
     >>> def get_weather(
-    ...     city: Annotated[str, "The city to get the weather for"],
+    ...     city: Annotated[str, Doc("The city to get the weather for")],
     ...     unit: Annotated[
     ...         Optional[str],
-    ...         "The unit to return the temperature in",
+    ...         Doc("The unit to return the temperature in"),
     ...         enum.Enum("Unit", "celcius fahrenheit")
     ...     ] = "celcius",
     ... ) -> str:
@@ -115,8 +109,8 @@ def get_function_schema(
         "required": [],
     }
     for name, param in params.items():
-        param_args = typing.get_args(param.annotation)
-        is_annotated = typing.get_origin(param.annotation) is typing.Annotated
+        param_args = get_args(param.annotation)
+        is_annotated = get_origin(param.annotation) is Annotated
 
         enum_ = None
         default_value = inspect._empty
@@ -126,10 +120,17 @@ def get_function_schema(
             (T, *_) = param_args
 
             # find description in param_args tuple
-            description = next(
-                (unwrap_doc(arg) for arg in param_args if isinstance(arg, (Doc, str))),
-                f"The {name} parameter",
-            )
+            try:
+                description = next(
+                    unwrap_doc(arg)
+                    for arg in param_args if isinstance(arg, Doc)
+                )
+            except StopIteration:
+                try:
+                    description = next(
+                        arg for arg in param_args if isinstance(arg, str))
+                except StopIteration:
+                    description = "The {name} parameter"
 
             # find enum in param_args tuple
             enum_ = next(
@@ -139,13 +140,13 @@ def get_function_schema(
                     if isinstance(arg, type) and issubclass(arg, enum.Enum)
                 ),
                 # use typing.Literal as enum if no enum found
-                typing.get_origin(T) is typing.Literal and typing.get_args(T) or None,
+                get_origin(T) is Literal and 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)
+            if get_origin(T) is Literal:
+                enum_ = get_args(T)
 
         # find default value for param
         if param.default is not inspect._empty:
@@ -157,20 +158,21 @@ def get_function_schema(
         }
 
         if enum_ is not None:
-            schema["properties"][name]["enum"] = [t for t in enum_ if t is not None]
+            schema["properties"][name]["enum"] = [
+                t for t in enum_ if t is not None]
 
         if default_value is not inspect._empty:
             schema["properties"][name]["default"] = default_value
 
         if (
-            typing.get_origin(T) is not typing.Literal
+            get_origin(T) is not Literal
             and not isinstance(None, T)
             and default_value is inspect._empty
         ):
             schema["required"].append(name)
 
-        if typing.get_origin(T) is typing.Literal:
-            if all(typing.get_args(T)):
+        if get_origin(T) is Literal:
+            if all(get_args(T)):
                 schema["required"].append(name)
 
     parms_key = "input_schema" if format == "claude" else "parameters"
@@ -185,24 +187,25 @@ def get_function_schema(
 
 
 def guess_type(
-    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"
+    T: Annotated[type, Doc("The type to guess the JSON schema type for")],
+) -> Annotated[
+    Union[str, list[str]], Doc(
+        "str | list of str that representing JSON schema type")
 ]:
     """Guesses the JSON schema type for the given python type."""
 
     # special case
-    if T is typing.Any:
+    if T is Any:
         return {}
 
-    origin = typing.get_origin(T)
+    origin = get_origin(T)
 
-    if origin is typing.Annotated:
-        return guess_type(typing.get_args(T)[0])
+    if origin is Annotated:
+        return guess_type(get_args(T)[0])
 
     # hacking around typing modules, `typing.Union` and `types.UnitonType`
-    if origin in [typing.Union, UnionType]:
-        union_types = [t for t in typing.get_args(T) if t is not type(None)]
+    if origin in [Union, UnionType]:
+        union_types = [t for t in get_args(T) if t is not type(None)]
         _types = [
             guess_type(union_type)
             for union_type in union_types
@@ -217,8 +220,8 @@ def guess_type(
             return _types[0]
         return _types
 
-    if origin is typing.Literal:
-        type_args = typing.Union[tuple(type(arg) for arg in typing.get_args(T))]
+    if origin is Literal:
+        type_args = Union[tuple(type(arg) for arg in get_args(T))]
         return guess_type(type_args)
     elif origin is list or origin is tuple:
         return "array"

pyproject.toml

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

test/test_pep_0727_doc.py

@@ -31,3 +31,29 @@ def func1(a: Annotated[str, Enum("Candidates", "a b c"), Doc("A string parameter
     assert schema["parameters"]["properties"]["a"]["description"] == "A string parameter", "parameter a should have a description"
     assert schema["parameters"]["properties"]["a"]["enum"] == [
         "a", "b", "c"], "parameter a should have enum values"
+
+
+def test_multiple_docs_in_annotation():
+    """Test a function with annotations with multiple Doc"""
+    def func1(a: Annotated[int, Doc("An integer parameter"), Doc("A number")]):
+        """My function"""
+        ...
+
+    schema = get_function_schema(func1)
+    assert schema["name"] == "func1", "Function name should be func1"
+    assert schema["description"] == "My function", "Function description should be there"
+    assert schema["parameters"]["properties"]["a"]["type"] == "number", "parameter a should be an integer"
+    assert schema["parameters"]["properties"]["a"]["description"] == "An integer parameter", "first description should be used"
+
+
+def test_mixed_docs_in_annotation():
+    """Test a function with annotations with mixed Doc and strings"""
+    def func1(a: Annotated[int, "An integer parameter", Doc("A number")]):
+        """My function"""
+        ...
+
+    schema = get_function_schema(func1)
+    assert schema["name"] == "func1", "Function name should be func1"
+    assert schema["description"] == "My function", "Function description should be there"
+    assert schema["parameters"]["properties"]["a"]["type"] == "number", "parameter a should be an integer"
+    assert schema["parameters"]["properties"]["a"]["description"] == "A number", "`Doc` should be used rather than string"