diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index cab2bed8a60..e3dc3988b5c 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -225,6 +225,7 @@ class VertexAITextEmbeddingConfig(LiteLLMBaseModel): # Import here to avoid circular import issues with litellm.__init__ from litellm.llms.vertex_ai.vertex_embeddings.bge import VertexBGEConfig + from litellm.types.utils import PromptTokensDetailsWrapper if VertexBGEConfig.is_bge_model(model): return VertexBGEConfig.transform_response(response=response, model=model, model_response=model_response) @@ -247,7 +248,21 @@ class VertexAITextEmbeddingConfig(LiteLLMBaseModel): model_response.object = "list" model_response.data = embedding_response model_response.model = model - usage: Final = Usage(prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens) + + # Extract billable character count from metadata if present + metadata: Final = response.get("metadata") or {} + billable_character_count: Final = metadata.get("billableCharacterCount") + if billable_character_count is not None: + prompt_tokens_details: Final = PromptTokensDetailsWrapper(character_count=billable_character_count) + usage: Final = Usage( + prompt_tokens=input_tokens, + completion_tokens=0, + total_tokens=input_tokens, + prompt_tokens_details=prompt_tokens_details, + ) + else: + usage: Final = Usage(prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens) + setattr(model_response, "usage", usage) return model_response diff --git a/tests/unit/llms/vertex_ai/test_vertex_text_embedding_billable_character_count.py b/tests/unit/llms/vertex_ai/test_vertex_text_embedding_billable_character_count.py new file mode 100644 index 00000000000..b1d449380b1 --- /dev/null +++ b/tests/unit/llms/vertex_ai/test_vertex_text_embedding_billable_character_count.py @@ -0,0 +1,123 @@ +""" +Test that Vertex AI text embedding responses preserve billableCharacterCount +in usage.prompt_tokens_details.character_count. + +Fixes: https://github.com/BerriAI/litellm/issues/45999 +""" + +import pytest + +from litellm.llms.vertex_ai.vertex_embeddings.transformation import ( + VertexAITextEmbeddingConfig, +) +from litellm.types.utils import EmbeddingResponse + + +class TestVertexTextEmbeddingBillableCharacterCount: + def setup_method(self): + self.config = VertexAITextEmbeddingConfig() + + def test_billable_character_count_mapped_to_usage_details(self): + """ + Test that metadata.billableCharacterCount from Vertex is mapped to + usage.prompt_tokens_details.character_count. + """ + # Sample Vertex AI text embedding response with billable character count + vertex_response = { + "predictions": [ + { + "embeddings": { + "statistics": {"truncated": False, "token_count": 8}, + "values": [0.1, 0.2, 0.3], # Sample embedding values + } + }, + { + "embeddings": { + "statistics": {"truncated": False, "token_count": 1}, + "values": [0.4, 0.5, 0.6], + } + }, + ], + "metadata": {"billableCharacterCount": 21}, + } + + model_response = EmbeddingResponse() + result = self.config.transform_vertex_response_to_openai( + response=vertex_response, + model="text-embedding-005", + model_response=model_response, + ) + + # Verify usage tokens + assert result.usage.prompt_tokens == 9 # 8 + 1 + assert result.usage.completion_tokens == 0 + assert result.usage.total_tokens == 9 + + # Verify billable character count is preserved + assert result.usage.prompt_tokens_details is not None + assert result.usage.prompt_tokens_details.character_count == 21 + + def test_no_character_count_when_metadata_missing(self): + """ + Test that usage.prompt_tokens_details is not set when metadata + is missing from the response. + """ + # Vertex response without metadata + vertex_response = { + "predictions": [ + { + "embeddings": { + "statistics": {"truncated": False, "token_count": 5}, + "values": [0.1, 0.2, 0.3], + } + } + ] + } + + model_response = EmbeddingResponse() + result = self.config.transform_vertex_response_to_openai( + response=vertex_response, + model="text-embedding-005", + model_response=model_response, + ) + + # Verify usage tokens + assert result.usage.prompt_tokens == 5 + assert result.usage.completion_tokens == 0 + assert result.usage.total_tokens == 5 + + # Verify no prompt_tokens_details when metadata is absent + assert result.usage.prompt_tokens_details is None + + def test_no_character_count_when_billable_count_missing(self): + """ + Test that usage.prompt_tokens_details is not set when metadata exists + but billableCharacterCount is missing. + """ + # Vertex response with metadata but no billableCharacterCount + vertex_response = { + "predictions": [ + { + "embeddings": { + "statistics": {"truncated": False, "token_count": 3}, + "values": [0.1, 0.2], + } + } + ], + "metadata": {"someOtherField": "value"}, + } + + model_response = EmbeddingResponse() + result = self.config.transform_vertex_response_to_openai( + response=vertex_response, + model="gemini-embedding-001", + model_response=model_response, + ) + + # Verify usage tokens + assert result.usage.prompt_tokens == 3 + assert result.usage.completion_tokens == 0 + assert result.usage.total_tokens == 3 + + # Verify no prompt_tokens_details when billableCharacterCount is absent + assert result.usage.prompt_tokens_details is None