From c0df1f7c3be35d1972438fdf6e595409bec7a33a Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Sun, 9 Jul 2023 11:00:25 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(callback.py):=20fix=20issue?= =?UTF-8?q?=20with=20splitting=20output=20into=20multiple=20ChatResponse?= =?UTF-8?q?=20instances=20=E2=9C=A8=20feat(callback.py):=20split=20the=20o?= =?UTF-8?q?utput=20into=20multiple=20ChatResponse=20instances=20to=20emula?= =?UTF-8?q?te=20a=20stream=20of=20tokens?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/langflow/api/v1/callback.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/src/backend/langflow/api/v1/callback.py b/src/backend/langflow/api/v1/callback.py index 4b82dc580..6250f7987 100644 --- a/src/backend/langflow/api/v1/callback.py +++ b/src/backend/langflow/api/v1/callback.py @@ -65,9 +65,11 @@ 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: ") - + split_output = output.split() + first_word = split_output[0] + rest_of_output = split_output[1:] # Create a formatted message. - intermediate_steps = f"{observation_prefix}{output}" + intermediate_steps = f"{observation_prefix}{first_word}" # Create a ChatResponse instance. resp = ChatResponse( @@ -75,10 +77,21 @@ class AsyncStreamingLLMCallbackHandler(AsyncCallbackHandler): type="stream", intermediate_steps=intermediate_steps, ) - + rest_of_resps = [ + ChatResponse( + message="", + type="stream", + intermediate_steps=f"{word}", + ) + for word in rest_of_output + ] + resps = [resp] + rest_of_resps # Try to send the response, handle potential errors. + try: - await self.websocket.send_json(resp.dict()) + # This is to emulate the stream of tokens + for resp in resps: + await self.websocket.send_json(resp.dict()) except Exception as e: logger.error(e)