diff --git a/src/backend/langflow/utils/payload.py b/src/backend/langflow/utils/payload.py
index 02cca5c71..0e2f0fc7a 100644
--- a/src/backend/langflow/utils/payload.py
+++ b/src/backend/langflow/utils/payload.py
@@ -28,16 +28,16 @@ def extract_input_variables(nodes):
     return nodes
 
 
-def get_root_node(graph):
+def get_root_vertex(graph):
     """
     Returns the root node of the template.
     """
-    incoming_edges = {edge.source for edge in graph.edges}
+    incoming_edges = {edge.source_id for edge in graph.edges}
 
-    if not incoming_edges and len(graph.nodes) == 1:
-        return graph.nodes[0]
+    if not incoming_edges and len(graph.vertices) == 1:
+        return graph.vertices[0]
 
-    return next((node for node in graph.nodes if node not in incoming_edges), None)
+    return next((node for node in graph.vertices if node.id not in incoming_edges), None)
 
 
 def build_json(root, graph) -> Dict: