From 60c7ea3e35a83fc31d2cfaff06891f3da69b2abd Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Tue, 20 Feb 2024 16:01:17 -0300 Subject: [PATCH] Sort vertices containing interface components first --- src/backend/langflow/graph/graph/base.py | 27 ++++++++++++++++++------ 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/src/backend/langflow/graph/graph/base.py b/src/backend/langflow/graph/graph/base.py index 3c772c897..370c45461 100644 --- a/src/backend/langflow/graph/graph/base.py +++ b/src/backend/langflow/graph/graph/base.py @@ -418,6 +418,7 @@ class Graph: if edge not in edges: edges.append(edge) vertices = self.layered_topological_sort(vertices, edges) + vertices = self.sort_interface_components_first(vertices) return self.sort_chat_inputs_first(vertices) def layered_topological_sort( @@ -523,11 +524,23 @@ class Graph: # and sort the layers accordingly # InterfaceComponentTypes is an enum # check all values of the enum and sort the layers - for layer in vertices: - layer.sort( - key=lambda x: any( - comp_type.value in x - for comp_type in InterfaceComponentTypes.__members__.values() - ) - ) + vertices = self.sort_interface_components_first(vertices) return self.sort_chat_inputs_first(vertices) + + def sort_interface_components_first(self, vertices: List[Vertex]) -> List[Vertex]: + """Sorts the vertices in the graph so that vertices containing ChatInput or ChatOutput come first.""" + + def contains_interface_component(vertex): + return any( + component.value in vertex for component in InterfaceComponentTypes + ) + + # Sort each inner list so that vertices containing ChatInput or ChatOutput come first + sorted_vertices = [ + sorted( + inner_list, + key=lambda vertex: not contains_interface_component(vertex), + ) + for inner_list in vertices + ] + return sorted_vertices