From bd06fcd9bee7318dc031dcd6fbc35c1dcccd0940 Mon Sep 17 00:00:00 2001
From: Debanjum <debanjum@gmail.com>
Date: Tue, 11 Mar 2025 23:38:27 +0530
Subject: [PATCH] Stop using old google generativeai package to raise, catch
 exceptions

---
 src/khoj/processor/conversation/google/utils.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/khoj/processor/conversation/google/utils.py b/src/khoj/processor/conversation/google/utils.py
index 4c01be2c..2a5f994f 100644
--- a/src/khoj/processor/conversation/google/utils.py
+++ b/src/khoj/processor/conversation/google/utils.py
@@ -4,8 +4,8 @@ from copy import deepcopy
 from threading import Thread
 
 from google import genai
+from google.genai import errors as gerrors
 from google.genai import types as gtypes
-from google.generativeai.types.generation_types import StopCandidateException
 from langchain.schema import ChatMessage
 from tenacity import (
     before_sleep_log,
@@ -73,7 +73,7 @@ def gemini_completion_with_backoff(
         # Generate the response
         response = client.models.generate_content(model=model_name, config=config, contents=formatted_messages)
         response_text = response.text
-    except StopCandidateException as e:
+    except gerrors.ClientError as e:
         response = None
         response_text, _ = handle_gemini_response(e.args)
         # Respond with reason for stopping
@@ -147,7 +147,7 @@ def gemini_llm_thread(
             aggregated_response += message
             g.send(message)
             if stopped:
-                raise StopCandidateException(message)
+                raise ValueError(message)
 
         # Calculate cost of chat
         input_tokens = chunk.usage_metadata.prompt_token_count
@@ -159,7 +159,7 @@ def gemini_llm_thread(
         tracer["temperature"] = temperature
         if is_promptrace_enabled():
             commit_conversation_trace(messages, aggregated_response, tracer)
-    except StopCandidateException as e:
+    except ValueError as e:
         logger.warning(
             f"LLM Response Prevented for {model_name}: {e.args[0]}.\n"
             + f"Last Message by {messages[-1].role}: {messages[-1].content}"