add kwargs support to all vectorizer embed methods

Author: tylerhutcherson

redisvl/utils/vectorize/text/azureopenai.py

@@ -205,7 +205,9 @@ def embed_many(
 
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
-            response = self._client.embeddings.create(input=batch, model=self.model)
+            response = self._client.embeddings.create(
+                input=batch, model=self.model, **kwargs
+            )
             embeddings += [
                 self._process_embedding(r.embedding, as_buffer, dtype)
                 for r in response.data
@@ -248,7 +250,9 @@ def embed(
 
         dtype = kwargs.pop("dtype", self.dtype)
 
-        result = self._client.embeddings.create(input=[text], model=self.model)
+        result = self._client.embeddings.create(
+            input=[text], model=self.model, **kwargs
+        )
         return self._process_embedding(result.data[0].embedding, as_buffer, dtype)
 
     @retry(
@@ -292,7 +296,7 @@ async def aembed_many(
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = await self._aclient.embeddings.create(
-                input=batch, model=self.model
+                input=batch, model=self.model, **kwargs
             )
             embeddings += [
                 self._process_embedding(r.embedding, as_buffer, dtype)
@@ -336,7 +340,9 @@ async def aembed(
 
         dtype = kwargs.pop("dtype", self.dtype)
 
-        result = await self._aclient.embeddings.create(input=[text], model=self.model)
+        result = await self._aclient.embeddings.create(
+            input=[text], model=self.model, **kwargs
+        )
         return self._process_embedding(result.data[0].embedding, as_buffer, dtype)
 
     @property

redisvl/utils/vectorize/text/bedrock.py

@@ -156,7 +156,7 @@ def embed(
             text = preprocess(text)
 
         response = self._client.invoke_model(
-            modelId=self.model, body=json.dumps({"inputText": text})
+            modelId=self.model, body=json.dumps({"inputText": text}), **kwargs
         )
         response_body = json.loads(response["body"].read())
         embedding = response_body["embedding"]
@@ -206,7 +206,7 @@ def embed_many(
             batch_embeddings = []
             for text in batch:
                 response = self._client.invoke_model(
-                    modelId=self.model, body=json.dumps({"inputText": text})
+                    modelId=self.model, body=json.dumps({"inputText": text}), **kwargs
                 )
                 response_body = json.loads(response["body"].read())
                 batch_embeddings.append(response_body["embedding"])

redisvl/utils/vectorize/text/cohere.py

@@ -156,7 +156,7 @@ def embed(
             TypeError: In an invalid input_type is provided.
 
         """
-        input_type = kwargs.get("input_type")
+        input_type = kwargs.pop("input_type", None)
 
         if not isinstance(text, str):
             raise TypeError("Must pass in a str value to embed.")
@@ -172,7 +172,7 @@ def embed(
         dtype = kwargs.pop("dtype", self.dtype)
 
         embedding = self._client.embed(
-            texts=[text], model=self.model, input_type=input_type
+            texts=[text], model=self.model, input_type=input_type, **kwargs
         ).embeddings[0]
         return self._process_embedding(embedding, as_buffer, dtype)
 
@@ -227,7 +227,7 @@ def embed_many(
             TypeError: In an invalid input_type is provided.
 
         """
-        input_type = kwargs.get("input_type")
+        input_type = kwargs.pop("input_type", None)
 
         if not isinstance(texts, list):
             raise TypeError("Must pass in a list of str values to embed.")
@@ -244,7 +244,7 @@ def embed_many(
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = self._client.embed(
-                texts=batch, model=self.model, input_type=input_type
+                texts=batch, model=self.model, input_type=input_type, **kwargs
             )
             embeddings += [
                 self._process_embedding(embedding, as_buffer, dtype)

redisvl/utils/vectorize/text/mistral.py

@@ -155,7 +155,9 @@ def embed_many(
 
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
-            response = self._client.embeddings.create(model=self.model, inputs=batch)
+            response = self._client.embeddings.create(
+                model=self.model, inputs=batch, **kwargs
+            )
             embeddings += [
                 self._process_embedding(r.embedding, as_buffer, dtype)
                 for r in response.data
@@ -198,7 +200,9 @@ def embed(
 
         dtype = kwargs.pop("dtype", self.dtype)
 
-        result = self._client.embeddings.create(model=self.model, inputs=[text])
+        result = self._client.embeddings.create(
+            model=self.model, inputs=[text], **kwargs
+        )
         return self._process_embedding(result.data[0].embedding, as_buffer, dtype)
 
     @retry(
@@ -242,7 +246,7 @@ async def aembed_many(
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = await self._client.embeddings.create_async(
-                model=self.model, inputs=batch
+                model=self.model, inputs=batch, **kwargs
             )
             embeddings += [
                 self._process_embedding(r.embedding, as_buffer, dtype)
@@ -287,7 +291,7 @@ async def aembed(
         dtype = kwargs.pop("dtype", self.dtype)
 
         result = await self._client.embeddings.create_async(
-            model=self.model, inputs=[text]
+            model=self.model, inputs=[text], **kwargs
         )
         return self._process_embedding(result.data[0].embedding, as_buffer, dtype)
 

redisvl/utils/vectorize/text/openai.py

@@ -156,7 +156,9 @@ def embed_many(
 
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
-            response = self._client.embeddings.create(input=batch, model=self.model)
+            response = self._client.embeddings.create(
+                input=batch, model=self.model, **kwargs
+            )
             embeddings += [
                 self._process_embedding(r.embedding, as_buffer, dtype)
                 for r in response.data
@@ -199,7 +201,9 @@ def embed(
 
         dtype = kwargs.pop("dtype", self.dtype)
 
-        result = self._client.embeddings.create(input=[text], model=self.model)
+        result = self._client.embeddings.create(
+            input=[text], model=self.model, **kwargs
+        )
         return self._process_embedding(result.data[0].embedding, as_buffer, dtype)
 
     @retry(
@@ -243,7 +247,7 @@ async def aembed_many(
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = await self._aclient.embeddings.create(
-                input=batch, model=self.model
+                input=batch, model=self.model, **kwargs
             )
             embeddings += [
                 self._process_embedding(r.embedding, as_buffer, dtype)
@@ -287,7 +291,9 @@ async def aembed(
 
         dtype = kwargs.pop("dtype", self.dtype)
 
-        result = await self._aclient.embeddings.create(input=[text], model=self.model)
+        result = await self._aclient.embeddings.create(
+            input=[text], model=self.model, **kwargs
+        )
         return self._process_embedding(result.data[0].embedding, as_buffer, dtype)
 
     @property

redisvl/utils/vectorize/text/vertexai.py

@@ -168,7 +168,7 @@ def embed_many(
 
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
-            response = self._client.get_embeddings(batch)
+            response = self._client.get_embeddings(batch, **kwargs)
             embeddings += [
                 self._process_embedding(r.values, as_buffer, dtype) for r in response
             ]
@@ -210,7 +210,7 @@ def embed(
 
         dtype = kwargs.pop("dtype", self.dtype)
 
-        result = self._client.get_embeddings([text])
+        result = self._client.get_embeddings([text], **kwargs)
         return self._process_embedding(result[0].values, as_buffer, dtype)
 
     @property

redisvl/utils/vectorize/text/voyageai.py

@@ -204,8 +204,8 @@ def embed_many(
             TypeError: If an invalid input_type is provided.
 
         """
-        input_type = kwargs.get("input_type")
-        truncation = kwargs.get("truncation")
+        input_type = kwargs.pop("input_type", None)
+        truncation = kwargs.pop("truncation", None)
         dtype = kwargs.pop("dtype", self.dtype)
 
         if not isinstance(texts, list):
@@ -235,7 +235,7 @@ def embed_many(
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = self._client.embed(
-                texts=batch, model=self.model, input_type=input_type
+                texts=batch, model=self.model, input_type=input_type, **kwargs
             )
             embeddings += [
                 self._process_embedding(embedding, as_buffer, dtype)
@@ -284,8 +284,8 @@ async def aembed_many(
             TypeError: In an invalid input_type is provided.
 
         """
-        input_type = kwargs.get("input_type")
-        truncation = kwargs.get("truncation")
+        input_type = kwargs.pop("input_type", None)
+        truncation = kwargs.pop("truncation", None)
         dtype = kwargs.pop("dtype", self.dtype)
 
         if not isinstance(texts, list):
@@ -315,7 +315,7 @@ async def aembed_many(
         embeddings: List = []
         for batch in self.batchify(texts, batch_size, preprocess):
             response = await self._aclient.embed(
-                texts=batch, model=self.model, input_type=input_type
+                texts=batch, model=self.model, input_type=input_type, **kwargs
             )
             embeddings += [
                 self._process_embedding(embedding, as_buffer, dtype)
@@ -360,7 +360,6 @@ async def aembed(
         Raises:
             TypeError: In an invalid input_type is provided.
         """
-
         result = await self.aembed_many(
             texts=[text], preprocess=preprocess, as_buffer=as_buffer, **kwargs
         )