From 29bcf3e43fb6ee742254bc095f5def6ecc2b9e0e Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Wed, 31 Jan 2024 17:22:04 -0300
Subject: [PATCH] Add duration tracking to build_vertex function

---
 src/backend/langflow/api/v1/chat.py    | 5 ++++-
 src/backend/langflow/api/v1/schemas.py | 4 +++-
 2 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/backend/langflow/api/v1/chat.py b/src/backend/langflow/api/v1/chat.py
index 97162991e..fbbbbefb1 100644
--- a/src/backend/langflow/api/v1/chat.py
+++ b/src/backend/langflow/api/v1/chat.py
@@ -287,6 +287,8 @@ async def build_vertex(
         cache = chat_service.get_cache(flow_id)
         graph = cache.get("result")
         result_dict = {}
+        duration = ""
+        start_time = time.perf_counter()
         if tweaks:
             graph = process_tweaks_on_graph(graph, tweaks)
         if not isinstance(graph, Graph):
@@ -303,7 +305,8 @@ async def build_vertex(
             # to the frontend
             vertex.set_artifacts()
             artifacts = vertex.artifacts
-            result_dict = ResultDict(results=result_dict, artifacts=artifacts)
+            duration = format_elapsed_time(time.perf_counter() - start_time)
+            result_dict = ResultDict(results=result_dict, artifacts=artifacts, duration=duration)
         except Exception as exc:
             params = str(exc)
             valid = False
diff --git a/src/backend/langflow/api/v1/schemas.py b/src/backend/langflow/api/v1/schemas.py
index 2b89da976..dee62547e 100644
--- a/src/backend/langflow/api/v1/schemas.py
+++ b/src/backend/langflow/api/v1/schemas.py
@@ -221,9 +221,11 @@ class VerticesOrderResponse(BaseModel):
 
 
 class ResultDict(BaseModel):
+    """Outputs of the vertex build process."""
+
     results: Optional[Any] = Field(default_factory=dict)
     artifacts: Optional[Any] = Field(default_factory=dict)
-    """Outputs of the vertex build process."""
+    duration: Optional[float] = None
 
 
 class VertexBuildResponse(BaseModel):