From 2d4642fc36cfd93f7e3d79ee9e0ce4a3c8f5a0b4 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sat, 8 Jul 2023 16:49:30 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20fix(callback.py):=20add=20error?= =?UTF-8?q?=20handling=20when=20sending=20response=20to=20websocket=20to?= =?UTF-8?q?=20prevent=20potential=20errors=20=F0=9F=94=A7=20fix(callback.p?= =?UTF-8?q?y):=20change=20intermediate=5Fsteps=20assignment=20to=20use=20f?= =?UTF-8?q?ormatted=20string=20for=20better=20readability=20and=20maintain?= =?UTF-8?q?ability=20=E2=9C=A8=20feat(callback.py):=20add=20observation=5F?= =?UTF-8?q?prefix=20parameter=20to=20on=5Ftool=5Fend=20method=20to=20allow?= =?UTF-8?q?=20customization=20of=20the=20observation=20prefix=20in=20the?= =?UTF-8?q?=20response=20message=20=E2=9C=A8=20feat(callback.py):=20add=20?= =?UTF-8?q?logger=20to=20handle=20potential=20errors=20when=20sending=20re?= =?UTF-8?q?sponse=20to=20websocket?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/api/v1/callback.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/src/backend/langflow/api/v1/callback.py b/src/backend/langflow/api/v1/callback.py index 03f76543e..4b82dc580 100644 --- a/src/backend/langflow/api/v1/callback.py +++ b/src/backend/langflow/api/v1/callback.py @@ -4,12 +4,14 @@ from langchain.callbacks.base import AsyncCallbackHandler, BaseCallbackHandler from langflow.api.v1.schemas import ChatResponse +from typing import Any, Dict from typing import Any, Dict, List, Union from fastapi import WebSocket from langchain.schema import AgentAction, LLMResult, AgentFinish +from langflow.utils.logger import logger # https://github.com/hwchase17/chat-langchain/blob/master/callback.py @@ -62,12 +64,23 @@ class AsyncStreamingLLMCallbackHandler(AsyncCallbackHandler): async def on_tool_end(self, output: str, **kwargs: Any) -> Any: """Run when tool ends running.""" + observation_prefix = kwargs.get("observation_prefix", "Tool output: ") + + # Create a formatted message. + intermediate_steps = f"{observation_prefix}{output}" + + # Create a ChatResponse instance. resp = ChatResponse( message="", type="stream", - intermediate_steps=f"Tool output: {output}", + intermediate_steps=intermediate_steps, ) - await self.websocket.send_json(resp.dict()) + + # Try to send the response, handle potential errors. + try: + await self.websocket.send_json(resp.dict()) + except Exception as e: + logger.error(e) async def on_tool_error( self, error: Union[Exception, KeyboardInterrupt], **kwargs: Any