diff --git a/src/backend/langflow/interface/custom/utils.py b/src/backend/langflow/interface/custom/utils.py index 700b98b53..11a155a00 100644 --- a/src/backend/langflow/interface/custom/utils.py +++ b/src/backend/langflow/interface/custom/utils.py @@ -1,4 +1,6 @@ import re +from types import GenericAlias +from typing import Any def extract_inner_type(return_type: str) -> str: @@ -10,6 +12,23 @@ def extract_inner_type(return_type: str) -> str: return return_type +def extract_inner_type_from_generic_alias(return_type: GenericAlias) -> Any: + """ + Extracts the inner type from a type hint that is a list. + """ + if return_type.__origin__ == list: + return list(return_type.__args__) + + return return_type + + +def extract_union_types_from_generic_alias(return_type: GenericAlias) -> tuple: + """ + Extracts the inner type from a type hint that is a Union. + """ + return return_type.__args__ + + def extract_union_types(return_type: str) -> list[str]: """ Extracts the inner type from a type hint that is a list.