From ba8eb6fca42105cb03335a9bf5c1d4fd761056fa Mon Sep 17 00:00:00 2001 From: pablodanswer Date: Mon, 2 Sep 2024 14:03:35 -0700 Subject: [PATCH] add stashed reversion for tool call chunks --- backend/danswer/llm/answering/answer.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/backend/danswer/llm/answering/answer.py b/backend/danswer/llm/answering/answer.py index 95c49c33856..684b0aab786 100644 --- a/backend/danswer/llm/answering/answer.py +++ b/backend/danswer/llm/answering/answer.py @@ -230,11 +230,24 @@ def _raw_output_for_explicit_tool_calling_llms( self.tools, self.force_use_tool ) ] - yield from self._process_llm_stream( + + for message in self.llm.stream( prompt=prompt, tools=final_tool_definitions if final_tool_definitions else None, tool_choice="required" if self.force_use_tool.force_use else None, - ) + ): + if isinstance(message, AIMessageChunk) and ( + message.tool_call_chunks or message.tool_calls + ): + if tool_call_chunk is None: + tool_call_chunk = message + else: + tool_call_chunk += message # type: ignore + else: + if message.content: + if self.is_cancelled: + return + yield cast(str, message.content) if not tool_call_chunk: return # no tool call needed