Add support for Gemini (jupyterlab#666)

* Add Gemini Provider * Rename GeminiAI * Use VertexAI * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Change back to langchain-google-genai * Add gemini as a partner provider * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add more gemini models * Add Gemini to docs --------- Co-authored-by: David L. Qiu <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Marchlak · Oct 28, 2024 · c06d71e · c06d71e
1 parent 07d2b42
commit c06d71e
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -10,7 +10,7 @@ in JupyterLab and the Jupyter Notebook. More specifically, Jupyter AI offers:
   This works anywhere the IPython kernel runs (JupyterLab, Jupyter Notebook, Google Colab, Kaggle, VSCode, etc.).
 * A native chat UI in JupyterLab that enables you to work with generative AI as a conversational assistant.
 * Support for a wide range of generative model providers, including AI21, Anthropic, AWS, Cohere,
-  Hugging Face, NVIDIA, and OpenAI.
+  Gemini, Hugging Face, NVIDIA, and OpenAI.
 * Local model support through GPT4All, enabling use of generative AI models on consumer grade machines
   with ease and privacy.
 

diff --git a/docs/source/index.md b/docs/source/index.md
@@ -8,7 +8,7 @@ in JupyterLab and the Jupyter Notebook. More specifically, Jupyter AI offers:
   This works anywhere the IPython kernel runs (JupyterLab, Jupyter Notebook, Google Colab, VSCode, etc.).
 * A native chat UI in JupyterLab that enables you to work with generative AI as a conversational assistant.
 * Support for a wide range of generative model providers and models
-  (AI21, Anthropic, Cohere, Hugging Face, OpenAI, SageMaker, NVIDIA, etc.).
+  (AI21, Anthropic, Cohere, Gemini, Hugging Face, OpenAI, SageMaker, NVIDIA, etc.).
 
 <img src="_static/jupyter-ai-screenshot.png"
     alt='A screenshot of Jupyter AI showing the chat interface and the magic commands'

diff --git a/docs/source/users/index.md b/docs/source/users/index.md
@@ -140,6 +140,7 @@ Jupyter AI supports the following model providers:
 | Bedrock (chat)      | `bedrock-chat`       | N/A                        | `boto3`                         |
 | Cohere              | `cohere`             | `COHERE_API_KEY`           | `cohere`                        |
 | ERNIE-Bot           | `qianfan`            | `QIANFAN_AK`, `QIANFAN_SK` | `qianfan`                       |
+| Gemini              | `gemini`             | `GOOGLE_API_KEY`           | `langchain-google-genai`        |
 | GPT4All             | `gpt4all`            | N/A                        | `gpt4all`                       |
 | Hugging Face Hub    | `huggingface_hub`    | `HUGGINGFACEHUB_API_TOKEN` | `huggingface_hub`, `ipywidgets`, `pillow` |
 | NVIDIA              | `nvidia-chat`        | `NVIDIA_API_KEY`           | `langchain_nvidia_ai_endpoints` |

diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/gemini.py b/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/gemini.py
@@ -0,0 +1,18 @@
+from jupyter_ai_magics.providers import BaseProvider, EnvAuthStrategy
+from langchain_google_genai import GoogleGenerativeAI
+
+
+class GeminiProvider(BaseProvider, GoogleGenerativeAI):
+    id = "gemini"
+    name = "Gemini"
+    models = [
+        "gemini-1.0-pro",
+        "gemini-1.0-pro-001",
+        "gemini-1.0-pro-latest",
+        "gemini-1.0-pro-vision-latest",
+        "gemini-pro",
+        "gemini-pro-vision",
+    ]
+    model_id_key = "model"
+    auth_strategy = EnvAuthStrategy(name="GOOGLE_API_KEY")
+    pypi_package_deps = ["langchain-google-genai"]
diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml
@@ -47,7 +47,7 @@ all = [
     "boto3",
     "qianfan",
     "together",
-
+    "langchain-google-genai",
 ]
 
 [project.entry-points."jupyter_ai.model_providers"]
@@ -66,6 +66,7 @@ amazon-bedrock-chat = "jupyter_ai_magics:BedrockChatProvider"
 qianfan = "jupyter_ai_magics:QianfanProvider"
 nvidia-chat = "jupyter_ai_magics.partner_providers.nvidia:ChatNVIDIAProvider"
 together-ai = "jupyter_ai_magics:TogetherAIProvider"
+gemini = "jupyter_ai_magics.partner_providers.gemini:GeminiProvider"
 
 [project.entry-points."jupyter_ai.embeddings_model_providers"]
 bedrock = "jupyter_ai_magics:BedrockEmbeddingsProvider"