🔀 refactor(chat.py): import InitResponse and BuiltResponse from schemas module

🚀 feat(chat.py): add response models to /build/init and /build/{flow_id}/status endpoints
The InitResponse and BuiltResponse models are now imported from the schemas module to improve code organization. The /build/init and /build/{flow_id}/status endpoints now have response models to provide a clear understanding of the expected response.
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-06-13 17:21:45 -03:00
commit 41f6aecda4
2 changed files with 20 additions and 9 deletions

View file

@ -6,7 +6,8 @@ from fastapi import (
WebSocketException, WebSocketException,
status, 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.chat.manager import ChatManager
from langflow.graph.graph.base import Graph 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)) 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): async def init_build(graph_data: dict):
"""Initialize the build by storing graph data and returning a unique session ID.""" """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 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): async def build_status(flow_id: str):
"""Check the flow_id is in the flow_data_store.""" """Check the flow_id is in the flow_data_store."""
try: 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 flow_data_store[flow_id], dict
): )
return JSONResponse(content={"built": True})
else: return BuiltResponse(
return JSONResponse(content={"built": False}) built=built,
)
except Exception as exc: except Exception as exc:
logger.error(exc) logger.error(exc)
return HTTPException(status_code=500, detail=str(exc)) return HTTPException(status_code=500, detail=str(exc))

View file

@ -68,3 +68,11 @@ class FileResponse(ChatMessage):
if v not in ["image", "csv"]: if v not in ["image", "csv"]:
raise ValueError("data_type must be image or csv") raise ValueError("data_type must be image or csv")
return v return v
class InitResponse(BaseModel):
flowId: str
class BuiltResponse(BaseModel):
built: bool