Merge pull request #672 from gquental/gemini-vertexai

Update Gemini version and SDK
vanna-ai · Oct 23, 2024 · b7b030a · b7b030a
2 parents 932131f + baf3dbe
commit b7b030a
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
@@ -15,7 +15,7 @@ def __init__(self, config=None):
         if "model_name" in config:
             model_name = config["model_name"]
         else:
-            model_name = "gemini-1.0-pro"
+            model_name = "gemini-1.5-pro"
 
         self.google_api_key = None
 
@@ -30,7 +30,7 @@ def __init__(self, config=None):
             self.chat_model = genai.GenerativeModel(model_name)
         else:
             # Authenticate using VertexAI
-            from vertexai.preview.generative_models import GenerativeModel
+            from vertexai.generative_models import GenerativeModel
             self.chat_model = GenerativeModel(model_name)
 
     def system_message(self, message: str) -> any: