From 07ab8ab9313bf6bc9be993658c81eb77a6ed53b2 Mon Sep 17 00:00:00 2001 From: sabaimran Date: Thu, 17 Oct 2024 09:00:01 -0700 Subject: [PATCH] Update handling of gemini response with new API changes. Per documentation: finish_reason (google.ai.generativelanguage_v1beta.types.Candidate.FinishReason): Optional. Output only. The reason why the model stopped generating tokens. If empty, the model has not stopped generating the tokens. --- src/khoj/processor/conversation/google/utils.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/khoj/processor/conversation/google/utils.py b/src/khoj/processor/conversation/google/utils.py index 7ef99a92..5679ba4d 100644 --- a/src/khoj/processor/conversation/google/utils.py +++ b/src/khoj/processor/conversation/google/utils.py @@ -148,6 +148,10 @@ def handle_gemini_response(candidates, prompt_feedback=None): elif candidates[0].finish_reason == FinishReason.SAFETY: message = generate_safety_response(candidates[0].safety_ratings) stopped = True + # Check if finish reason is empty, therefore generation is in progress + elif not candidates[0].finish_reason: + message = None + stopped = False # Check if the response was stopped due to reaching maximum token limit or other reasons elif candidates[0].finish_reason != FinishReason.STOP: message = f"\nI can't talk further about that because of **{candidates[0].finish_reason.name} issue.**"