Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Feature/support enumeration type hints #37

Merged
merged 20 commits into from
Feb 17, 2024
Merged

Feature/support enumeration type hints #37

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
20 commits
Select commit Hold shift + click to select a range
f3ee51f
Add support for typing.Literal #10
lorenzocelli Feb 8, 2024
aa8da6b
Add support for enum.Enum #10
lorenzocelli Feb 8, 2024
0347fa9
Merge branch 'main' into feature/support-enumeration-type-hints
lorenzocelli Feb 10, 2024
15ecfdf
Add LiteralParameterSchema and EnumParameterSchema
lorenzocelli Feb 10, 2024
b8f27c9
Merge branch 'main' into feature/support-enumeration-type-hints
lorenzocelli Feb 11, 2024
f61c784
Use ReferenceSchema for enum tests
lorenzocelli Feb 11, 2024
0ceab9d
Store ParameterSchema instances in FunctionSchema and convert enum va…
lorenzocelli Feb 11, 2024
2471860
Convert enum to names rather than values
lorenzocelli Feb 11, 2024
cb80329
Merge branch 'main' into feature/support-enumeration-type-hints
lorenzocelli Feb 11, 2024
cda7b51
Merge branch 'main' into feature/support-enumeration-type-hints
lorenzocelli Feb 12, 2024
3a9fd4a
Avoid populating the schema before calling to_json
lorenzocelli Feb 12, 2024
4f42b4c
Refactor ParameterSchema add-to logic to get logic
lorenzocelli Feb 15, 2024
6057e6f
Support positional arguments value parsing
lorenzocelli Feb 15, 2024
a420a7d
Support enum.Enum parameter with default value
lorenzocelli Feb 15, 2024
1354c64
Fix remove_param
lorenzocelli Feb 15, 2024
e72566f
Fix get_required_parameters type hint
lorenzocelli Feb 17, 2024
ad025df
Use dictionary comprehension and check value type
lorenzocelli Feb 17, 2024
4f5b663
Update README
lorenzocelli Feb 17, 2024
4c21510
Simplify dictionary comprehension
lorenzocelli Feb 17, 2024
6d659b4
Add comment in ParameterSchema._get_default
lorenzocelli Feb 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
247 changes: 243 additions & 4 deletions tests/test_schema.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from typing import List, Optional
from enum import Enum
from typing import List, Optional, Literal

from tool2schema import (
FindGPTEnabled,
Expand Down Expand Up @@ -288,9 +289,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
Expand Down Expand Up @@ -816,3 +817,241 @@ 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 == []


#################################################
# 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 == []
14 changes: 12 additions & 2 deletions tool2schema/schema.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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;
"""
Expand Down Expand Up @@ -232,6 +231,17 @@ 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 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:
Expand Down
Loading