From 98bacf5f74f8549c9d46febd507ad277e27ef337 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 22 Nov 2023 21:12:21 -0300 Subject: [PATCH] Refactor graph base module --- src/backend/langflow/graph/graph/base.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/src/backend/langflow/graph/graph/base.py b/src/backend/langflow/graph/graph/base.py index aeb7ec91d..01856d517 100644 --- a/src/backend/langflow/graph/graph/base.py +++ b/src/backend/langflow/graph/graph/base.py @@ -1,18 +1,14 @@ from typing import Dict, Generator, List, Type, Union +from langchain.chains.base import Chain from langflow.graph.edge.base import Edge from langflow.graph.graph.constants import lazy_load_vertex_dict from langflow.graph.graph.utils import process_flow from langflow.graph.vertex.base import Vertex -from langflow.graph.vertex.types import ( - FileToolVertex, - LLMVertex, - ToolkitVertex, -) +from langflow.graph.vertex.types import FileToolVertex, LLMVertex, ToolkitVertex from langflow.interface.tools.constants import FILE_TOOLS from langflow.utils import payload from loguru import logger -from langchain.chains.base import Chain class Graph: @@ -31,8 +27,8 @@ class Graph: for node in self._nodes: if node_id := node.get("id"): self.top_level_nodes.append(node_id) - self._graph_data = process_flow(self.raw_graph_data) + self._nodes = self._graph_data["nodes"] self._edges = self._graph_data["edges"] self._build_graph()