Add support for MistralAI (#823)

* Add support for MistralAI

* embeddings

* lint

Author: jtpio

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,
-  Gemini, Hugging Face, NVIDIA, and OpenAI.
+  Gemini, Hugging Face, MistralAI, NVIDIA, and OpenAI.
 * Local model support through GPT4All, enabling use of generative AI models on consumer grade machines
   with ease and privacy.
 

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, Gemini, Hugging Face, OpenAI, SageMaker, NVIDIA, etc.).
+  (AI21, Anthropic, Cohere, Gemini, Hugging Face, MistralAI, 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'

docs/source/users/index.md

@@ -148,6 +148,7 @@ Jupyter AI supports the following model providers:
 | 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` |
+| MistralAI           | `mistralai`          | `MISTRAL_API_KEY`          | `langchain-mistralai`           |
 | NVIDIA              | `nvidia-chat`        | `NVIDIA_API_KEY`           | `langchain_nvidia_ai_endpoints` |
 | OpenAI              | `openai`             | `OPENAI_API_KEY`           | `langchain-openai`              |
 | OpenAI (chat)       | `openai-chat`        | `OPENAI_API_KEY`           | `langchain-openai`              |

packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/mistralai.py

@@ -0,0 +1,32 @@
+from jupyter_ai_magics.providers import BaseProvider, EnvAuthStrategy
+from langchain_mistralai import ChatMistralAI, MistralAIEmbeddings
+
+from ..embedding_providers import BaseEmbeddingsProvider
+
+
+class MistralAIProvider(BaseProvider, ChatMistralAI):
+    id = "mistralai"
+    name = "MistralAI"
+    models = [
+        "open-mistral-7b",
+        "open-mixtral-8x7b",
+        "open-mixtral-8x22b",
+        "mistral-small-latest",
+        "mistral-medium-latest",
+        "mistral-large-latest",
+        "codestral-latest",
+    ]
+    model_id_key = "model"
+    auth_strategy = EnvAuthStrategy(name="MISTRAL_API_KEY")
+    pypi_package_deps = ["langchain-mistralai"]
+
+
+class MistralAIEmbeddingsProvider(BaseEmbeddingsProvider, MistralAIEmbeddings):
+    id = "mistralai"
+    name = "MistralAI"
+    models = [
+        "mistral-embed",
+    ]
+    model_id_key = "model"
+    pypi_package_deps = ["langchain-mistralai"]
+    auth_strategy = EnvAuthStrategy(name="MISTRAL_API_KEY")

packages/jupyter-ai-magics/pyproject.toml

@@ -41,13 +41,14 @@ all = [
     "huggingface_hub",
     "ipywidgets",
     "langchain_anthropic",
+    "langchain-mistralai",
     "langchain_nvidia_ai_endpoints",
+    "langchain-google-genai",
     "langchain-openai",
     "pillow",
     "boto3",
     "qianfan",
     "together",
-    "langchain-google-genai",
 ]
 
 [project.entry-points."jupyter_ai.model_providers"]
@@ -67,10 +68,12 @@ 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"
+mistralai = "jupyter_ai_magics.partner_providers.mistralai:MistralAIProvider"
 
 [project.entry-points."jupyter_ai.embeddings_model_providers"]
 bedrock = "jupyter_ai_magics:BedrockEmbeddingsProvider"
 cohere = "jupyter_ai_magics:CohereEmbeddingsProvider"
+mistralai = "jupyter_ai_magics.partner_providers.mistralai:MistralAIEmbeddingsProvider"
 gpt4all = "jupyter_ai_magics:GPT4AllEmbeddingsProvider"
 huggingface_hub = "jupyter_ai_magics:HfHubEmbeddingsProvider"
 openai = "jupyter_ai_magics.partner_providers.openai:OpenAIEmbeddingsProvider"