migrates mistral text vectorizer to new mistral client

Author: justin-cechmanek

docs/user_guide/vectorizers_04.ipynb

@@ -532,14 +532,14 @@
     }
    ],
    "source": [
-    "# from redisvl.utils.vectorize import MistralAITextVectorizer\n",
+    "from redisvl.utils.vectorize import MistralAITextVectorizer\n",
     "\n",
-    "# mistral = MistralAITextVectorizer()\n",
+    "mistral = MistralAITextVectorizer()\n",
     "\n",
-    "# # embed a sentence using their asyncronous method\n",
-    "# test = await mistral.aembed(\"This is a test sentence.\")\n",
-    "# print(\"Vector dimensions: \", len(test))\n",
-    "# print(test[:10])"
+    "# embed a sentence using their asyncronous method\n",
+    "test = await mistral.aembed(\"This is a test sentence.\")\n",
+    "print(\"Vector dimensions: \", len(test))\n",
+    "print(test[:10])"
    ]
   },
   {
@@ -588,9 +588,17 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Vector dimensions: 1024\n"
+     ]
+    }
+   ],
    "source": [
     "from redisvl.utils.vectorize import BedrockTextVectorizer\n",
     "\n",
@@ -836,7 +844,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.8.13 ('redisvl2')",
+   "display_name": "redisvl-dev",
    "language": "python",
    "name": "python3"
   },
@@ -852,12 +860,7 @@
    "pygments_lexer": "ipython3",
    "version": "3.12.2"
   },
-  "orig_nbformat": 4,
-  "vscode": {
-   "interpreter": {
-    "hash": "9b1e6e9c2967143209c2f955cb869d1d3234f92dc4787f49f155f3abbdfb1316"
-   }
-  }
+  "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 2

redisvl/utils/vectorize/text/mistral.py

@@ -44,7 +44,6 @@ class MistralAITextVectorizer(BaseVectorizer):
     """
 
     _client: Any = PrivateAttr()
-    _aclient: Any = PrivateAttr()
 
     def __init__(self, model: str = "mistral-embed", api_config: Optional[Dict] = None):
         """Initialize the MistralAI vectorizer.
@@ -69,8 +68,7 @@ def _initialize_clients(self, api_config: Optional[Dict]):
         """
         # Dynamic import of the mistralai module
         try:
-            from mistralai.async_client import MistralAsyncClient
-            from mistralai.client import MistralClient
+            from mistralai import Mistral
         except ImportError:
             raise ImportError(
                 "MistralAI vectorizer requires the mistralai library. \
@@ -88,13 +86,12 @@ def _initialize_clients(self, api_config: Optional[Dict]):
                     environment variable."
             )
 
-        self._client = MistralClient(api_key=api_key)
-        self._aclient = MistralAsyncClient(api_key=api_key)
+        self._client = Mistral(api_key=api_key)
 
     def _set_model_dims(self, model) -> int:
         try:
             embedding = (
-                self._client.embeddings(model=model, input=["dimension test"])
+                self._client.embeddings.create(model=model, inputs=["dimension test"])
                 .data[0]
                 .embedding
             )
@@ -144,7 +141,7 @@ def embed_many(
 
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
-            response = self._client.embeddings(model=self.model, input=batch)
+            response = self._client.embeddings.create(model=self.model, inputs=batch)
             embeddings += [
                 self._process_embedding(r.embedding, as_buffer, dtype)
                 for r in response.data
@@ -186,7 +183,7 @@ def embed(
 
         dtype = kwargs.pop("dtype", None)
 
-        result = self._client.embeddings(model=self.model, input=[text])
+        result = self._client.embeddings.create(model=self.model, inputs=[text])
         return self._process_embedding(result.data[0].embedding, as_buffer, dtype)
 
     @retry(
@@ -228,7 +225,9 @@ async def aembed_many(
 
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
-            response = await self._aclient.embeddings(model=self.model, input=batch)
+            response = await self._client.embeddings.create_async(
+                model=self.model, inputs=batch
+            )
             embeddings += [
                 self._process_embedding(r.embedding, as_buffer, dtype)
                 for r in response.data
@@ -270,7 +269,9 @@ async def aembed(
 
         dtype = kwargs.pop("dtype", None)
 
-        result = await self._aclient.embeddings(model=self.model, input=[text])
+        result = await self._client.embeddings.create_async(
+            model=self.model, inputs=[text]
+        )
         return self._process_embedding(result.data[0].embedding, as_buffer, dtype)
 
     @property

tests/integration/test_vectorizers.py

@@ -28,7 +28,7 @@ def skip_vectorizer() -> bool:
         CohereTextVectorizer,
         AzureOpenAITextVectorizer,
         BedrockTextVectorizer,
-        # MistralAITextVectorizer,
+        MistralAITextVectorizer,
         CustomTextVectorizer,
     ]
 )
@@ -242,7 +242,7 @@ def bad_return_type(text: str) -> str:
     params=[
         OpenAITextVectorizer,
         BedrockTextVectorizer,
-        # MistralAITextVectorizer,
+        MistralAITextVectorizer,
         CustomTextVectorizer,
     ]
 )