udpate embed request with embedding_types

Author: alekhya-n

cohere/client.py

@@ -395,21 +395,20 @@ def embed(
         texts: List[str],
         model: Optional[str] = None,
         truncate: Optional[str] = None,
-        compression: Optional[str] = None,
         input_type: Optional[str] = None,
+        embedding_types: Optional[List[str]] = None,
     ) -> Embeddings:
         """Returns an Embeddings object for the provided texts. Visit https://cohere.ai/embed to learn about embeddings.
 
         Args:
             text (List[str]): A list of strings to embed.
             model (str): (Optional) The model ID to use for embedding the text.
             truncate (str): (Optional) One of NONE|START|END, defaults to END. How the API handles text longer than the maximum token length.
-            compression (str): (Optional) One of "int8" or "binary". The type of compression to use for the embeddings.
             input_type (str): (Optional) One of "classification", "clustering", "search_document", "search_query". The type of input text provided to embed.
+            embedding_types (List[str]): (Optional) Specifies the types of embeddings you want to get back. Not required and default is None, which returns the float embeddings in the response's embeddings field. Can be one or more of the following types: "float", "int8", "uint8", "binary", "ubinary".
         """
         responses = {
             "embeddings": [],
-            "compressed_embeddings": [],
         }
         json_bodys = []
 
@@ -420,20 +419,18 @@ def embed(
                     "model": model,
                     "texts": texts_batch,
                     "truncate": truncate,
-                    "compression": compression,
                     "input_type": input_type,
+                    "embedding_types": embedding_types,
                 }
             )
 
         meta = None
         for result in self._executor.map(lambda json_body: self._request(cohere.EMBED_URL, json=json_body), json_bodys):
             responses["embeddings"].extend(result["embeddings"])
-            responses["compressed_embeddings"].extend(result.get("compressed_embeddings", []))
             meta = result["meta"] if not meta else meta
 
         return Embeddings(
             embeddings=responses["embeddings"],
-            compressed_embeddings=responses["compressed_embeddings"],
             meta=meta,
         )
 

cohere/client_async.py

@@ -279,25 +279,25 @@ async def embed(
         texts: List[str],
         model: Optional[str] = None,
         truncate: Optional[str] = None,
-        compression: Optional[str] = None,
         input_type: Optional[str] = None,
+        embedding_types: Optional[List[str]] = None,
     ) -> Embeddings:
         """Returns an Embeddings object for the provided texts. Visit https://cohere.ai/embed to learn about embeddings.
 
         Args:
             text (List[str]): A list of strings to embed.
             model (str): (Optional) The model ID to use for embedding the text.
             truncate (str): (Optional) One of NONE|START|END, defaults to END. How the API handles text longer than the maximum token length.
-            compression (str): (Optional) One of "int8" or "binary". The type of compression to use for the embeddings.
             input_type (str): (Optional) One of "classification", "clustering", "search_document", "search_query". The type of input text provided to embed.
+            embedding_types (List[str]): (Optional) Specifies the types of embeddings you want to get back. Not required and default is None, which returns the float embeddings in the response's embeddings field. Can be one or more of the following types: "float", "int8", "uint8", "binary", "ubinary".
         """
         json_bodys = [
             dict(
                 texts=texts[i : i + cohere.COHERE_EMBED_BATCH_SIZE],
                 model=model,
                 truncate=truncate,
-                compression=compression,
                 input_type=input_type,
+                embedding_types=embedding_types,
             )
             for i in range(0, len(texts), cohere.COHERE_EMBED_BATCH_SIZE)
         ]
@@ -306,9 +306,6 @@ async def embed(
 
         return Embeddings(
             embeddings=[e for res in responses for e in res["embeddings"]],
-            compressed_embeddings=[e for res in responses for e in res["compressed_embeddings"]]
-            if compression
-            else None,
             meta=meta,
         )