From 4e5181307fa25237bfb8b9ddde6ce54bccf96b61 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 25 Jan 2024 16:15:34 -0300 Subject: [PATCH] Add logger import and update edge type in Graph class --- src/backend/langflow/graph/graph/base.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/graph/graph/base.py b/src/backend/langflow/graph/graph/base.py index ffce2a2e7..17c402c51 100644 --- a/src/backend/langflow/graph/graph/base.py +++ b/src/backend/langflow/graph/graph/base.py @@ -2,6 +2,8 @@ from collections import defaultdict, deque from typing import Dict, Generator, List, Type, Union from langchain.chains.base import Chain +from loguru import logger + from langflow.graph.edge.base import ContractEdge, Edge from langflow.graph.graph.constants import lazy_load_vertex_dict from langflow.graph.graph.utils import process_flow @@ -9,7 +11,6 @@ from langflow.graph.vertex.base import Vertex 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 class Graph: @@ -192,13 +193,13 @@ class Graph: neighbors[neighbor] += 1 return neighbors - def _build_edges(self) -> List[Edge]: + def _build_edges(self) -> List[ContractEdge]: """Builds the edges of the graph.""" # Edge takes two vertices as arguments, so we need to build the vertices first # and then build the edges # if we can't find a vertex, we raise an error - edges: List[Edge] = [] + edges: List[ContractEdge] = [] for edge in self._edges: source = self.get_vertex(edge["source"]) target = self.get_vertex(edge["target"]) @@ -206,7 +207,7 @@ class Graph: raise ValueError(f"Source vertex {edge['source']} not found") if target is None: raise ValueError(f"Target vertex {edge['target']} not found") - edges.append(Edge(source, target, edge)) + edges.append(ContractEdge(source, target, edge)) return edges def _get_vertex_class(self, vertex_type: str, vertex_base_type: str) -> Type[Vertex]: