From f7ebfc5b76c4bfd03440c640c3f7c43dfd35d205 Mon Sep 17 00:00:00 2001 From: Kyle Chung Date: Wed, 3 Jul 2024 15:02:27 +0700 Subject: [PATCH] fix: model name for gemini chat is hardcoded when no api-key is used --- src/vanna/google/gemini_chat.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/vanna/google/gemini_chat.py b/src/vanna/google/gemini_chat.py index 2a857f00..b33ce5bd 100644 --- a/src/vanna/google/gemini_chat.py +++ b/src/vanna/google/gemini_chat.py @@ -7,7 +7,7 @@ def __init__(self, config=None): VannaBase.__init__(self, config=config) # default temperature - can be overrided using config - self.temperature = 0.7 + self.temperature = 0.7 if "temperature" in config: self.temperature = config["temperature"] @@ -31,7 +31,7 @@ def __init__(self, config=None): else: # Authenticate using VertexAI from vertexai.preview.generative_models import GenerativeModel - self.chat_model = GenerativeModel("gemini-pro") + self.chat_model = GenerativeModel(model_name) def system_message(self, message: str) -> any: return message