From ae3b02f79e8ba2f5dc77b8446ab42d03f636a382 Mon Sep 17 00:00:00 2001 From: Eric Gustin <34000337+EricGustin@users.noreply.github.com> Date: Mon, 7 Apr 2025 09:52:53 -0800 Subject: [PATCH] Support PEP 604 Union Types (#347) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ⚠️ Local Worker will not work with Slack toolkit until this is merged  Support [PEP 604](https://peps.python.org/pep-0604/) Union Types for tool input param types and return types. https://github.com/ArcadeAI/arcade-ai/pull/345 introduced PEP 604 Union Types into our toolkits, but the worker did not support this syntax, so it was failing on start up if a tool used bar syntax in its return type (Slack). Additionally, optionals defined with `Union[T, None]` were already supported, but didn't have any unit tests, so I added them. --- arcade/arcade/core/catalog.py | 23 ++- arcade/arcade/core/utils.py | 8 + arcade/tests/core/test_catalog.py | 8 +- .../core/utils/test_is_strict_optional.py | 26 ++++ .../tests/tool/test_create_tool_definition.py | 146 +++++++++++++++++- .../test_create_tool_definition_errors.py | 38 ++++- .../test_create_tool_definition_pydantic.py | 50 +++++- ..._create_tool_definition_pydantic_errors.py | 26 +++- 8 files changed, 296 insertions(+), 29 deletions(-) create mode 100644 arcade/tests/core/utils/test_is_strict_optional.py diff --git a/arcade/arcade/core/catalog.py b/arcade/arcade/core/catalog.py index 2e2fc653..1fc6e6b4 100644 --- a/arcade/arcade/core/catalog.py +++ b/arcade/arcade/core/catalog.py @@ -48,6 +48,7 @@ from arcade.core.toolkit import Toolkit from arcade.core.utils import ( does_function_return_value, first_or_none, + is_strict_optional, is_string_literal, is_union, snake_to_pascal_case, @@ -481,10 +482,10 @@ def create_output_definition(func: Callable) -> ToolOutput: return_type = return_type.__origin__ # Unwrap Optional types - is_optional = False - if get_origin(return_type) is Union and type(None) in get_args(return_type): + # Both Optional[T] and T | None are supported + is_optional = is_strict_optional(return_type) + if is_optional: return_type = next(arg for arg in get_args(return_type) if arg is not type(None)) - is_optional = True wire_type_info = get_wire_type_info(return_type) @@ -659,14 +660,9 @@ def extract_python_param_info(param: inspect.Parameter) -> ParamInfo: # Handle optional types # Both Optional[T] and T | None are supported - is_optional = False - if ( - is_union(field_type) - and len(get_args(field_type)) == 2 - and type(None) in get_args(field_type) - ): + is_optional = is_strict_optional(field_type) + if is_optional: field_type = next(arg for arg in get_args(field_type) if arg is not type(None)) - is_optional = True # Union types are not currently supported # (other than optional, which is handled above) @@ -703,10 +699,10 @@ def extract_pydantic_param_info(param: inspect.Parameter) -> ParamInfo: field_type = original_type # Unwrap Optional types - is_optional = False - if get_origin(field_type) is Union and type(None) in get_args(field_type): + # Both Optional[T] and T | None are supported + is_optional = is_strict_optional(field_type) + if is_optional: field_type = next(arg for arg in get_args(field_type) if arg is not type(None)) - is_optional = True return ParamInfo( name=param.name, @@ -732,7 +728,6 @@ def get_wire_type( float: "number", dict: "json", } - outer_type_mapping: dict[type, WireType] = { list: "array", dict: "json", diff --git a/arcade/arcade/core/utils.py b/arcade/arcade/core/utils.py index 0e4c2eaf..a199f93d 100644 --- a/arcade/arcade/core/utils.py +++ b/arcade/arcade/core/utils.py @@ -52,6 +52,14 @@ def is_union(_type: type) -> bool: return get_origin(_type) in {Union, UnionType} +def is_strict_optional(_type: type) -> bool: + """ + Returns True if the given type is a strict optional type, i.e. a union with exactly two types + where one type is None. This covers Optional[T], Union[T, None] and T | None + """ + return is_union(_type) and len(get_args(_type)) == 2 and type(None) in get_args(_type) + + def does_function_return_value(func: Callable) -> bool: """ Returns True if the given function returns a value, i.e. if it has a return statement with a value. diff --git a/arcade/tests/core/test_catalog.py b/arcade/tests/core/test_catalog.py index a2a0bb7a..49f86753 100644 --- a/arcade/tests/core/test_catalog.py +++ b/arcade/tests/core/test_catalog.py @@ -94,14 +94,10 @@ def test_add_toolkit_type_error(): mock_import.return_value = MagicMock() mock_getattr.return_value = InvalidTool() - # Assert that ToolDefinitionError is raised with the correct message - with pytest.raises(ToolDefinitionError) as exc_info: + # Assert that ToolDefinitionError is raised + with pytest.raises(ToolDefinitionError): catalog.add_toolkit(mock_toolkit) - assert "Type error encountered while adding tool invalid_tool from mock_module" in str( - exc_info.value - ) - def test_get_tool_by_name(): catalog = ToolCatalog() diff --git a/arcade/tests/core/utils/test_is_strict_optional.py b/arcade/tests/core/utils/test_is_strict_optional.py new file mode 100644 index 00000000..fdf6ccb5 --- /dev/null +++ b/arcade/tests/core/utils/test_is_strict_optional.py @@ -0,0 +1,26 @@ +from typing import Optional, Union + +import pytest + +from arcade.core.utils import is_strict_optional + + +@pytest.mark.parametrize( + "type_input, expected", + [ + (Union[int, None], True), + (Union[int, str, None], False), + (Union[int, str], False), + (Optional[int], True), + (int | None, True), + (None | int, True), + (int | str, False), + (int | str | None, False), + (int, False), + (str, False), + (list, False), + (dict, False), + ], +) +def test_is_optional_type(type_input, expected): + assert is_strict_optional(type_input) == expected diff --git a/arcade/tests/tool/test_create_tool_definition.py b/arcade/tests/tool/test_create_tool_definition.py index bc94bed0..539a3727 100644 --- a/arcade/tests/tool/test_create_tool_definition.py +++ b/arcade/tests/tool/test_create_tool_definition.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import Annotated, Literal, Optional +from typing import Annotated, Literal, Optional, Union import pytest @@ -172,12 +172,33 @@ def func_with_optional_param_with_default_value( @tool(desc="A function with an optional input parameter with bar syntax") -def func_with_optional_param_with_bar_syntax( +def func_with_optional_param_with_bar_syntax_1( param1: Annotated[str | None, "First param"] = None, ): pass +@tool(desc="A function with an optional input parameter with bar syntax") +def func_with_optional_param_with_bar_syntax_2( + param1: Annotated[None | str, "First param"] = None, +): + pass + + +@tool(desc="A function with an optional input parameter with union syntax") +def func_with_optional_param_with_union_syntax_1( + param1: Annotated[Union[str, None], "First param"] = None, +): + pass + + +@tool(desc="A function with an optional input parameter with union syntax") +def func_with_optional_param_with_union_syntax_2( + param1: Annotated[Union[None, str], "First param"] = None, +): + pass + + @tool(desc="A function with multiple parameters, some with default values") def func_with_mixed_params( context: ToolContext, @@ -240,6 +261,26 @@ def func_with_optional_return() -> Optional[str]: return "maybe output" +@tool(desc="A function with an optional return type that uses bar syntax") +def func_with_optional_return_with_bar_syntax_1() -> str | None: + return "maybe output" + + +@tool(desc="A function with an optional return type that uses bar syntax") +def func_with_optional_return_with_bar_syntax_2() -> None | str: + return "maybe output" + + +@tool(desc="A function with an optional return type that uses union syntax") +def func_with_optional_return_with_union_syntax_1() -> Union[str, None]: + return "maybe output" + + +@tool(desc="A function with an optional return type that uses union syntax") +def func_with_optional_return_with_union_syntax_2() -> Union[None, str]: + return "maybe output" + + @tool(desc="A function with a complex return type") def func_with_complex_return() -> dict[str, str]: return [{"key": "value"}] @@ -536,7 +577,7 @@ def func_with_complex_return() -> dict[str, str]: id="func_with_optional_param_with_default_value", ), pytest.param( - func_with_optional_param_with_bar_syntax, + func_with_optional_param_with_bar_syntax_1, { "input": ToolInput( parameters=[ @@ -555,6 +596,57 @@ def func_with_complex_return() -> dict[str, str]: }, id="func_with_optional_param_with_bar_syntax", ), + pytest.param( + func_with_optional_param_with_bar_syntax_2, + { + "input": ToolInput( + parameters=[ + InputParameter( + name="param1", + description="First param", + inferrable=True, + required=False, + value_schema=ValueSchema(val_type="string", enum=None), + ) + ] + ), + }, + id="func_with_optional_param_with_bar_syntax_2", + ), + pytest.param( + func_with_optional_param_with_union_syntax_1, + { + "input": ToolInput( + parameters=[ + InputParameter( + name="param1", + description="First param", + inferrable=True, + required=False, + value_schema=ValueSchema(val_type="string", enum=None), + ) + ] + ), + }, + id="func_with_optional_param_with_union_syntax_1", + ), + pytest.param( + func_with_optional_param_with_union_syntax_2, + { + "input": ToolInput( + parameters=[ + InputParameter( + name="param1", + description="First param", + inferrable=True, + required=False, + value_schema=ValueSchema(val_type="string", enum=None), + ) + ] + ), + }, + id="func_with_optional_param_with_union_syntax_2", + ), pytest.param( func_with_mixed_params, { @@ -724,6 +816,54 @@ def func_with_complex_return() -> dict[str, str]: }, id="func_with_optional_return", ), + pytest.param( + func_with_optional_return_with_bar_syntax_1, + { + "input": ToolInput(parameters=[]), + "output": ToolOutput( + value_schema=ValueSchema(val_type="string", enum=None), + available_modes=["value", "error", "null"], + description="No description provided.", + ), + }, + id="func_with_optional_return_with_bar_syntax_1", + ), + pytest.param( + func_with_optional_return_with_bar_syntax_2, + { + "input": ToolInput(parameters=[]), + "output": ToolOutput( + value_schema=ValueSchema(val_type="string", enum=None), + available_modes=["value", "error", "null"], + description="No description provided.", + ), + }, + id="func_with_optional_return_with_bar_syntax_2", + ), + pytest.param( + func_with_optional_return_with_union_syntax_1, + { + "input": ToolInput(parameters=[]), + "output": ToolOutput( + value_schema=ValueSchema(val_type="string", enum=None), + available_modes=["value", "error", "null"], + description="No description provided.", + ), + }, + id="func_with_optional_return_with_union_syntax_1", + ), + pytest.param( + func_with_optional_return_with_union_syntax_2, + { + "input": ToolInput(parameters=[]), + "output": ToolOutput( + value_schema=ValueSchema(val_type="string", enum=None), + available_modes=["value", "error", "null"], + description="No description provided.", + ), + }, + id="func_with_optional_return_with_union_syntax_2", + ), pytest.param( func_with_complex_return, { diff --git a/arcade/tests/tool/test_create_tool_definition_errors.py b/arcade/tests/tool/test_create_tool_definition_errors.py index 5e6d10bf..d2848b1c 100644 --- a/arcade/tests/tool/test_create_tool_definition_errors.py +++ b/arcade/tests/tool/test_create_tool_definition_errors.py @@ -1,4 +1,4 @@ -from typing import Annotated +from typing import Annotated, Union import pytest @@ -18,6 +18,16 @@ def func_with_missing_return_type(): return "hello world" +@tool(desc="A function with a union return type (illegal)") +def func_with_union_return_type_1() -> Union[str, int]: + return "hello world" + + +@tool(desc="A function with a union return type (illegal)") +def func_with_union_return_type_2() -> str | int: + return "hello world" + + @tool(desc="A function with a parameter type (illegal)") def func_with_missing_param_type(param1): pass @@ -34,7 +44,12 @@ def func_with_unsupported_param(param1: complex): @tool(desc="A function with a union parameter (illegal)") -def func_with_union_param(param1: str | int): +def func_with_union_param_1(param1: str | int): + pass + + +@tool(desc="A function with a union parameter (illegal)") +def func_with_union_param_2(param1: Union[str, int]): pass @@ -95,9 +110,14 @@ def func_with_secret_requirement_invalid_type(): id=func_with_unsupported_param.__name__, ), pytest.param( - func_with_union_param, + func_with_union_param_1, ToolDefinitionError, - id=func_with_union_param.__name__, + id=func_with_union_param_1.__name__, + ), + pytest.param( + func_with_union_param_2, + ToolDefinitionError, + id=func_with_union_param_2.__name__, ), pytest.param( func_with_multiple_context_params, @@ -119,6 +139,16 @@ def func_with_secret_requirement_invalid_type(): ToolDefinitionError, id=func_with_secret_requirement_invalid_type.__name__, ), + pytest.param( + func_with_union_return_type_1, + ToolDefinitionError, + id=func_with_union_return_type_1.__name__, + ), + pytest.param( + func_with_union_return_type_2, + ToolDefinitionError, + id=func_with_union_return_type_2.__name__, + ), ], ) def test_missing_info_raises_error(func_under_test, exception_type): diff --git a/arcade/tests/tool/test_create_tool_definition_pydantic.py b/arcade/tests/tool/test_create_tool_definition_pydantic.py index 0fc32542..e423501f 100644 --- a/arcade/tests/tool/test_create_tool_definition_pydantic.py +++ b/arcade/tests/tool/test_create_tool_definition_pydantic.py @@ -35,13 +35,27 @@ def func_takes_pydantic_field_with_description( return product_name -@tool(desc="A function that accepts an Pydantic Field") +@tool(desc="A function that accepts an optional Pydantic Field") def func_takes_pydantic_field_optional( product_name: Optional[str] = Field(None, description="The name of the product"), ) -> str: return product_name +@tool(desc="A function that accepts an optional Pydantic Field with bar syntax") +def func_takes_pydantic_field_optional_bar_syntax( + product_name: str | None = Field(None, description="The name of the product"), +) -> str: + return product_name + + +@tool(desc="A function that accepts an optional Pydantic Field with union syntax") +def func_takes_pydantic_field_optional_union_syntax( + product_name: Union[str, None] = Field(None, description="The name of the product"), +) -> str: + return product_name + + # Annotated[] takes precedence over Field() properties @tool(desc="A function that accepts an annotated Pydantic Field") def func_takes_pydantic_field_annotated_description( @@ -174,6 +188,40 @@ def read_products( }, id="func_takes_pydantic_field_optional", ), + pytest.param( + func_takes_pydantic_field_optional_bar_syntax, + { + "input": ToolInput( + parameters=[ + InputParameter( + name="product_name", + description="The name of the product", + required=False, + inferrable=True, + value_schema=ValueSchema(val_type="string", enum=None), + ) + ] + ) + }, + id="func_takes_pydantic_field_optional_bar_syntax", + ), + pytest.param( + func_takes_pydantic_field_optional_union_syntax, + { + "input": ToolInput( + parameters=[ + InputParameter( + name="product_name", + description="The name of the product", + required=False, + inferrable=True, + value_schema=ValueSchema(val_type="string", enum=None), + ) + ] + ) + }, + id="func_takes_pydantic_field_optional_union_syntax", + ), pytest.param( func_takes_pydantic_field_annotated_description, { diff --git a/arcade/tests/tool/test_create_tool_definition_pydantic_errors.py b/arcade/tests/tool/test_create_tool_definition_pydantic_errors.py index e05fb92a..f49a9090 100644 --- a/arcade/tests/tool/test_create_tool_definition_pydantic_errors.py +++ b/arcade/tests/tool/test_create_tool_definition_pydantic_errors.py @@ -1,4 +1,4 @@ -from typing import Annotated +from typing import Annotated, Union import pytest from pydantic import Field @@ -21,6 +21,20 @@ def field_with_literal_default_factory( pass +@tool(desc="A function that accepts an optional Pydantic Field with non-strict optional syntax") +def func_takes_pydantic_field_non_strict_optional_bar_syntax( + product_name: str | int | None = Field(None, description="The name of the product"), +) -> str: + return product_name + + +@tool(desc="A function that accepts an optional Pydantic Field with non-strict optional syntax") +def func_takes_pydantic_field_non_strict_optional_union_syntax( + product_name: Union[str, int, None] = Field(None, description="The name of the product"), +) -> str: + return product_name + + @pytest.mark.parametrize( "func_under_test, exception_type", [ @@ -29,6 +43,16 @@ def field_with_literal_default_factory( ToolDefinitionError, id=field_with_literal_default_factory.__name__, ), + pytest.param( + func_takes_pydantic_field_non_strict_optional_bar_syntax, + ToolDefinitionError, + id=func_takes_pydantic_field_non_strict_optional_bar_syntax.__name__, + ), + pytest.param( + func_takes_pydantic_field_non_strict_optional_union_syntax, + ToolDefinitionError, + id=func_takes_pydantic_field_non_strict_optional_union_syntax.__name__, + ), ], ) def test_missing_info_raises_error(func_under_test, exception_type):