From 76c91de563f64f5e8fcba1b1aa6da42fc77fe0bb Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
Date: Mon, 11 Nov 2024 11:22:04 -0300
Subject: [PATCH] fix: Enable sorting of chat inputs to improve UI message
 speed (#4480)

---
 src/backend/base/langflow/graph/graph/base.py | 3 ++-
 src/backend/tests/unit/test_chat_endpoint.py  | 2 +-
 src/backend/tests/unit/test_endpoints.py      | 2 +-
 3 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/src/backend/base/langflow/graph/graph/base.py b/src/backend/base/langflow/graph/graph/base.py
index 6384ce400..f064e65be 100644
--- a/src/backend/base/langflow/graph/graph/base.py
+++ b/src/backend/base/langflow/graph/graph/base.py
@@ -2006,7 +2006,8 @@ class Graph:
 
         vertices_layers = self.layered_topological_sort(vertices)
         vertices_layers = self.sort_by_avg_build_time(vertices_layers)
-        # vertices_layers = self.sort_chat_inputs_first(vertices_layers)
+        # Sort the chat inputs first to speed up sending the User message to the UI
+        vertices_layers = self.sort_chat_inputs_first(vertices_layers)
         # Now we should sort each layer in a way that we make sure
         # vertex V does not depend on vertex V+1
         vertices_layers = self.sort_layer_by_dependency(vertices_layers)
diff --git a/src/backend/tests/unit/test_chat_endpoint.py b/src/backend/tests/unit/test_chat_endpoint.py
index 9fe66e379..c892da463 100644
--- a/src/backend/tests/unit/test_chat_endpoint.py
+++ b/src/backend/tests/unit/test_chat_endpoint.py
@@ -73,7 +73,7 @@ async def consume_and_assert_stream(r):
             assert parsed["event"] == "vertices_sorted"
             ids = parsed["data"]["ids"]
             ids.sort()
-            assert ids == ["ChatInput-CIGht", "Memory-amN4Z"]
+            assert ids == ["ChatInput-CIGht"]
 
             to_run = parsed["data"]["to_run"]
             to_run.sort()
diff --git a/src/backend/tests/unit/test_endpoints.py b/src/backend/tests/unit/test_endpoints.py
index 22efe00bb..eb1863616 100644
--- a/src/backend/tests/unit/test_endpoints.py
+++ b/src/backend/tests/unit/test_endpoints.py
@@ -257,7 +257,7 @@ async def test_get_vertices(client, added_flow_webhook_test, logged_in_headers):
     # The important part is before the - (ConversationBufferMemory, PromptTemplate, ChatOpenAI, LLMChain)
     ids = [_id.split("-")[0] for _id in response.json()["ids"]]
 
-    assert set(ids) == {"Webhook", "ChatInput"}
+    assert set(ids) == {"ChatInput"}
 
 
 async def test_build_vertex_invalid_flow_id(client, logged_in_headers):