From 76a49c866354ee8354e7d10637095eef0d795421 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 12 Jun 2023 07:49:49 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(manager.py):=20change=20stat?= =?UTF-8?q?us=20parameter=20type=20from=20status=20to=20int=20in=20close?= =?UTF-8?q?=5Fconnection=20method=20=F0=9F=94=A5=20refactor(manager.py):?= =?UTF-8?q?=20remove=20unused=20build=20method=20and=20rename=20set=5Fcach?= =?UTF-8?q?e=20to=20set=5Fclient=5Fcache=20=F0=9F=9A=80=20chore(manager.py?= =?UTF-8?q?):=20add=20async=20keyword=20to=20close=5Fconnection=20method?= =?UTF-8?q?=20The=20close=5Fconnection=20method=20now=20accepts=20an=20int?= =?UTF-8?q?eger=20as=20the=20status=20parameter=20instead=20of=20a=20statu?= =?UTF-8?q?s=20enum.=20The=20build=20method=20was=20removed=20as=20it=20wa?= =?UTF-8?q?s=20not=20being=20used.=20The=20set=5Fcache=20method=20was=20re?= =?UTF-8?q?named=20to=20set=5Fclient=5Fcache=20to=20improve=20semantics.?= =?UTF-8?q?=20The=20close=5Fconnection=20method=20was=20updated=20to=20be?= =?UTF-8?q?=20an=20async=20method.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/chat/manager.py | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/src/backend/langflow/chat/manager.py b/src/backend/langflow/chat/manager.py index 4169ecc92..614588968 100644 --- a/src/backend/langflow/chat/manager.py +++ b/src/backend/langflow/chat/manager.py @@ -103,7 +103,7 @@ class ChatManager: websocket = self.active_connections[client_id] await websocket.send_json(message.dict()) - async def close_connection(self, client_id: str, code: status, reason: str): + async def close_connection(self, client_id: str, code: int, reason: str): if websocket := self.active_connections[client_id]: await websocket.close(code=code, reason=reason) self.disconnect(client_id) @@ -159,14 +159,11 @@ class ChatManager: await self.send_json(client_id, response) self.chat_history.add_message(client_id, response) - def build(self, client_id: str, graph_data: Dict) -> bool: + def set_cache(self, client_id: str, langchain_object: Any) -> bool: """ - Build the langchain object and set the streaming options, - then store it in the in-memory cache. + Set the cache for a client. """ - logger.debug("Building langchain object") - graph = Graph.from_payload(graph_data) - langchain_object = graph.build() + self.in_memory_cache.set(client_id, langchain_object) return client_id in self.in_memory_cache @@ -191,7 +188,7 @@ class ChatManager: with self.cache_manager.set_client_id(client_id): if client_id not in self.in_memory_cache: - self.close_connection( + await self.close_connection( client_id=client_id, code=status.WS_1011_INTERNAL_ERROR, reason="Please, build the flow before sending messages", @@ -210,10 +207,11 @@ class ChatManager: ) finally: try: - connection = self.active_connections.get(client_id) - if connection: - await connection.close(code=1000, reason="Client disconnected") - self.disconnect(client_id) + await self.close_connection( + client_id=client_id, + code=status.WS_1000_NORMAL_CLOSURE, + reason="Client disconnected", + ) except Exception as e: logger.exception(e) self.disconnect(client_id)