migrates mistral vectorizer to new client (#255)

Author: justin-cechmanek

docs/user_guide/vectorizers_04.ipynb

@@ -31,7 +31,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -305,33 +305,25 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/Users/tyler.hutcherson/redis/redis-vl-python/.venv/lib/python3.9/site-packages/torch/_utils.py:831: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
-      "  return self.fget.__get__(instance, owner)()\n"
-     ]
-    },
     {
      "data": {
       "text/plain": [
-       "[0.00037810884532518685,\n",
-       " -0.05080341175198555,\n",
-       " -0.03514723479747772,\n",
-       " -0.02325104922056198,\n",
-       " -0.044158220291137695,\n",
-       " 0.020487844944000244,\n",
-       " 0.0014617963461205363,\n",
-       " 0.031261757016181946,\n",
+       "[0.0003780885017476976,\n",
+       " -0.05080340430140495,\n",
+       " -0.035147231072187424,\n",
+       " -0.02325103059411049,\n",
+       " -0.04415831342339516,\n",
+       " 0.02048780582845211,\n",
+       " 0.0014618589775636792,\n",
+       " 0.03126184269785881,\n",
        " 0.05605152249336243,\n",
-       " 0.018815357238054276]"
+       " 0.018815429881215096]"
       ]
      },
-     "execution_count": 6,
+     "execution_count": 3,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -532,14 +524,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 +580,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 +836,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.8.13 ('redisvl2')",
+   "display_name": "redisvl-dev",
    "language": "python",
    "name": "python3"
   },
@@ -852,12 +852,7 @@
    "pygments_lexer": "ipython3",
    "version": "3.12.2"
   },
-  "orig_nbformat": 4,
-  "vscode": {
-   "interpreter": {
-    "hash": "9b1e6e9c2967143209c2f955cb869d1d3234f92dc4787f49f155f3abbdfb1316"
-   }
-  }
+  "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 2

poetry.lock

@@ -32,7 +32,7 @@ sentence-transformers = { version = ">=2.2.2", optional = true }
 google-cloud-aiplatform = { version = ">=1.26", optional = true }
 protobuf = { version = ">=5.29.1,<6.0.0.dev0", optional = true }
 cohere = { version = ">=4.44", optional = true }
-mistralai = { version = ">=0.2.0", optional = true }
+mistralai = { version = ">=1.0.0", optional = true }
 boto3 = { version = ">=1.34.0", optional = true }
 
 [tool.poetry.extras]

pyproject.toml

@@ -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

redisvl/utils/vectorize/text/mistral.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,
     ]
 )