jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

Refactor vertex lookup using a dictionary

Browse source
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-11-28 17:56:30 -03:00
commit 86ce36426a
2 changed files with 5 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

5
src/backend/langflow/graph/graph/base.py
View file

@ -70,7 +70,7 @@ class Graph:
def _build_graph(self) -> None:
"""Builds the graph from the vertices and edges."""
self.vertices = self._build_vertices()
self.vertex_ids = [vertex.id for vertex in self.vertices]
self.vertex_map = {vertex.id: vertex for vertex in self.vertices}
self.edges = self._build_edges()
# This is a hack to make sure that the LLM vertex is sent to
@ -107,7 +107,7 @@ class Graph:
def get_vertex(self, vertex_id: str) -> Union[None, Vertex]:
"""Returns a vertex by id."""
return next((vertex for vertex in self.vertices if vertex.id == vertex_id), None)
return self.vertex_map.get(vertex_id)
def get_vertex_edges(self, vertex_id: str) -> List[Edge]:
"""Returns a list of edges for a given vertex."""
@ -249,3 +249,4 @@ class Graph:
vertex_ids = [vertex.id for vertex in self.vertices]
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}"
return f"Graph:\nNodes: {vertex_ids}\nConnections:\n{edges_repr}"