diff --git a/src/dsa/nodes.py b/src/dsa/nodes.py
index 521d24e..f1b0501 100644
--- a/src/dsa/nodes.py
+++ b/src/dsa/nodes.py
@@ -62,28 +62,30 @@ 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
+    chat_input_layers = []
     for layer in vertices_layers:
         for vertex_id in layer:
-            if "ChatInput" in vertex_id and self.get_predecessors(
-                self.get_vertex(vertex_id)
-            ):
-                return vertices_layers
+            if "ChatInput" in vertex_id:
+                # Direct method call reduced
+                if self.get_predecessors(self.get_vertex(vertex_id)):
+                    return vertices_layers
+                chat_input_layers.append(vertex_id)
+
+    if not chat_input_layers:
+        return vertices_layers
+
+    # Efficiently produce new layers without in-place mutation
+    new_layers = []
+    chat_input_set = set(chat_input_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
+        filtered_layer = [
+            vertex_id for vertex_id in layer if vertex_id not in chat_input_set
         ]
-        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:
-        return vertices_layers
+        if filtered_layer:  # Skip empty layers
+            new_layers.append(filtered_layer)
 
-    return [chat_inputs_first, *vertices_layers]
+    return [chat_input_layers, *new_layers]
 
 
 # Function to find the node with highest degree (most connections)