Fix bug when partial embeddings by type requested (#440)

* Fix bug when partial embeddings by type requested

* Update src/cohere/utils.py

Co-authored-by: harry-cohere <[email protected]>

---------

Co-authored-by: harry-cohere <[email protected]>

Author: billytrend-cohere

src/cohere/utils.py

@@ -211,13 +211,16 @@ def merge_embed_responses(responses: typing.List[EmbedResponse]) -> EmbedRespons
             for response in embeddings_type
         ]
 
+        # only get set keys from the pydantic model (i.e. exclude fields that are set to 'None')
+        fields = embeddings_type[0].embeddings.dict(exclude_unset=True).keys()
+
         merged_dicts = {
             field: [
                 embedding
                 for embedding_by_type in embeddings_by_type
                 for embedding in getattr(embedding_by_type, field)
             ]
-            for field in EmbedByTypeResponseEmbeddings.__fields__
+            for field in fields
         }
 
         embeddings_by_type_merged = EmbedByTypeResponseEmbeddings.parse_obj(merged_dicts)

tests/test_embed_utils.py

@@ -50,6 +50,48 @@
     )
 )
 
+ebt_partial_1 = EmbedResponse_EmbeddingsByType(
+    response_type="embeddings_by_type",
+    id="1",
+    embeddings=EmbedByTypeResponseEmbeddings(
+        float=[[0, 1, 2], [3, 4, 5]],
+        int8=[[0, 1, 2], [3, 4, 5]],
+        binary=[[5, 6, 7], [8, 9, 10]],
+    ),
+    texts=["hello", "goodbye"],
+    meta=ApiMeta(
+        api_version=ApiMetaApiVersion(version="1"),
+        billed_units=ApiMetaBilledUnits(
+            input_tokens=1,
+            output_tokens=1,
+            search_units=1,
+            classifications=1
+        ),
+        warnings=["test_warning_1"]
+    )
+)
+
+ebt_partial_2 = EmbedResponse_EmbeddingsByType(
+    response_type="embeddings_by_type",
+    id="2",
+    embeddings=EmbedByTypeResponseEmbeddings(
+        float=[[7, 8, 9], [10, 11, 12]],
+        int8=[[7, 8, 9], [10, 11, 12]],
+        binary=[[14, 15, 16], [17, 18, 19]],
+    ),
+    texts=["bye", "seeya"],
+    meta=ApiMeta(
+        api_version=ApiMetaApiVersion(version="1"),
+        billed_units=ApiMetaBilledUnits(
+            input_tokens=2,
+            output_tokens=2,
+            search_units=2,
+            classifications=2
+        ),
+        warnings=["test_warning_1", "test_warning_2"]
+    )
+)
+
 ebf_1 = EmbedResponse_EmbeddingsFloats(
     response_type="embeddings_floats",
     id="1",
@@ -93,7 +135,6 @@ def test_merge_embeddings_by_type(self) -> None:
             ebt_2
         ])
 
-
         if resp.meta is None:
             raise Exception("this is just for mpy")
 
@@ -147,3 +188,34 @@ def test_merge_embeddings_floats(self) -> None:
                 warnings=resp.meta.warnings  # order ignored
             )
         ))
+
+    def test_merge_partial_embeddings_floats(self) -> None:
+        resp = merge_embed_responses([
+            ebt_partial_1,
+            ebt_partial_2
+        ])
+
+        if resp.meta is None:
+            raise Exception("this is just for mpy")
+
+        self.assertEqual(set(resp.meta.warnings or []), {"test_warning_1", "test_warning_2"})
+        self.assertEqual(resp, EmbedResponse_EmbeddingsByType(
+            response_type="embeddings_by_type",
+            id="1, 2",
+            embeddings=EmbedByTypeResponseEmbeddings(
+                float=[[0, 1, 2], [3, 4, 5], [7, 8, 9], [10, 11, 12]],
+                int8=[[0, 1, 2], [3, 4, 5], [7, 8, 9], [10, 11, 12]],
+                binary=[[5, 6, 7], [8, 9, 10], [14, 15, 16], [17, 18, 19]],
+            ),
+            texts=["hello", "goodbye", "bye", "seeya"],
+            meta=ApiMeta(
+                api_version=ApiMetaApiVersion(version="1"),
+                billed_units=ApiMetaBilledUnits(
+                    input_tokens=3,
+                    output_tokens=3,
+                    search_units=3,
+                    classifications=3
+                ),
+                warnings=resp.meta.warnings  # order ignored
+            )
+        ))