🐛 fix(manager.py): change error message from "Could not find a LangChain object for client_id" to "Could not find a build result for client_id" for better clarity

🐛 fix(manager.py): handle WebSocketState.DISCONNECTED state in the finally block to properly disconnect the client if the connection is already closed
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-09-28 17:17:59 -03:00 • committed by anovazzi1
commit 117d5acbf5

View file

@ -1,6 +1,7 @@
from collections import defaultdict from collections import defaultdict
import uuid import uuid
from fastapi import WebSocket, status from fastapi import WebSocket, status
from starlette.websockets import WebSocketState
from langflow.api.v1.schemas import ChatMessage, ChatResponse, FileResponse from langflow.api.v1.schemas import ChatMessage, ChatResponse, FileResponse
from langflow.interface.utils import pil_to_base64 from langflow.interface.utils import pil_to_base64
from langflow.services.base import Service from langflow.services.base import Service
@ -218,23 +219,29 @@ class ChatService(Service):
else: else:
raise RuntimeError( raise RuntimeError(
f"Could not find a LangChain object for client_id {client_id}" f"Could not find a build result for client_id {client_id}"
) )
except Exception as exc: except Exception as exc:
# Handle any exceptions that might occur # Handle any exceptions that might occur
logger.exception(f"Error handling websocket: {exc}") logger.exception(f"Error handling websocket: {exc}")
await self.close_connection( if websocket.client_state == WebSocketState.CONNECTED:
client_id=client_id,
code=status.WS_1011_INTERNAL_ERROR,
reason=str(exc)[:120],
)
finally:
try:
await self.close_connection( await self.close_connection(
client_id=client_id, client_id=client_id,
code=status.WS_1000_NORMAL_CLOSURE, code=status.WS_1011_INTERNAL_ERROR,
reason="Client disconnected", reason=str(exc)[:120],
) )
elif websocket.client_state == WebSocketState.DISCONNECTED:
self.disconnect(client_id)
finally:
try:
# first check if the connection is still open
if websocket.client_state == WebSocketState.CONNECTED:
await self.close_connection(
client_id=client_id,
code=status.WS_1000_NORMAL_CLOSURE,
reason="Client disconnected",
)
except Exception as exc: except Exception as exc:
logger.error(f"Error closing connection: {exc}") logger.error(f"Error closing connection: {exc}")
self.disconnect(client_id) self.disconnect(client_id)