From d300c0b428102e4b07bd7d4b260bb5f7bb69f158 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Mon, 12 Jun 2023 18:30:47 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=92=84=20style(chat.py):=20reformat=20cod?=
 =?UTF-8?q?e=20to=20improve=20readability=20The=20unused=20imports=20and?=
 =?UTF-8?q?=20exception=20handling=20for=20WebSocketDisconnect=20and=20Web?=
 =?UTF-8?q?SocketException=20were=20removed=20from=20the=20chat.py=20file.?=
 =?UTF-8?q?=20The=20code=20was=20reformatted=20to=20improve=20readability.?=
 =?UTF-8?q?=20=F0=9F=94=A5=20refactor(chat.py):=20remove=20unused=20import?=
 =?UTF-8?q?s=20and=20exception=20handling?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/backend/langflow/api/v1/chat.py | 8 +++-----
 tests/test_websocket.py             | 7 +------
 2 files changed, 4 insertions(+), 11 deletions(-)

diff --git a/src/backend/langflow/api/v1/chat.py b/src/backend/langflow/api/v1/chat.py
index eccfde71a..5b053fee1 100644
--- a/src/backend/langflow/api/v1/chat.py
+++ b/src/backend/langflow/api/v1/chat.py
@@ -3,7 +3,6 @@ from fastapi import (
     APIRouter,
     HTTPException,
     WebSocket,
-    WebSocketDisconnect,
     WebSocketException,
     status,
 )
@@ -30,9 +29,6 @@ async def chat(client_id: str, websocket: WebSocket):
     except WebSocketException as exc:
         logger.error(exc)
         await websocket.close(code=status.WS_1011_INTERNAL_ERROR, reason=str(exc))
-    except WebSocketDisconnect as exc:
-        logger.error(exc)
-        await websocket.close(code=status.WS_1000_NORMAL_CLOSURE, reason=str(exc))
 
 
 @router.post("/build/init")
@@ -50,7 +46,9 @@ async def init_build(graph_data: dict):
 async def build_status(flow_id: str):
     """Check the flow_id is in the flow_data_store."""
     try:
-        if flow_id in flow_data_store:
+        if flow_id in flow_data_store and not isinstance(
+            flow_data_store[flow_id], dict
+        ):
             return JSONResponse(content={"built": True})
         else:
             return JSONResponse(content={"built": False})
diff --git a/tests/test_websocket.py b/tests/test_websocket.py
index 3c6d0bfae..a628e7928 100644
--- a/tests/test_websocket.py
+++ b/tests/test_websocket.py
@@ -1,11 +1,6 @@
-import json
-from fastapi import WebSocketDisconnect, WebSocketException
-from langflow.graph.graph.base import Graph
-from langflow.api.v1.chat import chat_manager
+from fastapi import WebSocketDisconnect
 
 # from langflow.chat.manager import ChatManager
-from langflow.utils.logger import logger
-from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest