diff --git a/src/dsa/nodes.py b/src/dsa/nodes.py
index 521d24e..33059c4 100644
--- a/src/dsa/nodes.py
+++ b/src/dsa/nodes.py
@@ -61,29 +61,34 @@ def find_cycle_vertices(edges):
 
 # derived from https://github.com/langflow-ai/langflow/pull/5263
 def sort_chat_inputs_first(self, vertices_layers: list[list[str]]) -> list[list[str]]:
-    # First check if any chat inputs have dependencies
-    for layer in vertices_layers:
+    # First, prepare to check only ChatInputs for dependencies
+    chatinputs_indices = []  # (layer_idx, position) if needed for other uses
+    chatinputs_ids = []
+    layers_len = len(vertices_layers)
+
+    # Gather all ChatInputs along with their indices, and check dependencies immediately
+    for i in range(layers_len):
+        layer = vertices_layers[i]
         for vertex_id in layer:
-            if "ChatInput" in vertex_id and self.get_predecessors(
-                self.get_vertex(vertex_id)
-            ):
-                return vertices_layers
-
-    # If no chat inputs have dependencies, move them to first layer
-    chat_inputs_first = []
-    for layer in vertices_layers:
-        layer_chat_inputs_first = [
-            vertex_id for vertex_id in layer if "ChatInput" in vertex_id
-        ]
-        chat_inputs_first.extend(layer_chat_inputs_first)
-        for vertex_id in layer_chat_inputs_first:
-            # Remove the ChatInput from the layer
-            layer.remove(vertex_id)
-
-    if not chat_inputs_first:
+            if "ChatInput" in vertex_id:
+                chatinputs_ids.append(vertex_id)
+                # Check dependencies lazily (only when candidate is found)
+                vertex = self.get_vertex(vertex_id)
+                predecessors = self.get_predecessors(vertex)
+                if predecessors:
+                    # If any ChatInput has dependencies, return immediately
+                    return vertices_layers
+
+    if not chatinputs_ids:
         return vertices_layers
 
-    return [chat_inputs_first, *vertices_layers]
+    # Now, rebuild each layer omitting the ChatInputs and prepend them as a new layer
+    result_layers = []
+    for layer in vertices_layers:
+        new_layer = [vid for vid in layer if "ChatInput" not in vid]
+        if new_layer:
+            result_layers.append(new_layer)
+    return [chatinputs_ids, *result_layers]
 
 
 # Function to find the node with highest degree (most connections)