From 7abf697b7b41dbcd930af7b631f70bfe65f91c27 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sat, 6 Jul 2024 23:54:25 -0300 Subject: [PATCH] fix: update run_manager on state changes (#2560) * feat: update RunnableVerticesManager to include an update_run_state method * feat: refactor build_run_map method in RunnableVerticesManager This commit refactors the `build_run_map` method in the `RunnableVerticesManager` class. Instead of directly accessing the `graph` object, it now takes `predecessor_map` and `vertices_to_run` as parameters. This improves code readability and maintainability. * fix(state): set vertices_to_run and predecessor_map correctly when updating a state * fix: set inactive vertices to active when updating graph * fix: add conditions to avoid duplicated vertices in maps This commit fixes a bug in the Graph class where the predecessor and successor maps were not being updated correctly. The bug caused incorrect mapping of edges in the graph. This fix ensures that the maps are updated properly when adding edges to the graph. * refactor: fix vertex build response in chat.py This commit fixes the `build_vertex` function in `chat.py` to correctly handle the `VertexBuildResponse` object. The `inactivated_vertices`, `next_vertices_ids`, and `top_level_vertices` attributes are now converted to lists to remove any duplicate values. This ensures that the response object contains unique values for these attributes. The fix improves the accuracy and reliability of the vertex build process. * style(graph/base.py): update type annotations for predecessor_map and successor_map variables to improve code readability and maintain consistency * fix: uncomment code for stopping build in MenuBar component --- src/backend/base/langflow/api/v1/chat.py | 6 ++-- src/backend/base/langflow/graph/graph/base.py | 34 ++++++++++++------- .../graph/graph/runnable_vertices_manager.py | 13 ++++--- .../components/menuBar/index.tsx | 5 ++- 4 files changed, 36 insertions(+), 22 deletions(-) diff --git a/src/backend/base/langflow/api/v1/chat.py b/src/backend/base/langflow/api/v1/chat.py index 716afca72..bbc183da8 100644 --- a/src/backend/base/langflow/api/v1/chat.py +++ b/src/backend/base/langflow/api/v1/chat.py @@ -261,9 +261,9 @@ async def build_vertex( background_tasks.add_task(graph.end_all_traces) build_response = VertexBuildResponse( - inactivated_vertices=inactivated_vertices, - next_vertices_ids=next_runnable_vertices, - top_level_vertices=top_level_vertices, + inactivated_vertices=list(set(inactivated_vertices)), + next_vertices_ids=list(set(next_runnable_vertices)), + top_level_vertices=list(set(top_level_vertices)), valid=valid, params=params, id=vertex.id, diff --git a/src/backend/base/langflow/graph/graph/base.py b/src/backend/base/langflow/graph/graph/base.py index 253dc8484..1b383f7e9 100644 --- a/src/backend/base/langflow/graph/graph/base.py +++ b/src/backend/base/langflow/graph/graph/base.py @@ -16,7 +16,7 @@ from langflow.graph.graph.runnable_vertices_manager import RunnableVerticesManag from langflow.graph.graph.state_manager import GraphStateManager from langflow.graph.graph.utils import find_start_component_id, process_flow from langflow.graph.schema import InterfaceComponentTypes, RunOutputs -from langflow.graph.vertex.base import Vertex +from langflow.graph.vertex.base import Vertex, VertexStates from langflow.graph.vertex.types import InterfaceVertex, StateVertex from langflow.schema import Data from langflow.schema.schema import INPUT_FIELD_NAME, InputType @@ -131,7 +131,8 @@ class Graph: name (str): The name of the state. caller (str): The ID of the vertex that is updating the state. """ - vertices_ids = [] + vertices_ids = set() + new_predecessor_map = {} for vertex_id in self._is_state_vertices: if vertex_id == caller: continue @@ -142,7 +143,7 @@ class Graph: and vertex_id != caller and isinstance(vertex, StateVertex) ): - vertices_ids.append(vertex_id) + vertices_ids.add(vertex_id) successors = self.get_all_successors(vertex, flat=True) # Update run_manager.run_predecessors because we are activating vertices # The run_prdecessors is the predecessor map of the vertices @@ -153,12 +154,19 @@ class Graph: edges_set = set() for vertex in [vertex] + successors: edges_set.update(vertex.edges) + vertices_ids.add(vertex.id) + if vertex.state == VertexStates.INACTIVE: + vertex.set_state("ACTIVE") edges = list(edges_set) - new_predecessor_map, _ = self.build_adjacency_maps(edges) - self.run_manager.run_predecessors.update(new_predecessor_map) - self.vertices_to_run.update(list(map(lambda x: x.id, successors))) - self.activated_vertices = vertices_ids + predecessor_map, _ = self.build_adjacency_maps(edges) + new_predecessor_map.update(predecessor_map) + + self.activated_vertices = list(vertices_ids) self.vertices_to_run.update(vertices_ids) + self.run_manager.update_run_state( + run_predecessors=new_predecessor_map, + vertices_to_run=self.vertices_to_run, + ) def reset_activated_vertices(self): """ @@ -1459,7 +1467,7 @@ class Graph: Returns: None """ - self.run_manager.build_run_map(self) + self.run_manager.build_run_map(predecessor_map=self.predecessor_map, vertices_to_run=self.vertices_to_run) def find_runnable_predecessors_for_successors(self, vertex_id: str) -> List[str]: """ @@ -1483,9 +1491,11 @@ class Graph: def build_adjacency_maps(self, edges: List[ContractEdge]) -> Tuple[Dict[str, List[str]], Dict[str, List[str]]]: """Returns the adjacency maps for the graph.""" - predecessor_map = defaultdict(list) - successor_map = defaultdict(list) + predecessor_map: dict[str, list[str]] = defaultdict(list) + successor_map: dict[str, list[str]] = defaultdict(list) for edge in edges: - predecessor_map[edge.target_id].append(edge.source_id) - successor_map[edge.source_id].append(edge.target_id) + if edge.source_id not in predecessor_map[edge.target_id]: + predecessor_map[edge.target_id].append(edge.source_id) + if edge.target_id not in successor_map[edge.source_id]: + successor_map[edge.source_id].append(edge.target_id) return predecessor_map, successor_map 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 95805ced0..fc9970554 100644 --- a/src/backend/base/langflow/graph/graph/runnable_vertices_manager.py +++ b/src/backend/base/langflow/graph/graph/runnable_vertices_manager.py @@ -40,6 +40,11 @@ class RunnableVerticesManager: self.run_predecessors = state["run_predecessors"] self.vertices_to_run = state["vertices_to_run"] + 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) + self.build_run_map(self.run_predecessors, self.vertices_to_run) + def is_vertex_runnable(self, vertex_id: str, inactivated_vertices: set[str]) -> bool: """Determines if a vertex is runnable.""" @@ -68,14 +73,14 @@ class RunnableVerticesManager: if vertex_id in self.run_predecessors[predecessor]: self.run_predecessors[predecessor].remove(vertex_id) - def build_run_map(self, graph): + def build_run_map(self, predecessor_map, vertices_to_run): """Builds a map of vertices and their runnable successors.""" self.run_map = defaultdict(list) - for vertex_id, predecessors in graph.predecessor_map.items(): + for vertex_id, predecessors in predecessor_map.items(): for predecessor in predecessors: self.run_map[predecessor].append(vertex_id) - self.run_predecessors = graph.predecessor_map.copy() - self.vertices_to_run = graph.vertices_to_run + self.run_predecessors = predecessor_map.copy() + self.vertices_to_run = vertices_to_run def update_vertex_run_state(self, vertex_id: str, is_runnable: bool): """Updates the runnable state of a vertex.""" diff --git a/src/frontend/src/components/headerComponent/components/menuBar/index.tsx b/src/frontend/src/components/headerComponent/components/menuBar/index.tsx index 7376bfc4f..ba3b0e9ad 100644 --- a/src/frontend/src/components/headerComponent/components/menuBar/index.tsx +++ b/src/frontend/src/components/headerComponent/components/menuBar/index.tsx @@ -218,8 +218,7 @@ export const MenuBar = ({}: {}): JSX.Element => { />
{printByBuildStatus()}
- {/* Deactivating this until we find a better solution */} - {/* */} + )}