From 9a67a496226269477c7f9b1d0bb0043e323bad5f Mon Sep 17 00:00:00 2001 From: Gabriel Almeida Date: Fri, 5 May 2023 16:38:40 -0300 Subject: [PATCH] feat(chat.py, chat_manager.py): add try-except block to handle WebSocketException and WebSocketDisconnect This commit adds a try-except block to handle WebSocketException and WebSocketDisconnect in the `websocket_endpoint` function in `chat.py` and `handle_websocket` function in `chat_manager.py`. If an exception occurs, the function logs the error and closes the websocket connection with the appropriate status code and reason. Additionally, the `handle_websocket` function in `chat_manager.py` now waits for a final message after sending a message to the client. If the message is a string, it is a chat message, which is then added to the chat history. --- src/backend/langflow/api/chat.py | 11 +++++++++-- src/backend/langflow/api/chat_manager.py | 25 ++++++++++++++++-------- 2 files changed, 26 insertions(+), 10 deletions(-) diff --git a/src/backend/langflow/api/chat.py b/src/backend/langflow/api/chat.py index 48a195179..436f64f7c 100644 --- a/src/backend/langflow/api/chat.py +++ b/src/backend/langflow/api/chat.py @@ -2,6 +2,7 @@ from fastapi import APIRouter, WebSocket from langflow.api.chat_manager import ChatManager from langflow.utils.logger import logger +from fastapi import status, WebSocketDisconnect, WebSocketException router = APIRouter() chat_manager = ChatManager() @@ -10,5 +11,11 @@ chat_manager = ChatManager() @router.websocket("/chat/{client_id}") async def websocket_endpoint(client_id: str, websocket: WebSocket): """Websocket endpoint for chat.""" - - await chat_manager.handle_websocket(client_id, websocket) + try: + await chat_manager.handle_websocket(client_id, 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)) diff --git a/src/backend/langflow/api/chat_manager.py b/src/backend/langflow/api/chat_manager.py index 280f51b3a..e4a4a1b4e 100644 --- a/src/backend/langflow/api/chat_manager.py +++ b/src/backend/langflow/api/chat_manager.py @@ -106,18 +106,18 @@ class ChatManager: # Process the graph data and chat message chat_message = payload.pop("message", "") chat_message = ChatMessage(message=chat_message) - self.chat_history.add_message(client_id, chat_message) + await self.send_json(client_id, chat_message) graph_data = payload start_resp = ChatResponse(message=None, type="start", intermediate_steps="") - self.chat_history.add_message(client_id, start_resp) + await self.send_json(client_id, start_resp) is_first_message = len(self.chat_history.get_history(client_id=client_id)) == 0 # Generate result and thought try: logger.debug("Generating result and thought") - result, intermediate_steps = await process_graph( + _, intermediate_steps = await process_graph( graph_data=graph_data, is_first_message=is_first_message, chat_message=chat_message, @@ -149,7 +149,7 @@ class ChatManager: type="end", files=file_responses, ) - self.chat_history.add_message(client_id, response) + await self.send_json(client_id, response) async def handle_websocket(self, client_id: str, websocket: WebSocket): await self.connect(client_id, websocket) @@ -172,15 +172,24 @@ class ChatManager: with self.cache_manager.set_client_id(client_id): await self.process_message(client_id, payload) + # After the message is sent, wait for message built + final_message = await websocket.receive_json() + # If the message is a string, it is a chat message + chat_response = ChatResponse.parse_obj(final_message) + self.chat_history.add_message(client_id, chat_response) + except Exception as e: # Handle any exceptions that might occur logger.exception(e) # send a message to the client - await self.active_connections[client_id].close(code=1000, reason=str(e)) + raise e finally: - # await self.active_connections[client_id].close( - # code=1000, reason="Client disconnected" - # ) + try: + await self.active_connections[client_id].close( + code=1000, reason="Client disconnected" + ) + except Exception as e: + logger.exception(e) self.disconnect(client_id)