From f3ee51ffb1dd3d09a915b4a8f4df273c844d9c63 Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Thu, 8 Feb 2024 20:03:38 +0100 Subject: [PATCH 01/16] Add support for typing.Literal #10 --- tests/test_schema.py | 122 ++++++++++++++++++++++++++++++++++++++++-- tool2schema/schema.py | 8 ++- 2 files changed, 124 insertions(+), 6 deletions(-) diff --git a/tests/test_schema.py b/tests/test_schema.py index fe39e15..f370ac4 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -1,4 +1,4 @@ -from typing import List, Optional +from typing import List, Optional, Literal from tool2schema import ( FindGPTEnabled, @@ -288,9 +288,9 @@ def test_function_tags_tune(): assert function_tags.tags == ["test"] -######################################## -# Example function to test with enum # -######################################## +######################################################### +# Example function to test with enum (using add_enum) # +######################################################### @GPTEnabled @@ -816,3 +816,117 @@ def test_function_typing_list_no_type(): } assert function_typing_list_no_type.schema.to_json() == expected_schema assert function_typing_list_no_type.tags == [] + + +###################################################### +# Example functions with typing.Literal annotation # +###################################################### + + +@GPTEnabled +def function_typing_literal_int(a: Literal[1, 2, 3], b: str, c: bool = False, d: list[int] = [1, 2, 3]): + """ + This is a test function. + + :param a: This is a parameter; + :param b: This is another parameter; + :param c: This is a boolean parameter; + :param d: This is a list parameter; + """ + return a, b, c, d + + +def test_function_typing_literal_int(): + # Check schema + expected_schema = { + "type": "function", + "function": { + "name": "function_typing_literal_int", + "description": "This is a test function.", + "parameters": { + "type": "object", + "properties": { + "a": { + "type": "integer", + "description": "This is a parameter", + "enum": [1, 2, 3], + }, + "b": { + "type": "string", + "description": "This is another parameter", + }, + "c": { + "type": "boolean", + "description": "This is a boolean parameter", + "default": False, + }, + "d": { + "type": "array", + "description": "This is a list parameter", + "items": { + "type": "integer", + }, + "default": [1, 2, 3], + }, + }, + "required": ["a", "b"], + }, + }, + } + assert function_typing_literal_int.schema.to_json() == expected_schema + assert function_typing_literal_int.tags == [] + + +@GPTEnabled +def function_typing_literal_string(a: Literal["a", "b", "c"], b: str, c: bool = False, d: list[int] = [1, 2, 3]): + """ + This is a test function. + + :param a: This is a parameter; + :param b: This is another parameter; + :param c: This is a boolean parameter; + :param d: This is a list parameter; + """ + return a, b, c, d + + +def test_function_typing_literal_string(): + # Check schema + expected_schema = { + "type": "function", + "function": { + "name": "function_typing_literal_string", + "description": "This is a test function.", + "parameters": { + "type": "object", + "properties": { + "a": { + "type": "string", + "description": "This is a parameter", + "enum": ["a", "b", "c"], + }, + "b": { + "type": "string", + "description": "This is another parameter", + }, + "c": { + "type": "boolean", + "description": "This is a boolean parameter", + "default": False, + }, + "d": { + "type": "array", + "description": "This is a list parameter", + "items": { + "type": "integer", + }, + "default": [1, 2, 3], + }, + }, + "required": ["a", "b"], + }, + }, + } + assert function_typing_literal_string.schema.to_json() == expected_schema + assert function_typing_literal_string.tags == [] + diff --git a/tool2schema/schema.py b/tool2schema/schema.py index 05a1374..480cce2 100644 --- a/tool2schema/schema.py +++ b/tool2schema/schema.py @@ -5,7 +5,7 @@ from enum import Enum from inspect import Parameter from types import ModuleType -from typing import Callable, Optional +from typing import Callable, Optional, get_args class SchemaType(Enum): @@ -130,7 +130,6 @@ def add_enum(self, n: str, enum: list) -> "FunctionSchema": """ Add enum property to a particular function parameter. - :param schema: The schema to modify; :param n: The name of the parameter with the enum values; :param enum: The list of values for the enum parameter; """ @@ -232,6 +231,11 @@ def _param_type(o: Parameter, pschema: dict) -> dict: pschema["type"] = FunctionSchema.TYPE_MAP["list"] if (sub_type := FunctionSchema._sub_type(o)) is not None: pschema["items"] = {"type": sub_type} + elif re.match(r"typing\.Literal.*", str(o.annotation)): + pschema["type"] = FunctionSchema.TYPE_MAP.get( + type(get_args(o.annotation)[0]).__name__, "object" + ) + pschema["enum"] = list(get_args(o.annotation)) elif o.annotation.__name__ == "list": pschema["type"] = FunctionSchema.TYPE_MAP["list"] if (sub_type := FunctionSchema._sub_type(o)) is not None: From aa8da6b7b6e54a3c23373db3bb9ad4fd21198a73 Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Thu, 8 Feb 2024 22:02:00 +0100 Subject: [PATCH 02/16] Add support for enum.Enum #10 Subclasses of enum.Enum can now be be used as type hints --- tests/test_schema.py | 125 ++++++++++++++++++++++++++++++++++++++++++ tool2schema/schema.py | 6 ++ 2 files changed, 131 insertions(+) diff --git a/tests/test_schema.py b/tests/test_schema.py index f370ac4..23cf848 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -1,3 +1,4 @@ +from enum import Enum from typing import List, Optional, Literal from tool2schema import ( @@ -930,3 +931,127 @@ def test_function_typing_literal_string(): assert function_typing_literal_string.schema.to_json() == expected_schema assert function_typing_literal_string.tags == [] + +################################################# +# Example functions with enum.Enum annotation # +################################################# + + +class IntEnum(Enum): + A = 1 + B = 2 + C = 3 + + +@GPTEnabled +def function_enum_int(a: IntEnum, b: str, c: bool = False, d: list[int] = [1, 2, 3]): + """ + This is a test function. + + :param a: This is a parameter; + :param b: This is another parameter; + :param c: This is a boolean parameter; + :param d: This is a list parameter; + """ + return a, b, c, d + + +def test_function_enum_int(): + # Check schema + expected_schema = { + "type": "function", + "function": { + "name": "function_enum_int", + "description": "This is a test function.", + "parameters": { + "type": "object", + "properties": { + "a": { + "type": "integer", + "description": "This is a parameter", + "enum": [1, 2, 3], + }, + "b": { + "type": "string", + "description": "This is another parameter", + }, + "c": { + "type": "boolean", + "description": "This is a boolean parameter", + "default": False, + }, + "d": { + "type": "array", + "description": "This is a list parameter", + "items": { + "type": "integer", + }, + "default": [1, 2, 3], + }, + }, + "required": ["a", "b"], + }, + }, + } + assert function_enum_int.schema.to_json() == expected_schema + assert function_enum_int.tags == [] + + +class StrEnum(Enum): + A = "a" + B = "b" + C = "c" + + +@GPTEnabled +def function_enum_string(a: StrEnum, b: str, c: bool = False, d: list[int] = [1, 2, 3]): + """ + This is a test function. + + :param a: This is a parameter; + :param b: This is another parameter; + :param c: This is a boolean parameter; + :param d: This is a list parameter; + """ + return a, b, c, d + + +def test_function_enum_string(): + # Check schema + expected_schema = { + "type": "function", + "function": { + "name": "function_enum_string", + "description": "This is a test function.", + "parameters": { + "type": "object", + "properties": { + "a": { + "type": "string", + "description": "This is a parameter", + "enum": ["a", "b", "c"], + }, + "b": { + "type": "string", + "description": "This is another parameter", + }, + "c": { + "type": "boolean", + "description": "This is a boolean parameter", + "default": False, + }, + "d": { + "type": "array", + "description": "This is a list parameter", + "items": { + "type": "integer", + }, + "default": [1, 2, 3], + }, + }, + "required": ["a", "b"], + }, + }, + } + assert function_enum_string.schema.to_json() == expected_schema + assert function_enum_string.tags == [] diff --git a/tool2schema/schema.py b/tool2schema/schema.py index 480cce2..aae5536 100644 --- a/tool2schema/schema.py +++ b/tool2schema/schema.py @@ -236,6 +236,12 @@ def _param_type(o: Parameter, pschema: dict) -> dict: type(get_args(o.annotation)[0]).__name__, "object" ) pschema["enum"] = list(get_args(o.annotation)) + elif issubclass(o.annotation, Enum): + e_values = [e.value for e in o.annotation] + pschema["type"] = FunctionSchema.TYPE_MAP.get( + type(e_values[0]).__name__, "object" + ) + pschema["enum"] = e_values elif o.annotation.__name__ == "list": pschema["type"] = FunctionSchema.TYPE_MAP["list"] if (sub_type := FunctionSchema._sub_type(o)) is not None: From 15ecfdfe9a4219b76035a5f5dcad34565fb493f4 Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Sat, 10 Feb 2024 12:33:06 +0100 Subject: [PATCH 03/16] Add LiteralParameterSchema and EnumParameterSchema Add support for literal and enum type hints and related tests --- tests/test_schema.py | 8 +++-- tool2schema/parameter_schema.py | 52 ++++++++++++++++++++++++++++++++- tool2schema/schema.py | 2 +- 3 files changed, 58 insertions(+), 4 deletions(-) diff --git a/tests/test_schema.py b/tests/test_schema.py index 23cf848..1442292 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -825,7 +825,9 @@ def test_function_typing_list_no_type(): @GPTEnabled -def function_typing_literal_int(a: Literal[1, 2, 3], b: str, c: bool = False, d: list[int] = [1, 2, 3]): +def function_typing_literal_int( + a: Literal[1, 2, 3], b: str, c: bool = False, d: list[int] = [1, 2, 3] +): """ This is a test function. @@ -879,7 +881,9 @@ def test_function_typing_literal_int(): @GPTEnabled -def function_typing_literal_string(a: Literal["a", "b", "c"], b: str, c: bool = False, d: list[int] = [1, 2, 3]): +def function_typing_literal_string( + a: Literal["a", "b", "c"], b: str, c: bool = False, d: list[int] = [1, 2, 3] +): """ This is a test function. diff --git a/tool2schema/parameter_schema.py b/tool2schema/parameter_schema.py index 7c2b133..5d8c095 100644 --- a/tool2schema/parameter_schema.py +++ b/tool2schema/parameter_schema.py @@ -1,6 +1,7 @@ import re import typing -from inspect import Parameter +from enum import Enum +from inspect import Parameter, isclass TYPE_MAP = { "int": "integer", @@ -160,11 +161,60 @@ def _add_type(self, schema: dict): schema["type"] = sub_type +class EnumParameterSchema(ParameterSchema): + """ + Parameter schema for Enum types. + """ + + def __init__(self, parameter: Parameter, docstring: str = None): + super().__init__(parameter, docstring) + self.enum_values = [e.value for e in parameter.annotation] + + @staticmethod + def matches(parameter: Parameter) -> bool: + return ( + parameter.annotation != parameter.empty + and isclass(parameter.annotation) + and issubclass(parameter.annotation, Enum) + ) + + def _add_type(self, schema: dict): + schema["type"] = TYPE_MAP.get(type(self.enum_values[0]).__name__, "object") + + def _add_enum(self, schema: dict): + schema["enum"] = self.enum_values + + +class LiteralParameterSchema(ParameterSchema): + """ + Parameter schema for typing.Literal types. + """ + + def __init__(self, parameter: Parameter, docstring: str = None): + super().__init__(parameter, docstring) + self.enum_values = list(typing.get_args(parameter.annotation)) + + @staticmethod + def matches(parameter: Parameter) -> bool: + return ( + parameter.annotation != parameter.empty + and typing.get_origin(parameter.annotation) is typing.Literal + ) + + def _add_type(self, schema: dict): + schema["type"] = TYPE_MAP.get(type(self.enum_values[0]).__name__, "object") + + def _add_enum(self, schema: dict): + schema["enum"] = self.enum_values + + # Order matters: specific classes should appear before more generic ones; # for example, ListParameterSchema must precede ValueTypeSchema, # as they both match list types PARAMETER_SCHEMAS = [ OptionalParameterSchema, + LiteralParameterSchema, + EnumParameterSchema, ListParameterSchema, ValueTypeSchema, ] diff --git a/tool2schema/schema.py b/tool2schema/schema.py index bfe72cf..ed11005 100644 --- a/tool2schema/schema.py +++ b/tool2schema/schema.py @@ -5,7 +5,7 @@ from enum import Enum from inspect import Parameter from types import ModuleType -from typing import Callable, Optional, get_args +from typing import Callable, Optional from tool2schema.parameter_schema import PARAMETER_SCHEMAS From f61c784ac38112bd4d82227be866a508edcdda5d Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Sun, 11 Feb 2024 01:47:13 +0100 Subject: [PATCH 04/16] Use ReferenceSchema for enum tests --- tests/test_schema.py | 160 ++++--------------------------------------- 1 file changed, 14 insertions(+), 146 deletions(-) diff --git a/tests/test_schema.py b/tests/test_schema.py index adb0884..b2fc227 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -534,42 +534,9 @@ def function_typing_literal_int( def test_function_typing_literal_int(): # Check schema - expected_schema = { - "type": "function", - "function": { - "name": "function_typing_literal_int", - "description": "This is a test function.", - "parameters": { - "type": "object", - "properties": { - "a": { - "type": "integer", - "description": "This is a parameter", - "enum": [1, 2, 3], - }, - "b": { - "type": "string", - "description": "This is another parameter", - }, - "c": { - "type": "boolean", - "description": "This is a boolean parameter", - "default": False, - }, - "d": { - "type": "array", - "description": "This is a list parameter", - "items": { - "type": "integer", - }, - "default": [1, 2, 3], - }, - }, - "required": ["a", "b"], - }, - }, - } - assert function_typing_literal_int.schema.to_json() == expected_schema + rf = ReferenceSchema(function_typing_literal_int) + rf.get_param("a")["enum"] = [1, 2, 3] + assert function_typing_literal_int.schema.to_json() == rf.schema assert function_typing_literal_int.tags == [] @@ -590,42 +557,10 @@ def function_typing_literal_string( def test_function_typing_literal_string(): # Check schema - expected_schema = { - "type": "function", - "function": { - "name": "function_typing_literal_string", - "description": "This is a test function.", - "parameters": { - "type": "object", - "properties": { - "a": { - "type": "string", - "description": "This is a parameter", - "enum": ["a", "b", "c"], - }, - "b": { - "type": "string", - "description": "This is another parameter", - }, - "c": { - "type": "boolean", - "description": "This is a boolean parameter", - "default": False, - }, - "d": { - "type": "array", - "description": "This is a list parameter", - "items": { - "type": "integer", - }, - "default": [1, 2, 3], - }, - }, - "required": ["a", "b"], - }, - }, - } - assert function_typing_literal_string.schema.to_json() == expected_schema + rf = ReferenceSchema(function_typing_literal_string) + rf.get_param("a")["enum"] = ["a", "b", "c"] + rf.get_param("a")["type"] = "string" + assert function_typing_literal_string.schema.to_json() == rf.schema assert function_typing_literal_string.tags == [] @@ -654,43 +589,9 @@ def function_enum_int(a: IntEnum, b: str, c: bool = False, d: list[int] = [1, 2, def test_function_enum_int(): - # Check schema - expected_schema = { - "type": "function", - "function": { - "name": "function_enum_int", - "description": "This is a test function.", - "parameters": { - "type": "object", - "properties": { - "a": { - "type": "integer", - "description": "This is a parameter", - "enum": [1, 2, 3], - }, - "b": { - "type": "string", - "description": "This is another parameter", - }, - "c": { - "type": "boolean", - "description": "This is a boolean parameter", - "default": False, - }, - "d": { - "type": "array", - "description": "This is a list parameter", - "items": { - "type": "integer", - }, - "default": [1, 2, 3], - }, - }, - "required": ["a", "b"], - }, - }, - } - assert function_enum_int.schema.to_json() == expected_schema + rf = ReferenceSchema(function_enum_int) + rf.get_param("a")["enum"] = [1, 2, 3] + assert function_enum_int.schema.to_json() == rf.schema assert function_enum_int.tags == [] @@ -714,41 +615,8 @@ def function_enum_string(a: StrEnum, b: str, c: bool = False, d: list[int] = [1, def test_function_enum_string(): - # Check schema - expected_schema = { - "type": "function", - "function": { - "name": "function_enum_string", - "description": "This is a test function.", - "parameters": { - "type": "object", - "properties": { - "a": { - "type": "string", - "description": "This is a parameter", - "enum": ["a", "b", "c"], - }, - "b": { - "type": "string", - "description": "This is another parameter", - }, - "c": { - "type": "boolean", - "description": "This is a boolean parameter", - "default": False, - }, - "d": { - "type": "array", - "description": "This is a list parameter", - "items": { - "type": "integer", - }, - "default": [1, 2, 3], - }, - }, - "required": ["a", "b"], - }, - }, - } - assert function_enum_string.schema.to_json() == expected_schema + rf = ReferenceSchema(function_enum_string) + rf.get_param("a")["enum"] = ["a", "b", "c"] + rf.get_param("a")["type"] = "string" + assert function_enum_string.schema.to_json() == rf.schema assert function_enum_string.tags == [] From 0ceab9d77a5f9ae9d67461114fb0c76bc305a4ca Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Sun, 11 Feb 2024 11:57:00 +0100 Subject: [PATCH 05/16] Store ParameterSchema instances in FunctionSchema and convert enum values on invocation --- tests/test_schema.py | 24 +++++++ tool2schema/parameter_schema.py | 25 +++++++ tool2schema/schema.py | 115 +++++++++++++++++--------------- 3 files changed, 110 insertions(+), 54 deletions(-) diff --git a/tests/test_schema.py b/tests/test_schema.py index b2fc227..5f649b3 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -594,6 +594,18 @@ def test_function_enum_int(): assert function_enum_int.schema.to_json() == rf.schema assert function_enum_int.tags == [] + # Try invoking the function to verify that 1 is converted to StrEnum.A + a, _, _, _ = function_enum_int(a=IntEnum.A.value, b="", c=False, d=[]) + assert a == IntEnum.A + + # Verify it is possible to invoke the function with the Enum instance + a, _, _, _ = function_enum_int(a=IntEnum.A, b="", c=False, d=[]) + assert a == IntEnum.A + + # Verify it is possible to invoke the function with positional args + a, _, _, _ = function_enum_int(IntEnum.A, "", False, []) + assert a == IntEnum.A + class StrEnum(Enum): A = "a" @@ -620,3 +632,15 @@ def test_function_enum_string(): rf.get_param("a")["type"] = "string" assert function_enum_string.schema.to_json() == rf.schema assert function_enum_string.tags == [] + + # Try invoking the function to verify that "a" is converted to StrEnum.A + a, _, _, _ = function_enum_string(a=StrEnum.A.value, b="", c=False, d=[]) + assert a == StrEnum.A + + # Verify it is possible to invoke the function with the Enum instance + a, _, _, _ = function_enum_string(a=StrEnum.A, b="", c=False, d=[]) + assert a == StrEnum.A + + # Verify it is possible to invoke the function with positional args + a, _, _, _ = function_enum_string(StrEnum.A, "", False, []) + assert a == StrEnum.A diff --git a/tool2schema/parameter_schema.py b/tool2schema/parameter_schema.py index 5d8c095..8ca409e 100644 --- a/tool2schema/parameter_schema.py +++ b/tool2schema/parameter_schema.py @@ -92,6 +92,18 @@ def to_json(self) -> dict: self._add_enum(json) return json + def parse_value(self, value): + """ + Convert the given value from the JSON representation to an instance + that can be passed to the original method as a parameter. Overriding + methods should check whether the value needs to be converted, and return + it as is if no conversion is necessary. + + :param value: The value to be converted. + :return: An instance of the type required by the original method. + """ + return value + class ValueTypeSchema(ParameterSchema): """ @@ -184,6 +196,19 @@ def _add_type(self, schema: dict): def _add_enum(self, schema: dict): schema["enum"] = self.enum_values + def parse_value(self, value): + """ + Convert an enum value to an instance of the enum type. + + :param value: The value to be converted. + """ + if value in self.enum_values: + # convert to an enum instance + return self.parameter.annotation(value) + + # the user is invoking the method directly + return value + class LiteralParameterSchema(ParameterSchema): """ diff --git a/tool2schema/schema.py b/tool2schema/schema.py index ed11005..d2511b3 100644 --- a/tool2schema/schema.py +++ b/tool2schema/schema.py @@ -7,7 +7,7 @@ from types import ModuleType from typing import Callable, Optional -from tool2schema.parameter_schema import PARAMETER_SCHEMAS +from tool2schema.parameter_schema import PARAMETER_SCHEMAS, ParameterSchema class SchemaType(Enum): @@ -83,6 +83,14 @@ def __init__(self, func, **kwargs) -> None: functools.update_wrapper(self, func) def __call__(self, *args, **kwargs): + + for key in kwargs: + if key in self.schema.parameter_schemas: + # Convert the JSON value to the type expected by the method + kwargs[key] = self.schema.parameter_schemas[key].parse_value( + kwargs[key] + ) + return self.func(*args, **kwargs) def gpt_enabled(self) -> bool: @@ -105,11 +113,20 @@ def wrapper(function): class FunctionSchema: - """Automatically create a function schema.""" + """Automatically create a function schema for OpenAI.""" + + def __init__(self, f: Callable, schema_type: SchemaType = SchemaType.API): + """ + Initialize FunctionSchema for the given function. - def __init__(self, f: Callable): + :param f: The function to create a schema for; + :param schema_type: Type of schema; + """ self.f = f - self.schema = FunctionSchema.schema(f) + self.schema_type: SchemaType = schema_type + self.schema: dict = {} + self.parameter_schemas: dict[str, ParameterSchema] = {} + self._populate_schema() def to_json(self, schema_type: SchemaType = SchemaType.API) -> dict: """ @@ -117,7 +134,7 @@ def to_json(self, schema_type: SchemaType = SchemaType.API) -> dict: :param schema_type: Type of schema to return; """ if schema_type == SchemaType.TUNE: - return FunctionSchema.schema(self.f, schema_type)["function"] + return FunctionSchema(self.f, schema_type).to_json()["function"] return self.schema def add_enum(self, n: str, enum: list) -> "FunctionSchema": @@ -130,82 +147,72 @@ def add_enum(self, n: str, enum: list) -> "FunctionSchema": self.schema["function"]["parameters"]["properties"][n]["enum"] = enum return self - @staticmethod - def schema(f: Callable, schema_type=SchemaType.API) -> dict: + def _populate_schema(self) -> None: """ - Construct a function schema for OpenAI. - - :param f: Function for which to construct the schema; + Populate the schema dictionary. """ - fschema = {"type": "function", "function": {"name": f.__name__}} - fschema = FunctionSchema._function_description(f, fschema) - fschema = FunctionSchema._param_schema(f, fschema, schema_type) - return fschema + self.schema["type"] = "function" + self.schema["function"] = {"name": self.f.__name__} + + description = self._extract_description() + + # Add the function description even if it is an empty string + if description is not None: + self.schema["function"]["description"] = description - @staticmethod - def _function_description(f: Callable, fschema: dict) -> dict: + self._populate_parameter_schema() + + def _extract_description(self) -> Optional[str]: """ - Extract the function description. + Extract the function description, if present. - :param f: Function from which to extract description; + :return: The function description, or None if not present; """ - if docstring := f.__doc__: # Check if docstring exists + if docstring := self.f.__doc__: # Check if docstring exists docstring = " ".join( [x.strip() for x in docstring.replace("\n", " ").split()] ) if desc := re.findall(r"(.*?):param", docstring): - fschema["function"]["description"] = desc[0].strip() - return fschema - fschema["function"]["description"] = docstring.strip() - return fschema + return desc[0].strip() - @staticmethod - def _param_schema(f: Callable, fschema: dict, schema_type: SchemaType) -> dict: - """ - Construct the parameter schema for a function. + return docstring.strip() - :param f: Function to extra parameter schema from; + return None + + def _populate_parameter_schema(self) -> None: """ - param_schema = {"type": "object", "properties": {}} - if params := FunctionSchema._param_properties(f): - param_schema["properties"] = params - if required_params := FunctionSchema._param_required(f): - param_schema["required"] = required_params - fschema["function"]["parameters"] = param_schema - elif schema_type == SchemaType.TUNE: - fschema["function"]["parameters"] = param_schema - return fschema - - @staticmethod - def _param_properties(f: Callable) -> dict: + Populate the parameters' dictionary. """ - Construct the parameter properties for a function. + json_schema = dict() - :param f: Function to extra parameter properties from; - """ - pschema = dict() - for n, o in inspect.signature(f).parameters.items(): + for n, o in inspect.signature(self.f).parameters.items(): if n == "kwargs": continue # Skip kwargs for Param in PARAMETER_SCHEMAS: if Param.matches(o): - pschema[n] = Param(o, f.__doc__).to_json() + p = Param(o, self.f.__doc__) + json_schema[n] = p.to_json() + self.parameter_schemas[n] = p break - return pschema + if self.parameter_schemas or self.schema_type == SchemaType.TUNE: + self.schema["function"]["parameters"] = {"type": "object", "properties": {}} - @staticmethod - def _param_required(f: Callable) -> dict: - """ - Get the list of required parameters for a function. + if self.parameter_schemas: + self.schema["function"]["parameters"]["properties"] = json_schema + self._populate_required_parameters() - :param f: Function to extract required parameters from; + def _populate_required_parameters(self) -> None: + """ + Populate the list of required parameters. """ req_params = [] - for n, o in inspect.signature(f).parameters.items(): + for n, o in inspect.signature(self.f).parameters.items(): if n == "kwargs": continue # Skip kwargs if o.default == Parameter.empty: req_params.append(n) - return req_params + + if req_params: + self.schema["function"]["parameters"]["required"] = req_params From 247186088115e8a54c1c0ee05fc324c19164089b Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Sun, 11 Feb 2024 17:03:30 +0100 Subject: [PATCH 06/16] Convert enum to names rather than values Convert Enum type hints to names instead of values in the parameters schema --- tests/test_schema.py | 68 ++++++++------------------------- tool2schema/parameter_schema.py | 18 ++++----- 2 files changed, 25 insertions(+), 61 deletions(-) diff --git a/tests/test_schema.py b/tests/test_schema.py index 5f649b3..a8e669f 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -569,52 +569,16 @@ def test_function_typing_literal_string(): ################################################# -class IntEnum(Enum): +class CustomEnum(Enum): A = 1 B = 2 C = 3 @GPTEnabled -def function_enum_int(a: IntEnum, b: str, c: bool = False, d: list[int] = [1, 2, 3]): - """ - This is a test function. - - :param a: This is a parameter; - :param b: This is another parameter; - :param c: This is a boolean parameter; - :param d: This is a list parameter; - """ - return a, b, c, d - - -def test_function_enum_int(): - rf = ReferenceSchema(function_enum_int) - rf.get_param("a")["enum"] = [1, 2, 3] - assert function_enum_int.schema.to_json() == rf.schema - assert function_enum_int.tags == [] - - # Try invoking the function to verify that 1 is converted to StrEnum.A - a, _, _, _ = function_enum_int(a=IntEnum.A.value, b="", c=False, d=[]) - assert a == IntEnum.A - - # Verify it is possible to invoke the function with the Enum instance - a, _, _, _ = function_enum_int(a=IntEnum.A, b="", c=False, d=[]) - assert a == IntEnum.A - - # Verify it is possible to invoke the function with positional args - a, _, _, _ = function_enum_int(IntEnum.A, "", False, []) - assert a == IntEnum.A - - -class StrEnum(Enum): - A = "a" - B = "b" - C = "c" - - -@GPTEnabled -def function_enum_string(a: StrEnum, b: str, c: bool = False, d: list[int] = [1, 2, 3]): +def function_custom_enum( + a: CustomEnum, b: str, c: bool = False, d: list[int] = [1, 2, 3] +): """ This is a test function. @@ -626,21 +590,21 @@ def function_enum_string(a: StrEnum, b: str, c: bool = False, d: list[int] = [1, return a, b, c, d -def test_function_enum_string(): - rf = ReferenceSchema(function_enum_string) - rf.get_param("a")["enum"] = ["a", "b", "c"] +def test_function_custom_enum(): + rf = ReferenceSchema(function_custom_enum) rf.get_param("a")["type"] = "string" - assert function_enum_string.schema.to_json() == rf.schema - assert function_enum_string.tags == [] + rf.get_param("a")["enum"] = [x.name for x in CustomEnum] + assert function_custom_enum.schema.to_json() == rf.schema + assert function_custom_enum.tags == [] - # Try invoking the function to verify that "a" is converted to StrEnum.A - a, _, _, _ = function_enum_string(a=StrEnum.A.value, b="", c=False, d=[]) - assert a == StrEnum.A + # Try invoking the function to verify that "A" is converted to CustomEnum.A + a, _, _, _ = function_custom_enum(a=CustomEnum.A.name, b="", c=False, d=[]) + assert a == CustomEnum.A # Verify it is possible to invoke the function with the Enum instance - a, _, _, _ = function_enum_string(a=StrEnum.A, b="", c=False, d=[]) - assert a == StrEnum.A + a, _, _, _ = function_custom_enum(a=CustomEnum.A, b="", c=False, d=[]) + assert a == CustomEnum.A # Verify it is possible to invoke the function with positional args - a, _, _, _ = function_enum_string(StrEnum.A, "", False, []) - assert a == StrEnum.A + a, _, _, _ = function_custom_enum(CustomEnum.A, "", False, []) + assert a == CustomEnum.A diff --git a/tool2schema/parameter_schema.py b/tool2schema/parameter_schema.py index 8ca409e..baf2fc1 100644 --- a/tool2schema/parameter_schema.py +++ b/tool2schema/parameter_schema.py @@ -180,7 +180,7 @@ class EnumParameterSchema(ParameterSchema): def __init__(self, parameter: Parameter, docstring: str = None): super().__init__(parameter, docstring) - self.enum_values = [e.value for e in parameter.annotation] + self.enum_names = [e.name for e in parameter.annotation] @staticmethod def matches(parameter: Parameter) -> bool: @@ -191,22 +191,22 @@ def matches(parameter: Parameter) -> bool: ) def _add_type(self, schema: dict): - schema["type"] = TYPE_MAP.get(type(self.enum_values[0]).__name__, "object") + schema["type"] = TYPE_MAP["str"] def _add_enum(self, schema: dict): - schema["enum"] = self.enum_values + schema["enum"] = self.enum_names def parse_value(self, value): """ - Convert an enum value to an instance of the enum type. + Convert an enum name to an instance of the enum type. - :param value: The value to be converted. + :param value: The enum name to be converted """ - if value in self.enum_values: - # convert to an enum instance - return self.parameter.annotation(value) + if value in self.enum_names: + # Convert to an enum instance + return self.parameter.annotation[value] - # the user is invoking the method directly + # The user is invoking the method directly return value From 3a9fd4a33d0c550292bba44021e3698d7c05e8b3 Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Mon, 12 Feb 2024 23:24:41 +0100 Subject: [PATCH 07/16] Avoid populating the schema before calling to_json Store parameter_schemas on initialisation but do not create the schema dictionary until to_json is called --- tests/test_schema.py | 1 + tool2schema/parameter_schema.py | 54 +++++++------- tool2schema/schema.py | 122 ++++++++++++++++++++------------ 3 files changed, 105 insertions(+), 72 deletions(-) diff --git a/tests/test_schema.py b/tests/test_schema.py index 8ed235b..9b71f72 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -271,6 +271,7 @@ def test_function_enum(): rf = ReferenceSchema(function_enum) rf.get_param("a")["enum"] = [1, 2, 3] assert function_enum.schema.to_json() == rf.schema + assert function_enum.schema.to_json(SchemaType.TUNE) == rf.tune_schema assert function_enum.tags == [] diff --git a/tool2schema/parameter_schema.py b/tool2schema/parameter_schema.py index 3ffb211..2df8f0d 100644 --- a/tool2schema/parameter_schema.py +++ b/tool2schema/parameter_schema.py @@ -129,7 +129,7 @@ def _get_sub_type(self): return None -class ListParameterSchema(GenericParameterSchema): +class ListTypeParameterSchema(GenericParameterSchema): """ Parameter schema for list (array) types. """ @@ -148,7 +148,7 @@ def _add_items(self, schema: dict): schema["items"] = {"type": sub_type} -class OptionalParameterSchema(GenericParameterSchema): +class OptionalTypeParameterSchema(GenericParameterSchema): """ Parameter schema for typing.Optional types. """ @@ -170,12 +170,28 @@ def _add_type(self, schema: dict): class EnumParameterSchema(ParameterSchema): """ - Parameter schema for Enum types. + Parameter schema for enumeration types. """ - def __init__(self, parameter: Parameter, docstring: str = None): + def __init__(self, values: list, parameter: Parameter, docstring: str = None): super().__init__(parameter, docstring) - self.enum_names = [e.name for e in parameter.annotation] + self.enum_values = values + + def _add_type(self, schema: dict): + schema["type"] = TYPE_MAP.get(type(self.enum_values[0]).__name__, "object") + + def _add_enum(self, schema: dict): + schema["enum"] = self.enum_values + + +class EnumTypeParameterSchema(EnumParameterSchema): + """ + Parameter schema for enum.Enum types. + """ + + def __init__(self, parameter: Parameter, docstring: str = None): + values = [e.name for e in parameter.annotation] + super().__init__(values, parameter, docstring) @staticmethod def matches(parameter: Parameter) -> bool: @@ -185,19 +201,13 @@ def matches(parameter: Parameter) -> bool: and issubclass(parameter.annotation, Enum) ) - def _add_type(self, schema: dict): - schema["type"] = TYPE_MAP["str"] - - def _add_enum(self, schema: dict): - schema["enum"] = self.enum_names - def parse_value(self, value): """ Convert an enum name to an instance of the enum type. :param value: The enum name to be converted """ - if value in self.enum_names: + if value in self.enum_values: # Convert to an enum instance return self.parameter.annotation[value] @@ -205,14 +215,14 @@ def parse_value(self, value): return value -class LiteralParameterSchema(ParameterSchema): +class LiteralTypeParameterSchema(EnumParameterSchema): """ Parameter schema for typing.Literal types. """ def __init__(self, parameter: Parameter, docstring: str = None): - super().__init__(parameter, docstring) - self.enum_values = list(typing.get_args(parameter.annotation)) + values = list(typing.get_args(parameter.annotation)) + super().__init__(values, parameter, docstring) @staticmethod def matches(parameter: Parameter) -> bool: @@ -221,20 +231,14 @@ def matches(parameter: Parameter) -> bool: and typing.get_origin(parameter.annotation) is typing.Literal ) - def _add_type(self, schema: dict): - schema["type"] = TYPE_MAP.get(type(self.enum_values[0]).__name__, "object") - - def _add_enum(self, schema: dict): - schema["enum"] = self.enum_values - # Order matters: specific classes should appear before more generic ones; # for example, ListParameterSchema must precede ValueTypeSchema, # as they both match list types PARAMETER_SCHEMAS = [ - OptionalParameterSchema, - LiteralParameterSchema, - EnumParameterSchema, - ListParameterSchema, + OptionalTypeParameterSchema, + LiteralTypeParameterSchema, + EnumTypeParameterSchema, + ListTypeParameterSchema, ValueTypeSchema, ] diff --git a/tool2schema/schema.py b/tool2schema/schema.py index 74b1952..55f68e5 100644 --- a/tool2schema/schema.py +++ b/tool2schema/schema.py @@ -7,7 +7,7 @@ from types import ModuleType from typing import Callable, Optional -from tool2schema.parameter_schema import PARAMETER_SCHEMAS, ParameterSchema +from tool2schema.parameter_schema import PARAMETER_SCHEMAS, EnumParameterSchema, ParameterSchema class SchemaType(Enum): @@ -111,18 +111,14 @@ def wrapper(function): class FunctionSchema: """Automatically create a function schema for OpenAI.""" - def __init__(self, f: Callable, schema_type: SchemaType = SchemaType.API): + def __init__(self, f: Callable): """ Initialize FunctionSchema for the given function. :param f: The function to create a schema for; - :param schema_type: Type of schema; """ self.f = f - self.schema_type: SchemaType = schema_type - self.schema: dict = {} - self.parameter_schemas: dict[str, ParameterSchema] = {} - self._populate_schema() + self.parameter_schemas: dict[str, ParameterSchema] = self._get_parameter_schemas() def to_json(self, schema_type: SchemaType = SchemaType.API) -> dict: """ @@ -130,8 +126,9 @@ def to_json(self, schema_type: SchemaType = SchemaType.API) -> dict: :param schema_type: Type of schema to return """ if schema_type == SchemaType.TUNE: - return FunctionSchema(self.f, schema_type).to_json()["function"] - return self.schema + return self._get_function_schema(schema_type) + + return self._get_schema() def add_enum(self, n: str, enum: list) -> "FunctionSchema": """ @@ -140,44 +137,68 @@ def add_enum(self, n: str, enum: list) -> "FunctionSchema": :param n: The name of the parameter with the enum values :param enum: The list of values for the enum parameter """ - self.schema["function"]["parameters"]["properties"][n]["enum"] = enum + p = self.parameter_schemas[n] + self.parameter_schemas[n] = EnumParameterSchema(enum, p.parameter, p.docstring) return self - def _populate_schema(self) -> None: + def _get_schema(self) -> dict: """ - Populate the schema dictionary. + Get the complete schema dictionary. """ - self.schema["type"] = "function" - self.schema["function"] = {"name": self.f.__name__} + # This dictionary is only used with the API schema type + return {"type": "function", "function": self._get_function_schema(SchemaType.API)} - description = self._extract_description() + def _get_function_schema(self, schema_type: SchemaType) -> dict: + """ + Get the function schema dictionary. + """ + schema = {"name": self.f.__name__} - # Add the function description even if it is an empty string - if description is not None: - self.schema["function"]["description"] = description + if self.parameter_schemas or schema_type == SchemaType.TUNE: + # If the schema type is tune, add the dictionary even if there are no parameters + schema["parameters"] = self._get_parameters_schema(schema_type) - self._populate_parameter_schema() + if (description := self._get_description()) is not None: + # Add the function description even if it is an empty string + schema["description"] = description - def _extract_description(self) -> Optional[str]: + return schema + + def _get_parameters_schema(self, schema_type: SchemaType) -> dict: """ - Extract the function description, if present. + Get the parameters schema dictionary. + """ + schema = {"type": "object"} - :return: The function description, or None if not present + if self.parameter_schemas or schema_type == SchemaType.TUNE: + # If the schema type is tune, add the dictionary even if empty + schema["properties"] = self._get_parameter_properties_schema() + + if required := self._get_required_parameters(): + schema["required"] = required + + return schema + + def _get_parameter_properties_schema(self) -> dict: """ - if docstring := self.f.__doc__: # Check if docstring exists - docstring = " ".join([x.strip() for x in docstring.replace("\n", " ").split()]) - if desc := re.findall(r"(.*?):param", docstring): - return desc[0].strip() + Get the properties schema for the function. + """ + schema = dict() - return docstring.strip() + for n, p in self.parameter_schemas.items(): + schema[n] = p.to_json() - return None + return schema - def _populate_parameter_schema(self) -> None: + def _get_parameter_schemas(self) -> dict[str, ParameterSchema]: """ - Populate the parameters' dictionary. + Get a dictionary of parameter schemas for the function. + Ignored parameters are not included in the dictionary. + + :return: A dictionary with parameter names as keys and + parameter schemas as values """ - json_schema = dict() + parameters = dict() for n, o in inspect.signature(self.f).parameters.items(): if n == "kwargs": @@ -185,28 +206,35 @@ def _populate_parameter_schema(self) -> None: for Param in PARAMETER_SCHEMAS: if Param.matches(o): - p = Param(o, self.f.__doc__) - json_schema[n] = p.to_json() - self.parameter_schemas[n] = p + parameters[n] = Param(o, self.f.__doc__) break - if self.parameter_schemas or self.schema_type == SchemaType.TUNE: - self.schema["function"]["parameters"] = {"type": "object", "properties": {}} + return parameters - if self.parameter_schemas: - self.schema["function"]["parameters"]["properties"] = json_schema - self._populate_required_parameters() + def _get_description(self) -> Optional[str]: + """ + Extract the function description, if present. - def _populate_required_parameters(self) -> None: + :return: The function description, or None if not present """ - Populate the list of required parameters. + if docstring := self.f.__doc__: # Check if docstring exists + docstring = " ".join([x.strip() for x in docstring.replace("\n", " ").split()]) + if desc := re.findall(r"(.*?):param", docstring): + return desc[0].strip() + + return docstring.strip() + + return None + + def _get_required_parameters(self) -> list[str]: + """ + Get the list of required parameters. + + :return: The list of parameters without a default value """ req_params = [] - for n, o in inspect.signature(self.f).parameters.items(): - if n == "kwargs": - continue # Skip kwargs - if o.default == Parameter.empty: + for n, p in self.parameter_schemas.items(): + if p.parameter.default == Parameter.empty: req_params.append(n) - if req_params: - self.schema["function"]["parameters"]["required"] = req_params + return req_params From 4f42b4ce21bbe0934745f859c740fb4dbda18487 Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Thu, 15 Feb 2024 13:10:07 +0100 Subject: [PATCH 08/16] Refactor ParameterSchema add-to logic to get logic Change the internal logic of ParameterSchema to go from _add_...(dict) methods to _get_...(), for similarity with FunctionSchema --- tool2schema/parameter_schema.py | 101 ++++++++++++++++++-------------- 1 file changed, 57 insertions(+), 44 deletions(-) diff --git a/tool2schema/parameter_schema.py b/tool2schema/parameter_schema.py index 2df8f0d..91e15f1 100644 --- a/tool2schema/parameter_schema.py +++ b/tool2schema/parameter_schema.py @@ -2,6 +2,7 @@ import typing from enum import Enum from inspect import Parameter, isclass +from typing import Union TYPE_MAP = { "int": "integer", @@ -37,58 +38,69 @@ def matches(parameter: Parameter) -> bool: """ raise NotImplementedError() - def _add_type(self, schema: dict): + def _get_type(self) -> Union[str, Parameter.empty]: """ - Add the type of this parameter to the given schema. + Get the type of this parameter, to be added to the JSON schema. + Return `Parameter.empty` to omit the type from the schema. """ - raise NotImplementedError() + return Parameter.empty - def _add_items(self, schema: dict): + def _get_items(self) -> Union[str, Parameter.empty]: """ - Add the items property to the given schema. + Get the items property to be added to the JSON schema. + Return `Parameter.empty` to omit the items from the schema. """ - pass + return Parameter.empty - def _add_enum(self, schema: dict): + def _get_enum(self) -> Union[list[str], Parameter.empty]: """ - Add the enum property to the given schema. + Get the enum property to be added to the JSON schema. + Return `Parameter.empty` to omit the enum from the schema. """ - pass + return Parameter.empty - def _add_description(self, schema: dict): + def _get_description(self) -> Union[str, Parameter.empty]: """ - Add the description of this parameter, extracted from the function docstring, to the given - schema. + Get the description of this parameter, extracted from the function docstring, + to be added to the JSON schema. Return `Parameter.empty` to omit the description + from the schema. """ if self.docstring is None: - return + return Parameter.empty docstring = " ".join([x.strip() for x in self.docstring.replace("\n", " ").split()]) params = re.findall(r":param ([^:]*): (.*?)(?=:param|:type|:return|:rtype|$)", docstring) for name, desc in params: if name == self.parameter.name and desc: - schema["description"] = desc.strip() - return + return desc.strip() + + return Parameter.empty - def _add_default(self, schema: dict): + def _get_default(self) -> any: """ - Add the default value, when present, to the given schema. + Get the default value for this parameter, when present, to be added to the JSON schema. + Return `Parameter.empty` to omit the default value from the schema. """ - if self.parameter.default == Parameter.empty: - return - - schema["default"] = self.parameter.default + return self.parameter.default def to_json(self) -> dict: """ Return the json schema for this parameter. """ - json = {} - self._add_description(json) - self._add_default(json) - self._add_items(json) - self._add_type(json) - self._add_enum(json) + fields = { + "description": self._get_description, + "default": self._get_default, + "items": self._get_items, + "type": self._get_type, + "enum": self._get_enum, + } + + json = dict() + + for field in fields: + if (value := fields[field]()) != Parameter.empty: + json[field] = value + return json def parse_value(self, value): @@ -113,8 +125,8 @@ class ValueTypeSchema(ParameterSchema): def matches(parameter: Parameter) -> bool: return parameter.annotation != Parameter.empty and parameter.annotation.__name__ in TYPE_MAP - def _add_type(self, schema: dict): - schema["type"] = TYPE_MAP[self.parameter.annotation.__name__] + def _get_type(self) -> Union[str, Parameter.empty]: + return TYPE_MAP[self.parameter.annotation.__name__] class GenericParameterSchema(ParameterSchema): @@ -122,11 +134,11 @@ class GenericParameterSchema(ParameterSchema): Base class for generic parameter types supporting subscription. """ - def _get_sub_type(self): + def _get_sub_type(self) -> Union[str, Parameter.empty]: if args := typing.get_args(self.parameter.annotation): return TYPE_MAP.get(args[0].__name__, "object") - return None + return Parameter.empty class ListTypeParameterSchema(GenericParameterSchema): @@ -140,12 +152,14 @@ def matches(parameter: Parameter) -> bool: parameter.annotation is list or typing.get_origin(parameter.annotation) is list ) - def _add_type(self, schema: dict): - schema["type"] = TYPE_MAP["list"] + def _get_type(self) -> Union[str, Parameter.empty]: + return TYPE_MAP["list"] + + def _get_items(self) -> Union[str, Parameter.empty]: + if (sub_type := super()._get_sub_type()) != Parameter.empty: + return {"type": sub_type} - def _add_items(self, schema: dict): - if sub_type := super()._get_sub_type(): - schema["items"] = {"type": sub_type} + return Parameter.empty class OptionalTypeParameterSchema(GenericParameterSchema): @@ -158,14 +172,13 @@ def matches(parameter: Parameter) -> bool: args = typing.get_args(parameter.annotation) return ( parameter.annotation != parameter.empty - and typing.get_origin(parameter.annotation) is typing.Union + and typing.get_origin(parameter.annotation) is Union and len(args) == 2 and type(None) in args ) - def _add_type(self, schema: dict): - if sub_type := super()._get_sub_type(): - schema["type"] = sub_type + def _get_type(self) -> Union[str, Parameter.empty]: + return super()._get_sub_type() class EnumParameterSchema(ParameterSchema): @@ -177,11 +190,11 @@ def __init__(self, values: list, parameter: Parameter, docstring: str = None): super().__init__(parameter, docstring) self.enum_values = values - def _add_type(self, schema: dict): - schema["type"] = TYPE_MAP.get(type(self.enum_values[0]).__name__, "object") + def _get_type(self) -> Union[str, Parameter.empty]: + return TYPE_MAP.get(type(self.enum_values[0]).__name__, "object") - def _add_enum(self, schema: dict): - schema["enum"] = self.enum_values + def _get_enum(self) -> Union[list[str], Parameter.empty]: + return self.enum_values class EnumTypeParameterSchema(EnumParameterSchema): From 6057e6fe4b63cc125c450585d64a030bbc982f4b Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Thu, 15 Feb 2024 14:55:53 +0100 Subject: [PATCH 09/16] Support positional arguments value parsing GPT-enabled functions can now be invoked with json-converted values as positional arguments --- tests/test_schema.py | 7 ++++++- tool2schema/parameter_schema.py | 20 +++++++++++--------- tool2schema/schema.py | 14 +++++++++++--- 3 files changed, 28 insertions(+), 13 deletions(-) diff --git a/tests/test_schema.py b/tests/test_schema.py index 9b71f72..1e083a4 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -655,7 +655,12 @@ def test_function_custom_enum(): assert function_custom_enum.schema.to_json() == rf.schema assert function_custom_enum.tags == [] - # Try invoking the function to verify that "A" is converted to CustomEnum.A + # Try invoking the function to verify that "A" is converted to CustomEnum.A, + # passing the value as a positional argument + a, _, _, _ = function_custom_enum(CustomEnum.A.name, b="", c=False, d=[]) + assert a == CustomEnum.A + + # Same as above but passing the value as a keyword argument a, _, _, _ = function_custom_enum(a=CustomEnum.A.name, b="", c=False, d=[]) assert a == CustomEnum.A diff --git a/tool2schema/parameter_schema.py b/tool2schema/parameter_schema.py index 91e15f1..93a3657 100644 --- a/tool2schema/parameter_schema.py +++ b/tool2schema/parameter_schema.py @@ -19,15 +19,17 @@ class ParameterSchema: inspect.Parameter and a function documentation string. """ - def __init__(self, parameter: Parameter, docstring: str = None): + def __init__(self, parameter: Parameter, index: int, docstring: str = None): """ Create a new parameter schema. :param parameter: The parameter to create a schema for + :param index: The index of the parameter in the function signature :param docstring: The docstring for the function containing the parameter """ - self.parameter = parameter - self.docstring = docstring + self.parameter: Parameter = parameter + self.index: int = index + self.docstring: str = docstring @staticmethod def matches(parameter: Parameter) -> bool: @@ -186,8 +188,8 @@ class EnumParameterSchema(ParameterSchema): Parameter schema for enumeration types. """ - def __init__(self, values: list, parameter: Parameter, docstring: str = None): - super().__init__(parameter, docstring) + def __init__(self, values: list, parameter: Parameter, index: int, docstring: str = None): + super().__init__(parameter, index, docstring) self.enum_values = values def _get_type(self) -> Union[str, Parameter.empty]: @@ -202,9 +204,9 @@ class EnumTypeParameterSchema(EnumParameterSchema): Parameter schema for enum.Enum types. """ - def __init__(self, parameter: Parameter, docstring: str = None): + def __init__(self, parameter: Parameter, index: int, docstring: str = None): values = [e.name for e in parameter.annotation] - super().__init__(values, parameter, docstring) + super().__init__(values, parameter, index, docstring) @staticmethod def matches(parameter: Parameter) -> bool: @@ -233,9 +235,9 @@ class LiteralTypeParameterSchema(EnumParameterSchema): Parameter schema for typing.Literal types. """ - def __init__(self, parameter: Parameter, docstring: str = None): + def __init__(self, parameter: Parameter, index: int, docstring: str = None): values = list(typing.get_args(parameter.annotation)) - super().__init__(values, parameter, docstring) + super().__init__(values, parameter, index, docstring) @staticmethod def matches(parameter: Parameter) -> bool: diff --git a/tool2schema/schema.py b/tool2schema/schema.py index 55f68e5..d4eb5b5 100644 --- a/tool2schema/schema.py +++ b/tool2schema/schema.py @@ -82,6 +82,14 @@ def __init__(self, func, **kwargs) -> None: def __call__(self, *args, **kwargs): + args = list(args) # Tuple is immutable, thus convert to list + + for i, arg in enumerate(args): + for p in self.schema.parameter_schemas.values(): + if p.index == i: + # Convert the JSON value to the type expected by the method + args[i] = p.parse_value(arg) + for key in kwargs: if key in self.schema.parameter_schemas: # Convert the JSON value to the type expected by the method @@ -138,7 +146,7 @@ def add_enum(self, n: str, enum: list) -> "FunctionSchema": :param enum: The list of values for the enum parameter """ p = self.parameter_schemas[n] - self.parameter_schemas[n] = EnumParameterSchema(enum, p.parameter, p.docstring) + self.parameter_schemas[n] = EnumParameterSchema(enum, p.parameter, p.index, p.docstring) return self def _get_schema(self) -> dict: @@ -200,13 +208,13 @@ def _get_parameter_schemas(self) -> dict[str, ParameterSchema]: """ parameters = dict() - for n, o in inspect.signature(self.f).parameters.items(): + for i, (n, o) in enumerate(inspect.signature(self.f).parameters.items()): if n == "kwargs": continue # Skip kwargs for Param in PARAMETER_SCHEMAS: if Param.matches(o): - parameters[n] = Param(o, self.f.__doc__) + parameters[n] = Param(o, i, self.f.__doc__) break return parameters From a420a7d5e847df256e1f31de264d429ecf0a1a1b Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Thu, 15 Feb 2024 15:24:40 +0100 Subject: [PATCH 10/16] Support enum.Enum parameter with default value Add support for keyword arguments of type enum.Enum with a default value. Rename ParameterSchema.parse_value to decode_value, and add encode_value method to encode the default value when necessary. --- tests/test_schema.py | 41 ++++++++++++++++++++++++++++++--- tool2schema/parameter_schema.py | 25 ++++++++++++++++++-- tool2schema/schema.py | 4 ++-- 3 files changed, 63 insertions(+), 7 deletions(-) diff --git a/tests/test_schema.py b/tests/test_schema.py index 1e083a4..86c0937 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -165,7 +165,9 @@ def remove_param(self, param: str) -> None: :param param: Name of the parameter to remove """ self.schema["function"]["parameters"]["properties"].pop(param) - self.schema["function"]["parameters"]["required"].pop(param, None) + + if required := self.get_required_parameters(): + required.remove(param) def get_param(self, param: str) -> dict: """ @@ -185,6 +187,12 @@ def set_param(self, param, value: dict) -> None: """ self.schema["function"]["parameters"]["properties"][param] = value + def get_required_parameters(self) -> list[str]: + """ + Get the list of required parameters, or none if not present. + """ + return self.schema["function"]["parameters"].get("required") + ########################################### # Example function to test with no tags # @@ -650,8 +658,9 @@ def function_custom_enum(a: CustomEnum, b: str, c: bool = False, d: list[int] = def test_function_custom_enum(): rf = ReferenceSchema(function_custom_enum) - rf.get_param("a")["type"] = "string" - rf.get_param("a")["enum"] = [x.name for x in CustomEnum] + a = rf.get_param("a") + a["type"] = "string" + a["enum"] = [x.name for x in CustomEnum] assert function_custom_enum.schema.to_json() == rf.schema assert function_custom_enum.tags == [] @@ -671,3 +680,29 @@ def test_function_custom_enum(): # Verify it is possible to invoke the function with positional args a, _, _, _ = function_custom_enum(CustomEnum.A, "", False, []) assert a == CustomEnum.A + + +@GPTEnabled +def function_custom_enum_default_value( + a: int, b: CustomEnum = CustomEnum.B, c: bool = False, d: list[int] = [1, 2, 3] +): + """ + This is a test function. + + :param a: This is a parameter + :param b: This is another parameter + :param c: This is a boolean parameter + :param d: This is a list parameter + """ + return a, b, c, d + + +def test_function_custom_enum_default_value(): + rf = ReferenceSchema(function_custom_enum_default_value) + rf.get_required_parameters().remove("b") + b = rf.get_param("b") + b["type"] = "string" + b["default"] = "B" + b["enum"] = [x.name for x in CustomEnum] + assert function_custom_enum_default_value.schema.to_json() == rf.schema + assert function_custom_enum_default_value.tags == [] diff --git a/tool2schema/parameter_schema.py b/tool2schema/parameter_schema.py index 93a3657..95bc690 100644 --- a/tool2schema/parameter_schema.py +++ b/tool2schema/parameter_schema.py @@ -83,6 +83,9 @@ def _get_default(self) -> any: Get the default value for this parameter, when present, to be added to the JSON schema. Return `Parameter.empty` to omit the default value from the schema. """ + if self.parameter.default != Parameter.empty: + return self.encode_value(self.parameter.default) + return self.parameter.default def to_json(self) -> dict: @@ -105,7 +108,17 @@ def to_json(self) -> dict: return json - def parse_value(self, value): + def encode_value(self, value): + """ + Convert the given value to its JSON representation. Overriding methods should + also override `decode_value` to convert the value back to its original type. + + :param value: The value to be converted + :return: The JSON representation of the value + """ + return value + + def decode_value(self, value): """ Convert the given value from the JSON representation to an instance that can be passed to the original method as a parameter. Overriding @@ -216,7 +229,15 @@ def matches(parameter: Parameter) -> bool: and issubclass(parameter.annotation, Enum) ) - def parse_value(self, value): + def encode_value(self, value): + """ + Convert an enum instance to its name. + + :param value: The enum instance to be converted. + """ + return value.name + + def decode_value(self, value): """ Convert an enum name to an instance of the enum type. diff --git a/tool2schema/schema.py b/tool2schema/schema.py index d4eb5b5..b318f42 100644 --- a/tool2schema/schema.py +++ b/tool2schema/schema.py @@ -88,12 +88,12 @@ def __call__(self, *args, **kwargs): for p in self.schema.parameter_schemas.values(): if p.index == i: # Convert the JSON value to the type expected by the method - args[i] = p.parse_value(arg) + args[i] = p.decode_value(arg) for key in kwargs: if key in self.schema.parameter_schemas: # Convert the JSON value to the type expected by the method - kwargs[key] = self.schema.parameter_schemas[key].parse_value(kwargs[key]) + kwargs[key] = self.schema.parameter_schemas[key].decode_value(kwargs[key]) return self.func(*args, **kwargs) From 1354c647debea0725632b1985e4dbcc4eed691d0 Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Thu, 15 Feb 2024 15:51:54 +0100 Subject: [PATCH 11/16] Fix remove_param --- tests/test_schema.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_schema.py b/tests/test_schema.py index 86c0937..9b6c7a3 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -166,7 +166,7 @@ def remove_param(self, param: str) -> None: """ self.schema["function"]["parameters"]["properties"].pop(param) - if required := self.get_required_parameters(): + if (required := self.get_required_parameters()) and param in required: required.remove(param) def get_param(self, param: str) -> dict: From e72566fe361e6594dbd94bbc84645e51dda8d25e Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Sat, 17 Feb 2024 01:56:08 +0100 Subject: [PATCH 12/16] Fix get_required_parameters type hint --- tests/test_schema.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_schema.py b/tests/test_schema.py index 9b6c7a3..9260b71 100644 --- a/tests/test_schema.py +++ b/tests/test_schema.py @@ -187,7 +187,7 @@ def set_param(self, param, value: dict) -> None: """ self.schema["function"]["parameters"]["properties"][param] = value - def get_required_parameters(self) -> list[str]: + def get_required_parameters(self) -> Optional[list[str]]: """ Get the list of required parameters, or none if not present. """ From ad025dff823fb9da127563783f768c81382db0fb Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Sat, 17 Feb 2024 01:59:25 +0100 Subject: [PATCH 13/16] Use dictionary comprehension and check value type Use dictionary comprehension in ParameterSchema.to_json and check the value type in EnumTypeParameterSchema.decode_value --- tool2schema/parameter_schema.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/tool2schema/parameter_schema.py b/tool2schema/parameter_schema.py index 95bc690..26c080e 100644 --- a/tool2schema/parameter_schema.py +++ b/tool2schema/parameter_schema.py @@ -100,11 +100,7 @@ def to_json(self) -> dict: "enum": self._get_enum, } - json = dict() - - for field in fields: - if (value := fields[field]()) != Parameter.empty: - json[field] = value + json = {field: value for field in fields if (value := fields[field]()) != Parameter.empty} return json @@ -235,7 +231,7 @@ def encode_value(self, value): :param value: The enum instance to be converted. """ - return value.name + return value.name if isinstance(value, Enum) else value def decode_value(self, value): """ From 4f5b6637952f9166d30ebab73422f71c12849d77 Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Sat, 17 Feb 2024 12:10:51 +0100 Subject: [PATCH 14/16] Update README --- README.md | 51 ++++++++++++++++++++++++++++++++++++++++++++++++--- 1 file changed, 48 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index 281cb7b..c13eefc 100644 --- a/README.md +++ b/README.md @@ -114,7 +114,53 @@ Any other parameter types will be listed as `object` in the schema. ### Enumerations -If you want to limit the possible values of a parameter, you can use the `enum` keyword argument. +If you want to limit the possible values of a parameter, you can use a `typing.Literal` type hint or a +subclass of `enum.Enum`. For example, using `typing.Literal`: + +```python +import typing + + +@GPTEnabled +def my_function(a: int, b: typing.Literal["yes", "no"]): + """ + Example function description. + + :param a: First parameter + :param b: Second parameter + """ + # Function code here... +``` + +Equivalent example using `enum.Enum`: + +```python +from enum import Enum + +class MyEnum(Enum): + YES = 0 + NO = 1 + + +@GPTEnabled +def my_function(a: int, b: MyEnum): + """ + Example function description. + + :param a: First parameter + :param b: Second parameter + """ + # Function code here... +``` + +In the case of `Enum` subclasses, note that the schema will include the enumeration names rather than the values. +In the example above, the schema will include `["YES", "NO"]` rather than `[0, 1]`. + +The `@GPTEnabled` decorator also allows to invoke the function using the name of the enum member rather than an +instance of the class. For example, you may invoke `my_function(1, MyEnum.YES)` as `my_function(1, "YES")`. + +If the enumeration values are not known at the time of defining the function, +you can add them later using the `add_enum` method. ```python @GPTEnabled @@ -126,11 +172,10 @@ def my_function(a: int, b: str,): :param b: Second parameter """ # Function code here... + my_function.schema.add_enum("b", ["yes", "no"]) ``` -The schema will then be updated to include the `enum` keyword. - ### Tags The `GPTEnabled` decorator also supports the `tags` keyword argument. This allows you to add tags to your function schema. From 4c2151092ed0de00afedf2f61af32b0f9a45072d Mon Sep 17 00:00:00 2001 From: Lorenzo Celli Date: Sat, 17 Feb 2024 16:30:51 +0100 Subject: [PATCH 15/16] Simplify dictionary comprehension Co-authored-by: Angus Stewart --- tool2schema/parameter_schema.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tool2schema/parameter_schema.py b/tool2schema/parameter_schema.py index 26c080e..c2b5a7e 100644 --- a/tool2schema/parameter_schema.py +++ b/tool2schema/parameter_schema.py @@ -93,14 +93,14 @@ def to_json(self) -> dict: Return the json schema for this parameter. """ fields = { - "description": self._get_description, - "default": self._get_default, - "items": self._get_items, - "type": self._get_type, - "enum": self._get_enum, + "description": self._get_description(), + "default": self._get_default(), + "items": self._get_items(), + "type": self._get_type(), + "enum": self._get_enum(), } - json = {field: value for field in fields if (value := fields[field]()) != Parameter.empty} + json = {f: v for f, v in fields.items() if v != Parameter.empty} return json From 6d659b497ff3b2d4ad6e8197b682c2ed931f0357 Mon Sep 17 00:00:00 2001 From: LorenzoCelli Date: Sat, 17 Feb 2024 16:38:06 +0100 Subject: [PATCH 16/16] Add comment in ParameterSchema._get_default --- tool2schema/parameter_schema.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tool2schema/parameter_schema.py b/tool2schema/parameter_schema.py index c2b5a7e..7e67253 100644 --- a/tool2schema/parameter_schema.py +++ b/tool2schema/parameter_schema.py @@ -86,7 +86,9 @@ def _get_default(self) -> any: if self.parameter.default != Parameter.empty: return self.encode_value(self.parameter.default) - return self.parameter.default + # Not that the default value may be present but None, we use + # Parameter.empty to indicate that the default value is not present + return Parameter.empty def to_json(self) -> dict: """