Refactor graph class to add and update vertices and edges

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-03 20:23:34 -03:00
commit 13a073615b

View file

@ -12,8 +12,13 @@ from langflow.graph.graph.state_manager import GraphStateManager
from langflow.graph.graph.utils import process_flow from langflow.graph.graph.utils import process_flow
from langflow.graph.schema import INPUT_FIELD_NAME, InterfaceComponentTypes from langflow.graph.schema import INPUT_FIELD_NAME, InterfaceComponentTypes
from langflow.graph.vertex.base import Vertex from langflow.graph.vertex.base import Vertex
from langflow.graph.vertex.types import (ChatVertex, FileToolVertex, LLMVertex, from langflow.graph.vertex.types import (
RoutingVertex, ToolkitVertex) ChatVertex,
FileToolVertex,
LLMVertex,
RoutingVertex,
ToolkitVertex,
)
from langflow.interface.tools.constants import FILE_TOOLS from langflow.interface.tools.constants import FILE_TOOLS
from langflow.utils import payload from langflow.utils import payload
@ -285,9 +290,14 @@ class Graph:
# Add new vertices # Add new vertices
for vertex_id in new_vertex_ids: for vertex_id in new_vertex_ids:
new_vertex = other.get_vertex(vertex_id) new_vertex = other.get_vertex(vertex_id)
new_vertex.graph = self
self._add_vertex(new_vertex) self._add_vertex(new_vertex)
# Now update the edges
for vertex_id in new_vertex_ids:
new_vertex = other.get_vertex(vertex_id)
self._update_edges(new_vertex)
new_vertex.graph = self
# Update existing vertices that have changed # Update existing vertices that have changed
for vertex_id in existing_vertex_ids.intersection(other_vertex_ids): for vertex_id in existing_vertex_ids.intersection(other_vertex_ids):
self_vertex = self.get_vertex(vertex_id) self_vertex = self.get_vertex(vertex_id)
@ -334,12 +344,24 @@ class Graph:
_vertex._build_params() _vertex._build_params()
def _add_vertex(self, vertex: Vertex) -> None: def _add_vertex(self, vertex: Vertex) -> None:
"""Adds a new vertex to the graph.""" """Adds a vertex to the graph."""
self.vertices.append(vertex) self.vertices.append(vertex)
self.vertex_map[vertex.id] = vertex self.vertex_map[vertex.id] = vertex
def add_vertex(self, vertex: Vertex) -> None:
"""Adds a new vertex to the graph."""
self._add_vertex(vertex)
self._update_edges(vertex)
def _update_edges(self, vertex: Vertex) -> None:
"""Updates the edges of a vertex."""
# Vertex has edges, so we need to update the edges # Vertex has edges, so we need to update the edges
for edge in vertex.edges: for edge in vertex.edges:
if edge.source_id in self.vertex_map and edge.target_id in self.vertex_map: if (
edge not in self.edges
and edge.source_id in self.vertex_map
and edge.target_id in self.vertex_map
):
self.edges.append(edge) self.edges.append(edge)
def _build_graph(self) -> None: def _build_graph(self) -> None: