Add image_tokens field to ApiMetaBilledUnits

src/cohere/types/api_meta_billed_units.py

@@ -33,6 +33,11 @@ class ApiMetaBilledUnits(UncheckedBaseModel):
     The number of billed classifications units.
     """
 
+    image_tokens: typing.Optional[float] = pydantic.Field(default=None)
+    """
+    The number of billed image tokens.
+    """
+
     if IS_PYDANTIC_V2:
         model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow")  # type: ignore # Pydantic v2
     else:

src/cohere/utils.py

@@ -174,14 +174,16 @@ def merge_meta_field(metas: typing.List[ApiMeta]) -> ApiMeta:
     output_tokens = sum_fields_if_not_none(billed_units, "output_tokens")
     search_units = sum_fields_if_not_none(billed_units, "search_units")
     classifications = sum_fields_if_not_none(billed_units, "classifications")
+    image_tokens = sum_fields_if_not_none(billed_units, "image_tokens")
     warnings = {warning for meta in metas if meta.warnings for warning in meta.warnings}
     return ApiMeta(
         api_version=api_version,
         billed_units=ApiMetaBilledUnits(
             input_tokens=input_tokens,
             output_tokens=output_tokens,
             search_units=search_units,
-            classifications=classifications
+            classifications=classifications,
+            image_tokens=image_tokens
         ),
         warnings=list(warnings)
     )

tests/test_embed_utils.py

@@ -219,3 +219,62 @@ def test_merge_partial_embeddings_floats(self) -> None:
                 warnings=resp.meta.warnings  # order ignored
             )
         ))
+
+    def test_image_tokens_field(self) -> None:
+        """Test that image_tokens field is properly handled in ApiMetaBilledUnits.
+
+        This is a regression test for issue #711 where the image_tokens field
+        was missing from the response model.
+        """
+        # Test that image_tokens can be set and accessed
+        billed_units = ApiMetaBilledUnits(
+            input_tokens=100,
+            image_tokens=2716
+        )
+        self.assertEqual(billed_units.image_tokens, 2716)
+        self.assertEqual(billed_units.input_tokens, 100)
+
+        # Test serialization includes image_tokens
+        dumped = billed_units.model_dump()
+        self.assertEqual(dumped["image_tokens"], 2716)
+
+    def test_merge_with_image_tokens(self) -> None:
+        """Test that image_tokens are properly merged across responses."""
+        ebt_with_images_1 = EmbeddingsByTypeEmbedResponse(
+            response_type="embeddings_by_type",
+            id="1",
+            embeddings=EmbedByTypeResponseEmbeddings(
+                float_=[[0, 1, 2]],
+            ),
+            texts=["hello"],
+            meta=ApiMeta(
+                api_version=ApiMetaApiVersion(version="1"),
+                billed_units=ApiMetaBilledUnits(
+                    input_tokens=1,
+                    image_tokens=100
+                ),
+            )
+        )
+
+        ebt_with_images_2 = EmbeddingsByTypeEmbedResponse(
+            response_type="embeddings_by_type",
+            id="2",
+            embeddings=EmbedByTypeResponseEmbeddings(
+                float_=[[3, 4, 5]],
+            ),
+            texts=["goodbye"],
+            meta=ApiMeta(
+                api_version=ApiMetaApiVersion(version="1"),
+                billed_units=ApiMetaBilledUnits(
+                    input_tokens=2,
+                    image_tokens=200
+                ),
+            )
+        )
+
+        resp = merge_embed_responses([ebt_with_images_1, ebt_with_images_2])
+
+        self.assertIsNotNone(resp.meta)
+        self.assertIsNotNone(resp.meta.billed_units)
+        self.assertEqual(resp.meta.billed_units.image_tokens, 300)
+        self.assertEqual(resp.meta.billed_units.input_tokens, 3)