From b3201ce8d77d783917ae7e2535ca5069a87bce81 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Tue, 20 Feb 2024 14:32:03 -0300 Subject: [PATCH] add error handling for missing vertices --- src/backend/langflow/graph/graph/base.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/backend/langflow/graph/graph/base.py b/src/backend/langflow/graph/graph/base.py index 3c772c897..2c1d7eaf3 100644 --- a/src/backend/langflow/graph/graph/base.py +++ b/src/backend/langflow/graph/graph/base.py @@ -7,12 +7,8 @@ from langflow.graph.graph.constants import lazy_load_vertex_dict from langflow.graph.graph.utils import process_flow from langflow.graph.schema import InterfaceComponentTypes from langflow.graph.vertex.base import Vertex -from langflow.graph.vertex.types import ( - ChatVertex, - FileToolVertex, - LLMVertex, - ToolkitVertex, -) +from langflow.graph.vertex.types import (ChatVertex, FileToolVertex, LLMVertex, + ToolkitVertex) from langflow.interface.tools.constants import FILE_TOOLS from langflow.utils import payload from loguru import logger @@ -216,7 +212,10 @@ class Graph: def get_vertex(self, vertex_id: str) -> Union[None, Vertex]: """Returns a vertex by id.""" - return self.vertex_map.get(vertex_id) + try: + return self.vertex_map.get(vertex_id) + except KeyError: + raise ValueError(f"Vertex {vertex_id} not found") def get_vertex_edges(self, vertex_id: str) -> List[ContractEdge]: """Returns a list of edges for a given vertex."""