From 4aefcb03a79b635978c846b0a3532aed1038bca9 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sun, 25 Feb 2024 12:17:47 -0300 Subject: [PATCH] Reorder Vertex class assignment --- src/backend/langflow/graph/graph/base.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/backend/langflow/graph/graph/base.py b/src/backend/langflow/graph/graph/base.py index 3535c0e2f..453d7a98f 100644 --- a/src/backend/langflow/graph/graph/base.py +++ b/src/backend/langflow/graph/graph/base.py @@ -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