embeddings

jupyterlab · Jun 9, 2024 · ff5c545 · ff5c545
1 parent 3753bf0
commit ff5c545
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/mistralai.py b/packages/jupyter-ai-magics/jupyter_ai_magics/partner_providers/mistralai.py
@@ -1,5 +1,7 @@
 from jupyter_ai_magics.providers import BaseProvider, EnvAuthStrategy
-from langchain_mistralai import ChatMistralAI
+from langchain_mistralai import ChatMistralAI, MistralAIEmbeddings
+
+from ..embedding_providers import BaseEmbeddingsProvider
 
 
 class MistralAIProvider(BaseProvider, ChatMistralAI):
@@ -12,9 +14,19 @@ class MistralAIProvider(BaseProvider, ChatMistralAI):
         "mistral-small-latest",
         "mistral-medium-latest",
         "mistral-large-latest",
-        "mistral-embed",
         "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")
diff --git a/packages/jupyter-ai-magics/pyproject.toml b/packages/jupyter-ai-magics/pyproject.toml
@@ -73,6 +73,7 @@ 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"