From 41f6aecda4e2358cd8e428fe9debff550c5eadc7 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Tue, 13 Jun 2023 17:21:45 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=80=20refactor(chat.py):=20import=20In?=
 =?UTF-8?q?itResponse=20and=20BuiltResponse=20from=20schemas=20module=20?=
 =?UTF-8?q?=F0=9F=9A=80=20feat(chat.py):=20add=20response=20models=20to=20?=
 =?UTF-8?q?/build/init=20and=20/build/{flow=5Fid}/status=20endpoints=20The?=
 =?UTF-8?q?=20InitResponse=20and=20BuiltResponse=20models=20are=20now=20im?=
 =?UTF-8?q?ported=20from=20the=20schemas=20module=20to=20improve=20code=20?=
 =?UTF-8?q?organization.=20The=20/build/init=20and=20/build/{flow=5Fid}/st?=
 =?UTF-8?q?atus=20endpoints=20now=20have=20response=20models=20to=20provid?=
 =?UTF-8?q?e=20a=20clear=20understanding=20of=20the=20expected=20response.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/backend/langflow/api/v1/chat.py    | 21 ++++++++++++---------
 src/backend/langflow/api/v1/schemas.py |  8 ++++++++
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/src/backend/langflow/api/v1/chat.py b/src/backend/langflow/api/v1/chat.py
index 5ce1ff2ab..52c8e9a67 100644
--- a/src/backend/langflow/api/v1/chat.py
+++ b/src/backend/langflow/api/v1/chat.py
@@ -6,7 +6,8 @@ from fastapi import (
     WebSocketException,
     status,
 )
-from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.responses import StreamingResponse
+from langflow.api.v1.schemas import BuiltResponse, InitResponse
 
 from langflow.chat.manager import ChatManager
 from langflow.graph.graph.base import Graph
@@ -31,7 +32,7 @@ async def chat(client_id: str, websocket: WebSocket):
         await websocket.close(code=status.WS_1011_INTERNAL_ERROR, reason=str(exc))
 
 
-@router.post("/build/init")
+@router.post("/build/init", response_model=InitResponse)
 async def init_build(graph_data: dict):
     """Initialize the build by storing graph data and returning a unique session ID."""
 
@@ -39,19 +40,21 @@ async def init_build(graph_data: dict):
 
     flow_data_store[flow_id] = graph_data
 
-    return JSONResponse(content={"flowId": flow_id})
+    return InitResponse(flowId=flow_id)
 
 
-@router.get("/build/{flow_id}/status")
+@router.get("/build/{flow_id}/status", response_model=BuiltResponse)
 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 and not isinstance(
+        built = 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})
+        )
+
+        return BuiltResponse(
+            built=built,
+        )
+
     except Exception as exc:
         logger.error(exc)
         return HTTPException(status_code=500, detail=str(exc))
diff --git a/src/backend/langflow/api/v1/schemas.py b/src/backend/langflow/api/v1/schemas.py
index f73b0642d..161704738 100644
--- a/src/backend/langflow/api/v1/schemas.py
+++ b/src/backend/langflow/api/v1/schemas.py
@@ -68,3 +68,11 @@ class FileResponse(ChatMessage):
         if v not in ["image", "csv"]:
             raise ValueError("data_type must be image or csv")
         return v
+
+
+class InitResponse(BaseModel):
+    flowId: str
+
+
+class BuiltResponse(BaseModel):
+    built: bool