Use VertexAI

jupyterlab · Mar 3, 2024 · 7c9ff52 · 7c9ff52
1 parent 341ccbd
commit 7c9ff52
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py b/packages/jupyter-ai-magics/jupyter_ai_magics/providers.py
@@ -38,7 +38,7 @@
     OpenAI,
     SagemakerEndpoint,
 )
-from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAI
+from langchain_google_vertexai import VertexAI
 
 # this is necessary because `langchain.pydantic_v1.main` does not include
 # `ModelMetaclass`, as it is not listed in `__all__` by the `pydantic.main`
@@ -864,14 +864,14 @@ class QianfanProvider(BaseProvider, QianfanChatEndpoint):
     auth_strategy = MultiEnvAuthStrategy(names=["QIANFAN_AK", "QIANFAN_SK"])
 
 
-class GeminiProvider(BaseProvider, GoogleGenerativeAI):
+class GeminiProvider(BaseProvider, VertexAI):
     id = "gemini"
     name = "Gemini"
     models = [
         "gemini-pro",
     ]
-    model_id_key = "model"
-    pypi_package_deps = ["langchain-google-genai"]
+    model_id_key = "model_name"
+    pypi_package_deps = ["langchain-google-vertexai"]
     auth_strategy = EnvAuthStrategy(name="GOOGLE_API_KEY")
 
     @classmethod

diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml
@@ -46,7 +46,7 @@ all = [
     "openai~=1.6.1",
     "boto3",
     "qianfan",
-    "langchain_google_genai",
+    "langchain-google-vertexai",
 ]
 
 [project.entry-points."jupyter_ai.model_providers"]