Merge pull request #535 from everdark/fix/gemini-chat-model-default

fix: model name for gemini chat is hardcoded when no api-key is used
vanna-ai · Jul 11, 2024 · f0c3c18 · f0c3c18
2 parents 2ab4e87 + f7ebfc5
commit f0c3c18
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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