diff --git a/src/backend/base/langflow/custom/custom_component/custom_component.py b/src/backend/base/langflow/custom/custom_component/custom_component.py index 3393e356d..1d428932f 100644 --- a/src/backend/base/langflow/custom/custom_component/custom_component.py +++ b/src/backend/base/langflow/custom/custom_component/custom_component.py @@ -17,10 +17,7 @@ from langflow.services.deps import get_storage_service, get_variable_service, se from langflow.services.storage.service import StorageService from langflow.services.tracing.schema import Log from langflow.template.utils import update_frontend_node_with_template_values -from langflow.type_extraction.type_extraction import ( - extract_inner_type_from_generic_alias, - extract_union_types_from_generic_alias, -) +from langflow.type_extraction.type_extraction import post_process_type from langflow.utils import validate if TYPE_CHECKING: @@ -365,19 +362,7 @@ class CustomComponent(BaseComponent): return self.get_method_return_type(self._function_entrypoint_name) def _extract_return_type(self, return_type: Any) -> List[Any]: - if hasattr(return_type, "__origin__") and return_type.__origin__ in [ - list, - List, - ]: - return_type = extract_inner_type_from_generic_alias(return_type) - - # If the return type is not a Union, then we just return it as a list - inner_type = return_type[0] if isinstance(return_type, list) else return_type - if not hasattr(inner_type, "__origin__") or inner_type.__origin__ != Union: - return return_type if isinstance(return_type, list) else [return_type] - # If the return type is a Union, then we need to parse it - return_type = extract_union_types_from_generic_alias(return_type) - return return_type + return post_process_type(return_type) @property def get_main_class_name(self): diff --git a/src/backend/base/langflow/initial_setup/starter_projects/Complex Agent.json b/src/backend/base/langflow/initial_setup/starter_projects/Complex Agent.json index dea218fc0..100013bf4 100644 --- a/src/backend/base/langflow/initial_setup/starter_projects/Complex Agent.json +++ b/src/backend/base/langflow/initial_setup/starter_projects/Complex Agent.json @@ -4037,8 +4037,7 @@ "name": "api_run_model", "selected": "Data", "types": [ - "Data", - "list" + "Data" ], "value": "__UNDEFINED__" }, @@ -4049,8 +4048,7 @@ "name": "api_build_tool", "selected": "Tool", "types": [ - "Tool", - "Sequence" + "Tool" ], "value": "__UNDEFINED__" } diff --git a/src/backend/base/langflow/initial_setup/starter_projects/Hierarchical Agent.json b/src/backend/base/langflow/initial_setup/starter_projects/Hierarchical Agent.json index c8433504f..b23483baf 100644 --- a/src/backend/base/langflow/initial_setup/starter_projects/Hierarchical Agent.json +++ b/src/backend/base/langflow/initial_setup/starter_projects/Hierarchical Agent.json @@ -2615,8 +2615,7 @@ "name": "api_run_model", "selected": "Data", "types": [ - "Data", - "list" + "Data" ], "value": "__UNDEFINED__" }, @@ -2627,8 +2626,7 @@ "name": "api_build_tool", "selected": "Tool", "types": [ - "Tool", - "Sequence" + "Tool" ], "value": "__UNDEFINED__" } diff --git a/src/backend/base/langflow/initial_setup/starter_projects/Sequential Agent.json b/src/backend/base/langflow/initial_setup/starter_projects/Sequential Agent.json index e73aad988..b46cfe2ee 100644 --- a/src/backend/base/langflow/initial_setup/starter_projects/Sequential Agent.json +++ b/src/backend/base/langflow/initial_setup/starter_projects/Sequential Agent.json @@ -2953,8 +2953,7 @@ "name": "api_run_model", "selected": "Data", "types": [ - "Data", - "list" + "Data" ], "value": "__UNDEFINED__" }, @@ -2965,8 +2964,7 @@ "name": "api_build_tool", "selected": "Tool", "types": [ - "Tool", - "Sequence" + "Tool" ], "value": "__UNDEFINED__" } diff --git a/src/backend/base/langflow/type_extraction/type_extraction.py b/src/backend/base/langflow/type_extraction/type_extraction.py index f22a8eebe..6474977a3 100644 --- a/src/backend/base/langflow/type_extraction/type_extraction.py +++ b/src/backend/base/langflow/type_extraction/type_extraction.py @@ -1,4 +1,6 @@ import re +from collections.abc import Sequence as SequenceABC +from itertools import chain from types import GenericAlias from typing import Any, List, Union @@ -7,7 +9,7 @@ def extract_inner_type_from_generic_alias(return_type: GenericAlias) -> Any: """ Extracts the inner type from a type hint that is a list or a Optional. """ - if return_type.__origin__ == list: + if return_type.__origin__ in [list, SequenceABC]: return list(return_type.__args__) return return_type @@ -57,10 +59,7 @@ def post_process_type(_type): Union[List[Any], Any]: The processed return type. """ - if hasattr(_type, "__origin__") and _type.__origin__ in [ - list, - List, - ]: + if hasattr(_type, "__origin__") and _type.__origin__ in [list, List, SequenceABC]: _type = extract_inner_type_from_generic_alias(_type) # If the return type is not a Union, then we just return it as a list @@ -69,7 +68,8 @@ def post_process_type(_type): return _type if isinstance(_type, list) else [_type] # If the return type is a Union, then we need to parse it _type = extract_union_types_from_generic_alias(_type) - return _type + _type = set(chain.from_iterable([post_process_type(t) for t in _type])) + return list(_type) def extract_union_types_from_generic_alias(return_type: GenericAlias) -> list: diff --git a/src/backend/tests/test_schema.py b/src/backend/tests/test_schema.py index 1a2b69c34..75c2002d6 100644 --- a/src/backend/tests/test_schema.py +++ b/src/backend/tests/test_schema.py @@ -1,4 +1,4 @@ -from typing import Union +from typing import Sequence, Union import pytest from pydantic import ValidationError @@ -42,6 +42,8 @@ class TestInput: assert post_process_type(int) == [int] assert post_process_type(list[int]) == [int] assert post_process_type(Union[int, str]) == [int, str] + assert post_process_type(Union[int, Sequence[str]]) == [int, str] + assert post_process_type(Union[int, Sequence[int]]) == [int] def test_input_to_dict(self): input_obj = Input(field_type="str") @@ -126,4 +128,4 @@ class TestPostProcessType: class CustomType: pass - assert post_process_type(Union[CustomType, int]) == [CustomType, int] + assert set(post_process_type(Union[CustomType, int])) == {CustomType, int}