From d6dcfd526673feb3a5483009f83479d1f755dbda Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Wed, 27 Sep 2023 11:46:22 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(chat.py):=20replace=20direct?=
 =?UTF-8?q?=20assignment=20with=20update=5Fbuild=5Fstatus=20function=20to?=
 =?UTF-8?q?=20update=20build=20status=20in=20cache=5Fservice=20dictionary?=
 =?UTF-8?q?=20=F0=9F=94=A7=20chore(utils.py):=20add=20update=5Fbuild=5Fsta?=
 =?UTF-8?q?tus=20function=20to=20update=20build=20status=20in=20cache=5Fse?=
 =?UTF-8?q?rvice=20dictionary?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/backend/langflow/api/v1/chat.py          | 10 ++++++----
 src/backend/langflow/services/cache/utils.py |  7 +++++++
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/src/backend/langflow/api/v1/chat.py b/src/backend/langflow/api/v1/chat.py
index a7a50a39f..c25c4bde2 100644
--- a/src/backend/langflow/api/v1/chat.py
+++ b/src/backend/langflow/api/v1/chat.py
@@ -13,6 +13,7 @@ from langflow.api.v1.schemas import BuildStatus, BuiltResponse, InitResponse, St
 
 from langflow.graph.graph.base import Graph
 from langflow.services.auth.utils import get_current_active_user, get_current_user
+from langflow.services.cache.utils import update_build_status
 from loguru import logger
 from langflow.services.getters import get_chat_service, get_session, get_cache_service
 from sqlmodel import Session
@@ -157,7 +158,7 @@ async def stream_build(
             graph = Graph.from_payload(graph_data)
 
             number_of_nodes = len(graph.nodes)
-            cache_service[flow_id]["status"] = BuildStatus.IN_PROGRESS
+            update_build_status(cache_service, flow_id, BuildStatus.IN_PROGRESS)
 
             for i, vertex in enumerate(graph.generator_build(), 1):
                 try:
@@ -184,7 +185,7 @@ async def stream_build(
                     logger.exception(exc)
                     params = str(exc)
                     valid = False
-                    cache_service[flow_id]["status"] = BuildStatus.FAILURE
+                    update_build_status(cache_service, flow_id, BuildStatus.FAILURE)
 
                 response = {
                     "valid": valid,
@@ -211,11 +212,12 @@ async def stream_build(
             chat_service.set_cache(flow_id, langchain_object)
             # We need to reset the chat history
             chat_service.chat_history.empty_history(flow_id)
-            cache_service[flow_id]["status"] = BuildStatus.SUCCESS
+            update_build_status(cache_service, flow_id, BuildStatus.SUCCESS)
         except Exception as exc:
             logger.exception(exc)
             logger.error("Error while building the flow: %s", exc)
-            cache_service[flow_id]["status"] = BuildStatus.FAILURE
+
+            update_build_status(cache_service, flow_id, BuildStatus.FAILURE)
             yield str(StreamData(event="error", data={"error": str(exc)}))
         finally:
             yield str(StreamData(event="message", data=final_response))
diff --git a/src/backend/langflow/services/cache/utils.py b/src/backend/langflow/services/cache/utils.py
index 83cde2240..a708c6e7e 100644
--- a/src/backend/langflow/services/cache/utils.py
+++ b/src/backend/langflow/services/cache/utils.py
@@ -9,6 +9,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict
 from appdirs import user_cache_dir
 from fastapi import UploadFile
+from langflow.api.v1.schemas import BuildStatus
 from langflow.services.database.models.base import orjson_dumps
 
 if TYPE_CHECKING:
@@ -202,3 +203,9 @@ def save_uploaded_file(file: UploadFile, folder_name):
             new_file.write(chunk)
 
     return file_path
+
+
+def update_build_status(cache_service, flow_id: str, status: BuildStatus):
+    cached_flow = cache_service[flow_id]
+    cached_flow["status"] = status
+    cache_service[flow_id] = cached_flow