From ca749746ad3629f836ddb1f723bb46d3caf3b13e Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 23 Feb 2024 10:08:34 -0300 Subject: [PATCH] Refactor vertex type lookup logic in Graph class --- src/backend/langflow/graph/graph/base.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/backend/langflow/graph/graph/base.py b/src/backend/langflow/graph/graph/base.py index c44ad8afb..6e2d60bea 100644 --- a/src/backend/langflow/graph/graph/base.py +++ b/src/backend/langflow/graph/graph/base.py @@ -400,11 +400,10 @@ class Graph: node_name = node_id.split("-")[0] if node_name in ["ChatOutput", "ChatInput"]: return ChatVertex - if node_base_type in lazy_load_vertex_dict.VERTEX_TYPE_MAP: - return lazy_load_vertex_dict.VERTEX_TYPE_MAP[node_base_type] - if node_name in lazy_load_vertex_dict.VERTEX_TYPE_MAP: return lazy_load_vertex_dict.VERTEX_TYPE_MAP[node_name] + elif node_base_type in lazy_load_vertex_dict.VERTEX_TYPE_MAP: + return lazy_load_vertex_dict.VERTEX_TYPE_MAP[node_base_type] if node_type in FILE_TOOLS: return FileToolVertex