From 537e358b65f807d80ad90e5e1fdfaf3fe331d078 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 15 Jul 2024 13:13:49 -0300 Subject: [PATCH] fix: update condition to run end_all_traces (#2707) --- src/backend/base/langflow/api/v1/chat.py | 2 +- src/backend/base/langflow/graph/graph/base.py | 2 ++ .../base/langflow/graph/graph/runnable_vertices_manager.py | 3 +++ 3 files changed, 6 insertions(+), 1 deletion(-) diff --git a/src/backend/base/langflow/api/v1/chat.py b/src/backend/base/langflow/api/v1/chat.py index 8c925e71f..ada8e416e 100644 --- a/src/backend/base/langflow/api/v1/chat.py +++ b/src/backend/base/langflow/api/v1/chat.py @@ -260,7 +260,7 @@ async def build_vertex( if graph.stop_vertex and graph.stop_vertex in next_runnable_vertices: next_runnable_vertices = [graph.stop_vertex] - if not graph.run_manager.vertices_to_run and not next_runnable_vertices: + if graph.run_manager.all_predecessors_are_fulfilled() and not next_runnable_vertices: background_tasks.add_task(graph.end_all_traces) build_response = VertexBuildResponse( diff --git a/src/backend/base/langflow/graph/graph/base.py b/src/backend/base/langflow/graph/graph/base.py index 4c4042578..87e602c56 100644 --- a/src/backend/base/langflow/graph/graph/base.py +++ b/src/backend/base/langflow/graph/graph/base.py @@ -539,6 +539,8 @@ class Graph: """Marks a vertex in the graph.""" vertex = self.get_vertex(vertex_id) vertex.set_state(state) + if state == VertexStates.INACTIVE: + self.run_manager.remove_from_predecessors(vertex_id) def mark_branch(self, vertex_id: str, state: str, visited: Optional[set] = None, output_name: Optional[str] = None): """Marks a branch of the graph.""" diff --git a/src/backend/base/langflow/graph/graph/runnable_vertices_manager.py b/src/backend/base/langflow/graph/graph/runnable_vertices_manager.py index b1dd8990f..4ae846230 100644 --- a/src/backend/base/langflow/graph/graph/runnable_vertices_manager.py +++ b/src/backend/base/langflow/graph/graph/runnable_vertices_manager.py @@ -41,6 +41,9 @@ class RunnableVerticesManager: self.run_predecessors = state["run_predecessors"] self.vertices_to_run = state["vertices_to_run"] + def all_predecessors_are_fulfilled(self) -> bool: + return all(not value for value in self.run_predecessors.values()) + def update_run_state(self, run_predecessors: dict, vertices_to_run: set): self.run_predecessors.update(run_predecessors) self.vertices_to_run.update(vertices_to_run)