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

Reorder Vertex class assignment

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-02-25 12:17:47 -03:00
commit 4aefcb03a7
1 changed files with 5 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -13,6 +13,7 @@ from langflow.graph.vertex.types import (
ChatVertex,
FileToolVertex,
LLMVertex,
RoutingVertex,
ToolkitVertex,
)
from langflow.interface.tools.constants import FILE_TOOLS
@ -402,10 +403,12 @@ class Graph:
node_name = node_id.split("-")[0]
if node_name in ["ChatOutput", "ChatInput"]:
return ChatVertex
if node_name in lazy_load_vertex_dict.VERTEX_TYPE_MAP:
return lazy_load_vertex_dict.VERTEX_TYPE_MAP[node_name]
elif node_name in ["ShouldRunNext"]:
return RoutingVertex
elif node_base_type in lazy_load_vertex_dict.VERTEX_TYPE_MAP:
return lazy_load_vertex_dict.VERTEX_TYPE_MAP[node_base_type]
elif node_name in lazy_load_vertex_dict.VERTEX_TYPE_MAP:
return lazy_load_vertex_dict.VERTEX_TYPE_MAP[node_name]
if node_type in FILE_TOOLS:
return FileToolVertex