jyapayne/langflow
1
0
Fork 0
Code Issues Pull requests Projects Releases 1 Packages Wiki Activity Actions

feat(chat.py, chat_manager.py): add try-except block to handle WebSocketException and WebSocketDisconnect

Browse source 
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.
This commit is contained in:
Gabriel Almeida 2023-05-05 16:38:40 -03:00
commit 9a67a49622
2 changed files with 26 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

11
src/backend/langflow/api/chat.py
View file

@ -2,6 +2,7 @@ from fastapi import APIRouter, WebSocket
from langflow.api.chat_manager import ChatManager from langflow.api.chat_manager import ChatManager
from langflow.utils.logger import logger from langflow.utils.logger import logger
from fastapi import status, WebSocketDisconnect, WebSocketException
router = APIRouter() router = APIRouter()
chat_manager = ChatManager() chat_manager = ChatManager()
@ -10,5 +11,11 @@ chat_manager = ChatManager()
@router.websocket("/chat/{client_id}") @router.websocket("/chat/{client_id}")
async def websocket_endpoint(client_id: str, websocket: WebSocket): async def websocket_endpoint(client_id: str, websocket: WebSocket):
"""Websocket endpoint for chat.""" """Websocket endpoint for chat."""
try:
await chat_manager.handle_websocket(client_id, websocket) 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))

25
src/backend/langflow/api/chat_manager.py
View file

@ -106,18 +106,18 @@ class ChatManager:
# Process the graph data and chat message # Process the graph data and chat message
chat_message = payload.pop("message", "") chat_message = payload.pop("message", "")
chat_message = ChatMessage(message=chat_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 graph_data = payload
start_resp = ChatResponse(message=None, type="start", intermediate_steps="") 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 is_first_message = len(self.chat_history.get_history(client_id=client_id)) == 0
# Generate result and thought # Generate result and thought
try: try:
logger.debug("Generating result and thought") logger.debug("Generating result and thought")
result, intermediate_steps = await process_graph( _, intermediate_steps = await process_graph(
graph_data=graph_data, graph_data=graph_data,
is_first_message=is_first_message, is_first_message=is_first_message,
chat_message=chat_message, chat_message=chat_message,
@ -149,7 +149,7 @@ class ChatManager:
type="end", type="end",
files=file_responses, 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): async def handle_websocket(self, client_id: str, websocket: WebSocket):
await self.connect(client_id, websocket) await self.connect(client_id, websocket)
@ -172,15 +172,24 @@ class ChatManager:
with self.cache_manager.set_client_id(client_id): with self.cache_manager.set_client_id(client_id):
await self.process_message(client_id, payload) 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: except Exception as e:
# Handle any exceptions that might occur # Handle any exceptions that might occur
logger.exception(e) logger.exception(e)
# send a message to the client # send a message to the client
await self.active_connections[client_id].close(code=1000, reason=str(e)) raise e
finally: finally:
# await self.active_connections[client_id].close( try:
# code=1000, reason="Client disconnected" await self.active_connections[client_id].close(
# ) code=1000, reason="Client disconnected"
)
except Exception as e:
logger.exception(e)
self.disconnect(client_id) self.disconnect(client_id)