diff --git a/src/dsa/nodes.py b/src/dsa/nodes.py
index 521d24e..77d01e2 100644
--- a/src/dsa/nodes.py
+++ b/src/dsa/nodes.py
@@ -5,7 +5,10 @@
 # derived from https://github.com/langflow-ai/langflow/pull/5261
 def find_last_node(nodes, edges):
     """This function receives a flow and returns the last node."""
-    return next((n for n in nodes if all(e["source"] != n["id"] for e in edges)), None)
+    # Create a set of all node ids that are used as source in edges
+    source_ids = {e["source"] for e in edges}
+    # Return the first node whose id is not in the set of source_ids
+    return next((n for n in nodes if n["id"] not in source_ids), None)
 
 
 # Function to find all leaf nodes (nodes with no outgoing edges)