From 842628771bed31884925e2ae64fafef5cfee85e7 Mon Sep 17 00:00:00 2001 From: pablodanswer Date: Mon, 16 Sep 2024 12:07:52 -0700 Subject: [PATCH] minor robustification for search --- backend/danswer/llm/answering/answer.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/backend/danswer/llm/answering/answer.py b/backend/danswer/llm/answering/answer.py index 9fd8503f3..6d180b234 100644 --- a/backend/danswer/llm/answering/answer.py +++ b/backend/danswer/llm/answering/answer.py @@ -204,7 +204,11 @@ class Answer: ]: tool_calls = 0 initiated = False - while tool_calls < MAX_TOOL_CALLS: + forced_tool_use = ( + self.force_use_tool.force_use and self.force_use_tool.args is not None + ) + + while tool_calls < (1 if forced_tool_use else MAX_TOOL_CALLS): if initiated: yield StreamStopInfo(stop_reason=StreamStopReason.NEW_RESPONSE) initiated = True @@ -213,7 +217,7 @@ class Answer: tool_call_chunk: AIMessageChunk | None = None - if self.force_use_tool.force_use and self.force_use_tool.args is not None: + if forced_tool_use: tool_call_chunk = AIMessageChunk(content="") tool_call_chunk.tool_calls = [ { @@ -438,7 +442,7 @@ class Answer: ]: tool_calls = 0 initiated = False - while tool_calls < MAX_TOOL_CALLS: + while tool_calls < (1 if self.force_use_tool.force_use else MAX_TOOL_CALLS): if initiated: yield StreamStopInfo(stop_reason=StreamStopReason.NEW_RESPONSE)