From 0fa85687f102bf8cfe7d369fe3d33eb1d099d6a8 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Thu, 25 Jan 2024 11:43:27 -0300 Subject: [PATCH] Refactor layered_topological_sort method in Graph class --- src/backend/langflow/graph/graph/base.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/backend/langflow/graph/graph/base.py b/src/backend/langflow/graph/graph/base.py index ef747ec33..1037150a9 100644 --- a/src/backend/langflow/graph/graph/base.py +++ b/src/backend/langflow/graph/graph/base.py @@ -1,5 +1,5 @@ -from typing import Dict, Generator, List, Type, Union from collections import defaultdict, deque +from typing import Dict, Generator, List, Type, Union from langchain.chains.base import Chain from loguru import logger @@ -256,7 +256,7 @@ class Graph: return f"Graph:\nNodes: {vertex_ids}\nConnections:\n{edges_repr}" def layered_topological_sort(self): - in_degree = {vertex: 0 for vertex in self.vertices} # Initialize in-degrees + in_degree = {vertex.id: 0 for vertex in self.vertices} # Initialize in-degrees graph = defaultdict(list) # Adjacency list representation # Build graph and compute in-degrees @@ -265,7 +265,7 @@ class Graph: in_degree[edge.target_id] += 1 # Queue for vertices with no incoming edges - queue = deque(vertex for vertex in self.vertices if in_degree[vertex] == 0) + queue = deque(vertex.id for vertex in self.vertices if in_degree[vertex.id] == 0) layers = [] current_layer = 0 @@ -273,9 +273,9 @@ class Graph: layers.append([]) # Start a new layer layer_size = len(queue) for _ in range(layer_size): - vertex = queue.popleft() - layers[current_layer].append(vertex.id) - for neighbor in graph[vertex]: + vertex_id = queue.popleft() + layers[current_layer].append(vertex_id) + for neighbor in graph[vertex_id]: in_degree[neighbor] -= 1 # 'remove' edge if in_degree[neighbor] == 0: queue.append(neighbor)