From 7ac2539b29bcfe099de770057375d2d93e94d22e Mon Sep 17 00:00:00 2001
From: Qiusheng Wu <giswqs@gmail.com>
Date: Tue, 5 Mar 2024 19:52:05 -0500
Subject: [PATCH] Backport PR #666: Add support for Gemini

---
 README.md                                      |  2 +-
 docs/source/index.md                           |  2 +-
 docs/source/users/index.md                     |  1 +
 .../partner_providers/gemini.py                | 18 ++++++++++++++++++
 packages/jupyter-ai-magics/pyproject.toml      |  3 ++-
 5 files changed, 23 insertions(+), 3 deletions(-)
 create mode 100644 packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/gemini.py

diff --git a/README.md b/README.md
index 4a5dfd0a5..901c78a2d 100644
--- 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, 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
index 8d70538e3..a2c060f6a 100644
--- 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
index a481bbc3e..ee4e335c8 100644
--- 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
new file mode 100644
index 000000000..c52e7ebdb
--- /dev/null
+++ 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
index ec0af6ba2..f6637d36c 100644
--- 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"