embedding types response update

Author: alekhya-n

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,
@@ -407,9 +407,7 @@ def embed(
             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": [],
-        }
+        embedding_responses = EmbeddingResponses()
         json_bodys = []
 
         for i in range(0, len(texts), self.batch_size):
@@ -426,11 +424,12 @@ def embed(
 
         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"])
+            embedding_responses.add_response(result)
             meta = result["meta"] if not meta else meta
 
         return Embeddings(
-            embeddings=responses["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,
@@ -303,9 +304,13 @@ async def embed(
         ]
         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"]],
+            embeddings=embedding_responses.get_embeddings(),
+            response_type=embedding_responses.response_type,
             meta=meta,
         )
 

cohere/responses/__init__.py

@@ -5,7 +5,7 @@
 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 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)