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

Refactor _get_vertex_class method to handle node_base_type

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-07 14:07:13 -03:00
commit 0abeda51e9
1 changed files with 8 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

11
src/backend/langflow/graph/graph/base.py
View file

@ -237,8 +237,12 @@ class Graph:
edges.append(ContractEdge(source, target, edge)) edges.append(ContractEdge(source, target, edge))
return edges return edges
def _get_vertex_class(self, node_type: str, node_lc_type: str, node_id: str) -> Type[Vertex]: def _get_vertex_class(self, node_type: str, node_base_type: str, node_id: str) -> Type[Vertex]:
"""Returns the node class based on the node type.""" """Returns the node class based on the node type."""
# First we check for the node_base_type
if node_base_type in lazy_load_vertex_dict.VERTEX_TYPE_MAP:
return lazy_load_vertex_dict.VERTEX_TYPE_MAP[node_base_type]
node_name = node_id.split("-")[0] node_name = node_id.split("-")[0]
if node_name in lazy_load_vertex_dict.VERTEX_TYPE_MAP: if node_name in lazy_load_vertex_dict.VERTEX_TYPE_MAP:
return lazy_load_vertex_dict.VERTEX_TYPE_MAP[node_name] return lazy_load_vertex_dict.VERTEX_TYPE_MAP[node_name]
@ -248,8 +252,8 @@ class Graph:
if node_type in lazy_load_vertex_dict.VERTEX_TYPE_MAP: if node_type in lazy_load_vertex_dict.VERTEX_TYPE_MAP:
return lazy_load_vertex_dict.VERTEX_TYPE_MAP[node_type] return lazy_load_vertex_dict.VERTEX_TYPE_MAP[node_type]
return ( return (
lazy_load_vertex_dict.VERTEX_TYPE_MAP[node_lc_type] lazy_load_vertex_dict.VERTEX_TYPE_MAP[node_base_type]
if node_lc_type in lazy_load_vertex_dict.VERTEX_TYPE_MAP if node_base_type in lazy_load_vertex_dict.VERTEX_TYPE_MAP
else Vertex else Vertex
) )
@ -312,3 +316,4 @@ class Graph:
return layers return layers
return layers return layers
return layers return layers
return layers