udpate embed request with embedding_types

CHANGELOG.md

@@ -1,4 +1,9 @@
 # Changelog
+## 4.42
+ - [#359] (https://github.com/cohere-ai/cohere-python/pull/359)
+    - Add embedding_types to embed request
+    - Update embed response with embeddings_by_types response
+
 ## 4.41
  - [#358] (https://github.com/cohere-ai/cohere-python/pull/363)
     - Update dataset route to datasets

cohere/client.py

@@ -47,7 +47,7 @@
 from cohere.responses.dataset import BaseDataset, Dataset, DatasetUsage, ParseInfo
 from cohere.responses.detectlang import DetectLanguageResponse, Language
 from cohere.responses.embed_job import EmbedJob
-from cohere.responses.embeddings import Embeddings
+from cohere.responses.embeddings import EmbeddingResponses, Embeddings
 from cohere.responses.feedback import (
     GenerateFeedbackResponse,
     GeneratePreferenceFeedbackResponse,
@@ -395,22 +395,19 @@ 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": [],
-        }
+        embedding_responses = EmbeddingResponses()
         json_bodys = []
 
         for i in range(0, len(texts), self.batch_size):
@@ -420,20 +417,19 @@ 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", []))
+            embedding_responses.add_response(result)
             meta = result["meta"] if not meta else meta
 
         return Embeddings(
-            embeddings=responses["embeddings"],
-            compressed_embeddings=responses["compressed_embeddings"],
+            embeddings=embedding_responses.get_embeddings(),
+            response_type=embedding_responses.response_type,
             meta=meta,
         )
 

cohere/client_async.py

@@ -30,6 +30,7 @@
     Connector,
     DetectLanguageResponse,
     Detokenization,
+    EmbeddingResponses,
     Embeddings,
     GenerateFeedbackResponse,
     GeneratePreferenceFeedbackResponse,
@@ -279,36 +280,37 @@ 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)
         ]
         responses = await asyncio.gather(*[self._request(cohere.EMBED_URL, json) for json in json_bodys])
         meta = responses[0]["meta"] if responses else None
+        embedding_responses = EmbeddingResponses()
+        for response in responses:
+            embedding_responses.add_response(response)
 
         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,
+            embeddings=embedding_responses.get_embeddings(),
+            response_type=embedding_responses.response_type,
             meta=meta,
         )
 

cohere/responses/__init__.py

@@ -5,7 +5,12 @@
 from cohere.responses.connector import Connector, ConnectorOAuth, ConnectorServiceAuth
 from cohere.responses.dataset import AsyncDataset, Dataset
 from cohere.responses.detectlang import DetectLanguageResponse, Language
-from cohere.responses.embeddings import Embeddings
+from cohere.responses.embeddings import (
+    EMBEDDINGS_BY_TYPE_RESPONSE_TYPE,
+    EMBEDDINGS_FLOATS_RESPONSE_TYPE,
+    EmbeddingResponses,
+    Embeddings,
+)
 from cohere.responses.feedback import (
     GenerateFeedbackResponse,
     GeneratePreferenceFeedbackResponse,

cohere/responses/embeddings.py

@@ -1,21 +1,65 @@
-from typing import Any, Dict, Iterator, List, Optional
+from typing import Any, Dict, Iterator, List, Optional, Union
 
 from cohere.responses.base import CohereObject
 
+EMBEDDINGS_FLOATS_RESPONSE_TYPE = "embeddings_floats"
+EMBEDDINGS_BY_TYPE_RESPONSE_TYPE = "embeddings_by_type"
+
+
+class EmbeddingsByType(CohereObject):
+    def __init__(
+        self,
+        float: Optional[List[List[float]]] = None,
+        int8: Optional[List[List[int]]] = None,
+        uint8: Optional[List[List[int]]] = None,
+        binary: Optional[List[List[int]]] = None,
+        ubinary: Optional[List[List[int]]] = None,
+    ) -> None:
+        self.float = float
+        self.int8 = int8
+        self.uint8 = uint8
+        self.binary = binary
+        self.ubinary = ubinary
+
 
 class Embeddings(CohereObject):
     def __init__(
         self,
-        embeddings: List[List[float]],
-        compressed_embeddings: List[List[int]] = None,
+        embeddings: Union[List[List[float]], EmbeddingsByType],
+        response_type: str,
         meta: Optional[Dict[str, Any]] = None,
     ) -> None:
+        self.response_type = response_type
         self.embeddings = embeddings
-        self.compressed_embeddings = compressed_embeddings
         self.meta = meta
 
     def __iter__(self) -> Iterator:
         return iter(self.embeddings)
 
     def __len__(self) -> int:
         return len(self.embeddings)
+
+
+class EmbeddingResponses:
+    def __init__(
+        self,
+    ) -> None:
+        self.response_type = None
+        self.embeddings_floats = []
+        self.embeddings_by_type = {}
+
+    def add_response(self, response):
+        self.response_type = response["response_type"]
+        if self.response_type == EMBEDDINGS_FLOATS_RESPONSE_TYPE:
+            self.embeddings_floats.extend(response["embeddings"])
+        elif self.response_type == EMBEDDINGS_BY_TYPE_RESPONSE_TYPE:
+            for k, v in response["embeddings"].items():
+                if k not in self.embeddings_by_type:
+                    self.embeddings_by_type[k] = []
+                self.embeddings_by_type[k].extend(v)
+
+    def get_embeddings(self) -> Union[List[List[float]], EmbeddingsByType]:
+        if self.response_type == EMBEDDINGS_FLOATS_RESPONSE_TYPE:
+            return self.embeddings_floats
+        elif self.response_type == EMBEDDINGS_BY_TYPE_RESPONSE_TYPE:
+            return EmbeddingsByType(**self.embeddings_by_type)

pyproject.toml

@@ -1,5 +1,6 @@
 [tool.poetry]
 name = "cohere"
+version = "4.42"
 version = "4.41"
 description = "Python SDK for the Cohere API"
 authors = ["Cohere"]

tests/async/test_async_embed.py

@@ -1,13 +1,14 @@
 import pytest
 
-from cohere.responses import Embeddings
+from cohere.responses import EMBEDDINGS_FLOATS_RESPONSE_TYPE, Embeddings
 
 
 @pytest.mark.asyncio
 async def test_async_embed_basic(async_client):
     embs = await async_client.embed(model="small", texts=["co:here", "cohere"])
     assert len(embs) == 2
     assert isinstance(embs, Embeddings)
+    assert embs.response_type == EMBEDDINGS_FLOATS_RESPONSE_TYPE
     assert embs.meta
     assert embs.meta["api_version"]
     assert embs.meta["api_version"]["version"]

tests/sync/test_embed.py

@@ -1,10 +1,16 @@
 import random
 import string
 import unittest
+from typing import List
 
 from utils import get_api_key
 
 import cohere
+from cohere.responses.embeddings import (
+    EMBEDDINGS_BY_TYPE_RESPONSE_TYPE,
+    EMBEDDINGS_FLOATS_RESPONSE_TYPE,
+    EmbeddingsByType,
+)
 
 API_KEY = get_api_key()
 co = cohere.Client(API_KEY)
@@ -102,3 +108,25 @@ def test_success_multiple_batches_in_order(self):
         for predictionExpected, predictionActual in zip(predictionsExpected, list(predictionsActual)):
             for elementExpected, elementAcutal in zip(predictionExpected, predictionActual):
                 self.assertAlmostEqual(elementExpected, elementAcutal, places=1)
+
+    def test_success_embeddings_floats_response_type(self):
+        prediction = co.embed(
+            model="small",
+            texts=["python", "golang", "typescript"],
+        )
+        self.assertEqual(prediction.response_type, EMBEDDINGS_FLOATS_RESPONSE_TYPE)
+        assert isinstance(prediction.embeddings, List)
+
+    def test_success_embeddings_by_type_response_type(self):
+        prediction = co.embed(
+            model="small",
+            texts=["python", "golang", "typescript"],
+            embedding_types=["float"],
+        )
+        self.assertEqual(prediction.response_type, EMBEDDINGS_BY_TYPE_RESPONSE_TYPE)
+        assert isinstance(prediction.embeddings, EmbeddingsByType)
+        assert prediction.embeddings.float
+        assert not prediction.embeddings.int8
+        assert not prediction.embeddings.uint8
+        assert not prediction.embeddings.binary
+        assert not prediction.embeddings.ubinary