Fix handling of list parameters (#16)

* Check if parameter is enum

* Check if parameter is enum

* Check array type

* Fix ruff errors

* Fix ruff and black

* Fix handling of primitive and typing list parameters

* Update gh workflow to use ruff format

* Remove unecessary include

* Remove useless change

* Fix ruff lint and format

* Make compatible with python 3.9

Author: jakecyr

.github/workflows/test-application.yml

@@ -24,13 +24,13 @@ jobs:
       - name: Install dependencies
         run: |
           python3 -m pip install poetry
-          poetry install --with dev
+          poetry install
       - name: Check formatting with ruff
         run: |
           poetry run ruff check .
       - name: Check formatting with black
         run: |
-          poetry run black --check .
+          poetry run ruff format --check .
       - name: Test Coverage
         run: |
           poetry run coverage run --source openai_function_calling -m pytest

.gitignore

@@ -7,3 +7,4 @@ dist
 .coverage
 htmlcov
 .mypy_cache
+.python-version

examples/weather_functions_infer.py

@@ -6,6 +6,7 @@
 
 import json
 from collections.abc import Callable
+from enum import Enum
 from typing import Any
 
 from openai import OpenAI
@@ -21,8 +22,15 @@
 from openai_function_calling.tool_helpers import ToolHelpers
 
 
+class TemperatureUnit(Enum):
+    """Temperature units."""
+
+    CELSIUS = "celsius"
+    FAHRENHEIT = "fahrenheit"
+
+
 # Define our functions.
-def get_current_weather(location: str, unit: str) -> str:
+def get_current_weather(location: str, unit: TemperatureUnit) -> str:
     """Get the current weather and return a summary."""
     return f"It is currently sunny in {location} and 75 degrees {unit}."
 

openai_function_calling/function.py

@@ -46,6 +46,7 @@ def __init__(
             parameters: A list of parameters.
             required_parameters: A list of parameter names that are required to run the\
                 function.
+
         """
         self.name: str = name
         self.description: str = description
@@ -80,6 +81,7 @@ def to_dict(self) -> FunctionDict:
 
         Returns:
             A JSON schema representation of the function.
+
         """
         return self.to_json_schema()
 
@@ -91,6 +93,7 @@ def to_json_schema(self) -> FunctionDict:
 
         Returns:
             A JSON schema representation of the function.
+
         """
         self.validate()
 
@@ -119,6 +122,7 @@ def merge(self, other_function: Function) -> None:
 
         Args:
             other_function: The other function to merge into the current.
+
         """
         if not self.name:
             self.name = other_function.name

openai_function_calling/function_inferrer.py

@@ -1,16 +1,24 @@
 """Function inferrer class definition."""
 
+from __future__ import annotations
+
+import dataclasses
 import inspect
-from collections.abc import Callable
-from typing import Any
+import typing
+from enum import EnumMeta
+from typing import TYPE_CHECKING, Any, get_args, get_origin, get_type_hints
 from warnings import warn
 
 from docstring_parser import Docstring, parser
 
 from openai_function_calling.function import Function
 from openai_function_calling.helper_functions import python_type_to_json_schema_type
+from openai_function_calling.json_schema_type import JsonSchemaType
 from openai_function_calling.parameter import Parameter
 
+if TYPE_CHECKING:  # pragma: no cover
+    from collections.abc import Callable
+
 
 class FunctionInferrer:
     """Class to help inferring a function definition from a reference."""
@@ -26,6 +34,7 @@ def infer_from_function_reference(function_reference: Callable) -> Function:
 
         Return:
             An instance of Function with inferred values.
+
         """
         inferred_from_annotations: Function = FunctionInferrer._infer_from_annotations(
             function_reference
@@ -51,6 +60,7 @@ def _infer_from_docstring(function_reference: Callable) -> Function:
 
         Returns:
             The inferred Function instance.
+
         """
         function_definition = Function(
             name=function_reference.__name__,
@@ -91,29 +101,53 @@ def _infer_from_annotations(function_reference: Callable) -> Function:
 
         Returns:
             The inferred Function instance.
-        """
-        function_definition = Function(
-            name=function_reference.__name__,
-            description="",
-            parameters=[],
-        )
-
-        if hasattr(function_reference, "__annotations__"):
-            annotations: dict[str, Any] = function_reference.__annotations__
 
-            for key in annotations:
-                if key == "return":
-                    continue
-
-                parameter_type: str = python_type_to_json_schema_type(
-                    annotations[key].__name__,
+        """
+        annotations: dict[str, Any] = get_type_hints(function_reference)
+        parameters: list[Parameter] = []
+
+        for param_name, annotation_type in annotations.items():
+            if param_name == "return":
+                continue
+
+            origin = get_origin(annotation_type) or annotation_type
+            args: tuple[Any, ...] = get_args(annotation_type)
+
+            if origin in [list, typing.List]:  # noqa: UP006
+                if not args:
+                    raise ValueError(
+                        f"Expected array parameter '{param_name}' to have an item type."
+                    )
+                item_type = args[0]
+                parameter_type = JsonSchemaType.ARRAY.value
+                array_item_type = python_type_to_json_schema_type(
+                    item_type.__name__ if hasattr(item_type, "__name__") else "Any"
+                )
+            elif origin in [dict, typing.Dict]:  # noqa: UP006
+                parameter_type = JsonSchemaType.OBJECT.value
+                array_item_type = None
+            else:
+                parameter_type = python_type_to_json_schema_type(
+                    annotation_type.__name__
+                    if hasattr(annotation_type, "__name__")
+                    else "Any"
                 )
 
-                function_definition.parameters.append(
-                    Parameter(name=key, type=parameter_type)
+                array_item_type = None
+
+            parameters.append(
+                Parameter(
+                    name=param_name,
+                    type=parameter_type,
+                    array_item_type=array_item_type,
                 )
+            )
 
-        return function_definition
+        return Function(
+            name=function_reference.__name__,
+            description="",
+            parameters=parameters,
+        )
 
     @staticmethod
     def _infer_from_inspection(function_reference: Callable) -> Function:
@@ -124,6 +158,7 @@ def _infer_from_inspection(function_reference: Callable) -> Function:
 
         Returns:
             The inferred Function instance.
+
         """
         function_definition = Function(
             name=function_reference.__name__,
@@ -135,9 +170,31 @@ def _infer_from_inspection(function_reference: Callable) -> Function:
 
         for name, parameter in inspected_parameters.items():
             parameter_type: str = python_type_to_json_schema_type(parameter.kind.name)
+            enum_values: list[str] | None = None
+
+            if parameter_type == "null":
+                if isinstance(parameter.annotation, EnumMeta):
+                    enum_values = list(
+                        parameter.annotation._value2member_map_.keys()  # noqa: SLF001
+                    )
+                    parameter_type = FunctionInferrer._infer_list_item_type(enum_values)
+                elif dataclasses.is_dataclass(parameter.annotation):
+                    parameter_type = JsonSchemaType.OBJECT.value
 
             function_definition.parameters.append(
-                Parameter(name=name, type=parameter_type)
+                Parameter(name=name, type=parameter_type, enum=enum_values)
             )
 
         return function_definition
+
+    @staticmethod
+    def _infer_list_item_type(list_of_items: list[Any]) -> str:
+        if len(list_of_items) == 0:
+            return JsonSchemaType.NULL.value
+
+        # Check if all items are the same type.
+        if len({type(item).__name__ for item in list_of_items}) == 1:
+            item: Any = type(list_of_items[0]).__name__
+            return python_type_to_json_schema_type(item)
+
+        return JsonSchemaType.ANY.value

openai_function_calling/helper_functions.py

@@ -13,6 +13,7 @@ def python_type_to_json_schema_type(python_type: str | None) -> str:
 
     Returns:
         A JSON schema type value.
+
     """
     json_schema_type: str = JsonSchemaType.NULL.value
 

openai_function_calling/parameter.py

@@ -52,6 +52,7 @@ def __init__(
                 is not set.
             ValueError: If the 'array_item_type' argument is set, but the 'type' is not\
                 'array'.
+
         """
         self.name: str = name
         self.type: str = type
@@ -67,6 +68,7 @@ def validate(self) -> None:
         Raises:
             ValueError: If 'array_item_type' is not set, but 'type' is array.
             ValueError: If 'array_item_type' is set, but 'type' is not array.
+
         """
         if self.type == JsonSchemaType.ARRAY and self.array_item_type is None:
             raise ValueError(
@@ -86,6 +88,7 @@ def to_json_schema(self) -> ParameterDict:
 
         Raises:
             ValueError: If there are validation errors. See the validate method.
+
         """
         self.validate()
 
@@ -111,6 +114,7 @@ def merge(self, other_parameter: Parameter) -> None:
 
         Args:
             other_parameter: The other parameter instance to merge into the current.
+
         """
         if not isinstance(other_parameter, Parameter):
             raise TypeError("Cannot merge non-parameter type into parameter.")
@@ -138,6 +142,7 @@ def __eq__(self, other: object) -> bool:
 
         Returns:
             If the other object is equal to the current instance.
+
         """
         if not isinstance(other, Parameter):
             return False

openai_function_calling/tool_helpers.py

@@ -26,6 +26,7 @@ def from_functions(functions: list[Function]) -> list[ChatCompletionToolParam]:
 
         Returns:
             A list of OpenAI chat completion tool parameters.
+
         """
         json_schemas: list[FunctionDict] = [f.to_json_schema() for f in functions]
         tool_params: list[ChatCompletionToolParam] = [
@@ -45,6 +46,7 @@ def infer_from_function_refs(
 
         Returns:
             A list of OpenAI chat completion tool parameters.
+
         """
         functions: list[Function] = [
             FunctionInferrer.infer_from_function_reference(f) for f in function_refs
@@ -60,5 +62,6 @@ def json_schema_to_tool_param(json_schema: FunctionDict) -> ChatCompletionToolPa
 
         Returns:
             An OpenAI chat completion tool parameter.
+
         """
         return {"type": "function", "function": cast(FunctionDefinition, json_schema)}