From cf0c31ef3b2637c7c8c8c1bfe924b728bb4bea41 Mon Sep 17 00:00:00 2001 From: yurekami Date: Mon, 29 Dec 2025 06:17:36 +0900 Subject: [PATCH] Add image_tokens field to ApiMetaBilledUnits MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The embedding API returns an `image_tokens` field in the billed_units response that was not present in the model definition. This caused the field to only be accessible through extra attributes. Changes: - Added `image_tokens` field to `ApiMetaBilledUnits` model - Updated `merge_meta_field` utility to properly merge image_tokens - Added tests for image_tokens field and merging Fixes #711 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- src/cohere/types/api_meta_billed_units.py | 5 ++ src/cohere/utils.py | 4 +- tests/test_embed_utils.py | 59 +++++++++++++++++++++++ 3 files changed, 67 insertions(+), 1 deletion(-) diff --git a/src/cohere/types/api_meta_billed_units.py b/src/cohere/types/api_meta_billed_units.py index 7b7e38af0..6772595d1 100644 --- a/src/cohere/types/api_meta_billed_units.py +++ b/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: diff --git a/src/cohere/utils.py b/src/cohere/utils.py index 35d8b852b..80ee52f9d 100644 --- a/src/cohere/utils.py +++ b/src/cohere/utils.py @@ -174,6 +174,7 @@ 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, @@ -181,7 +182,8 @@ def merge_meta_field(metas: typing.List[ApiMeta]) -> ApiMeta: input_tokens=input_tokens, output_tokens=output_tokens, search_units=search_units, - classifications=classifications + classifications=classifications, + image_tokens=image_tokens ), warnings=list(warnings) ) diff --git a/tests/test_embed_utils.py b/tests/test_embed_utils.py index 40c712177..212e35f35 100644 --- a/tests/test_embed_utils.py +++ b/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)