feat: refactor graph vertex sorting (#2583)
* refactor: extract method from class to new func * test: add new tests * refactor: simplify funcs to improve readability * refactor: extract new func from larger func * refactor: remove recursion from func * refactor: remove coupling with graph and vertex * refactor: create adapter funcs to use new code * refactor: add test for sorting up to vertex N with is_start=True --------- Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
This commit is contained in:
Ítalo Johnny
GitHub
parent
3406575c67
commit
aa1958a4ad
3 changed files with 190 additions and 71 deletions
|
|
@ -14,7 +14,7 @@ from langflow.graph.edge.base import ContractEdge
|
|||
from langflow.graph.graph.constants import lazy_load_vertex_dict
|
||||
from langflow.graph.graph.runnable_vertices_manager import RunnableVerticesManager
|
||||
from langflow.graph.graph.state_manager import GraphStateManager
|
||||
from langflow.graph.graph.utils import find_start_component_id, process_flow
|
||||
from langflow.graph.graph.utils import find_start_component_id, process_flow, sort_up_to_vertex
|
||||
from langflow.graph.schema import InterfaceComponentTypes, RunOutputs
|
||||
from langflow.graph.vertex.base import Vertex, VertexStates
|
||||
from langflow.graph.vertex.types import InterfaceVertex, StateVertex
|
||||
|
|
@ -1197,74 +1197,6 @@ class Graph:
|
|||
edges_repr = "\n".join([f"{edge.source_id} --> {edge.target_id}" for edge in self.edges])
|
||||
return f"Graph:\nNodes: {vertex_ids}\nConnections:\n{edges_repr}"
|
||||
|
||||
def sort_up_to_vertex(self, vertex_id: str, is_start: bool = False) -> List[Vertex]:
|
||||
"""Cuts the graph up to a given vertex and sorts the resulting subgraph."""
|
||||
# Initial setup
|
||||
visited = set() # To keep track of visited vertices
|
||||
excluded = set() # To keep track of vertices that should be excluded
|
||||
|
||||
def get_successors(vertex, recursive=True):
|
||||
# Recursively get the successors of the current vertex
|
||||
successors = vertex.successors
|
||||
if not successors:
|
||||
return []
|
||||
successors_result = []
|
||||
for successor in successors:
|
||||
# Just return a list of successors
|
||||
if recursive:
|
||||
next_successors = get_successors(successor)
|
||||
successors_result.extend(next_successors)
|
||||
successors_result.append(successor)
|
||||
return successors_result
|
||||
|
||||
try:
|
||||
stop_or_start_vertex = self.get_vertex(vertex_id)
|
||||
stack = [vertex_id] # Use a list as a stack for DFS
|
||||
except ValueError:
|
||||
stop_or_start_vertex = self.get_root_of_group_node(vertex_id)
|
||||
stack = [stop_or_start_vertex.id]
|
||||
vertex_id = stop_or_start_vertex.id
|
||||
stop_predecessors = [pre.id for pre in stop_or_start_vertex.predecessors]
|
||||
# DFS to collect all vertices that can reach the specified vertex
|
||||
while stack:
|
||||
current_id = stack.pop()
|
||||
if current_id not in visited and current_id not in excluded:
|
||||
visited.add(current_id)
|
||||
current_vertex = self.get_vertex(current_id)
|
||||
# Assuming get_predecessors is a method that returns all vertices with edges to current_vertex
|
||||
for predecessor in current_vertex.predecessors:
|
||||
stack.append(predecessor.id)
|
||||
|
||||
if current_id == vertex_id:
|
||||
# We should add to visited all the vertices that are successors of the current vertex
|
||||
# and their successors and so on
|
||||
# if the vertex is a start, it means we are starting from the beginning
|
||||
# and getting successors
|
||||
for successor in current_vertex.successors:
|
||||
if is_start:
|
||||
stack.append(successor.id)
|
||||
else:
|
||||
excluded.add(successor.id)
|
||||
all_successors = get_successors(successor, recursive=False)
|
||||
for successor in all_successors:
|
||||
if is_start:
|
||||
stack.append(successor.id)
|
||||
else:
|
||||
excluded.add(successor.id)
|
||||
elif current_id not in stop_predecessors and is_start:
|
||||
# If the current vertex is not the target vertex, we should add all its successors
|
||||
# to the stack if they are not in visited
|
||||
|
||||
# If we are starting from the beginning, we should add all successors
|
||||
for successor in current_vertex.successors:
|
||||
if successor.id not in visited:
|
||||
stack.append(successor.id)
|
||||
|
||||
# Filter the original graph's vertices and edges to keep only those in `visited`
|
||||
vertices_to_keep = [self.get_vertex(vid) for vid in visited]
|
||||
|
||||
return vertices_to_keep
|
||||
|
||||
def layered_topological_sort(
|
||||
self,
|
||||
vertices: List[Vertex],
|
||||
|
|
@ -1395,6 +1327,21 @@ class Graph:
|
|||
max_index = max(max_index, index_map[successor.id])
|
||||
return max_index
|
||||
|
||||
def __to_dict(self) -> Dict[str, Dict[str, List[str]]]:
|
||||
"""Converts the graph to a dictionary."""
|
||||
result: Dict = dict()
|
||||
for vertex in self.vertices:
|
||||
vertex_id = vertex.id
|
||||
sucessors = [i.id for i in self.get_all_successors(vertex)]
|
||||
predecessors = [i.id for i in self.get_predecessors(vertex)]
|
||||
result |= {vertex_id: {"successors": sucessors, "predecessors": predecessors}}
|
||||
return result
|
||||
|
||||
def __filter_vertices(self, vertex_id: str, is_start: bool = False):
|
||||
dictionaryized_graph = self.__to_dict()
|
||||
vertex_ids = sort_up_to_vertex(dictionaryized_graph, vertex_id, is_start)
|
||||
return [self.get_vertex(vertex_id) for vertex_id in vertex_ids]
|
||||
|
||||
def sort_vertices(
|
||||
self,
|
||||
stop_component_id: Optional[str] = None,
|
||||
|
|
@ -1404,9 +1351,11 @@ class Graph:
|
|||
self.mark_all_vertices("ACTIVE")
|
||||
if stop_component_id is not None:
|
||||
self.stop_vertex = stop_component_id
|
||||
vertices = self.sort_up_to_vertex(stop_component_id)
|
||||
vertices = self.__filter_vertices(stop_component_id)
|
||||
|
||||
elif start_component_id:
|
||||
vertices = self.sort_up_to_vertex(start_component_id, is_start=True)
|
||||
vertices = self.__filter_vertices(start_component_id, is_start=True)
|
||||
|
||||
else:
|
||||
vertices = self.vertices
|
||||
# without component_id we are probably running in the chat
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
from typing import List, Dict
|
||||
import copy
|
||||
from collections import deque
|
||||
|
||||
|
||||
PRIORITY_LIST_OF_INPUTS = ["webhook", "chat"]
|
||||
|
||||
|
||||
|
|
@ -224,3 +226,49 @@ def get_updated_edges(base_flow, g_nodes, g_edges, group_node_id):
|
|||
if edge["target"] == group_node_id or edge["source"] == group_node_id:
|
||||
updated_edges.append(new_edge)
|
||||
return updated_edges
|
||||
|
||||
|
||||
def get_successors(graph: Dict[str, Dict[str, List[str]]], vertex_id: str) -> List[str]:
|
||||
successors_result = []
|
||||
stack = [vertex_id]
|
||||
while stack:
|
||||
current_id = stack.pop()
|
||||
successors_result.append(current_id)
|
||||
stack.extend(graph[current_id]["successors"])
|
||||
return successors_result
|
||||
|
||||
|
||||
def sort_up_to_vertex(graph: Dict[str, Dict[str, List[str]]], vertex_id: str, is_start: bool = False) -> List[str]:
|
||||
"""Cuts the graph up to a given vertex and sorts the resulting subgraph."""
|
||||
try:
|
||||
stop_or_start_vertex = graph[vertex_id]
|
||||
except KeyError:
|
||||
raise ValueError(f"Vertex {vertex_id} not found into graph")
|
||||
|
||||
visited, excluded = set(), set()
|
||||
stack = [vertex_id]
|
||||
stop_predecessors = set(stop_or_start_vertex["predecessors"])
|
||||
|
||||
while stack:
|
||||
current_id = stack.pop()
|
||||
if current_id in visited or current_id in excluded:
|
||||
continue
|
||||
|
||||
visited.add(current_id)
|
||||
current_vertex = graph[current_id]
|
||||
|
||||
stack.extend(current_vertex["predecessors"])
|
||||
|
||||
if current_id == vertex_id or (current_id not in stop_predecessors and is_start):
|
||||
for successor_id in current_vertex["successors"]:
|
||||
if is_start:
|
||||
stack.append(successor_id)
|
||||
else:
|
||||
excluded.add(successor_id)
|
||||
for succ_id in get_successors(graph, successor_id):
|
||||
if is_start:
|
||||
stack.append(succ_id)
|
||||
else:
|
||||
excluded.add(succ_id)
|
||||
|
||||
return list(visited)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue