Merge db6656df1e into 492272a051

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
 

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