jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

fix: update type extraction to extract inner types correctly (#3446)

Browse source 
* refactor: Improve post-processing of return type in type extraction.

* feat: Add support for Union[int, Sequence[str]] in post_process_type.

* refactor(langflow): Simplify return type extraction in CustomComponent.

* Refactor: Remove 'Sequence' and 'list' type options in starter projects.

* test: fix assertion

---------

Co-authored-by: italojohnny <italojohnnydosanjos@gmail.com>
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-08-20 09:48:18 -03:00 • committed by GitHub
commit 037c0902d5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 18 additions and 37 deletions
Download patch file Download diff file Expand all files Collapse all files

19
src/backend/base/langflow/custom/custom_component/custom_component.py
View file

@ -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):

6
src/backend/base/langflow/initial_setup/starter_projects/Complex Agent.json
View file

@ -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__"
}

6
src/backend/base/langflow/initial_setup/starter_projects/Hierarchical Agent.json
View file

@ -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__"
}

6
src/backend/base/langflow/initial_setup/starter_projects/Sequential Agent.json
View file

@ -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__"
}

12
src/backend/base/langflow/type_extraction/type_extraction.py
View file

@ -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:

6
src/backend/tests/test_schema.py
View file

@ -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}