feat: GenAI - Added the GenerationResponse.usage_metadata property

PiperOrigin-RevId: 625552531

Author: Ark-kun

tests/unit/vertexai/test_generative_models.py

@@ -246,6 +246,20 @@ def mock_generate_content(
         finish_reason = gapic_content_types.Candidate.FinishReason.OTHER
         response.candidates[0].finish_reason = finish_reason
 
+    request_token_count = sum(
+        len(gapic_content_types.Content.to_json(content).split())
+        for content in request.contents
+    )
+    response_token_count = sum(
+        len(gapic_content_types.Content.to_json(candidate.content).split())
+        for candidate in response.candidates
+    )
+    response.usage_metadata.prompt_token_count = request_token_count
+    response.usage_metadata.candidates_token_count = response_token_count
+    response.usage_metadata.total_token_count = (
+        request_token_count + response_token_count
+    )
+
     return response
 
 
@@ -371,6 +385,8 @@ def test_generate_content(self, generative_models: generative_models):
         model = generative_models.GenerativeModel("gemini-pro")
         response = model.generate_content("Why is sky blue?")
         assert response.text
+        # TODO(avolkov): Add usage metadata to the mock
+        assert response.usage_metadata.total_token_count
 
         model2 = generative_models.GenerativeModel(
             "gemini-pro",

vertexai/generative_models/_generative_models.py

@@ -1648,6 +1648,12 @@ def prompt_feedback(
     ) -> gapic_prediction_service_types.GenerateContentResponse.PromptFeedback:
         return self._raw_response.prompt_feedback
 
+    @property
+    def usage_metadata(
+        self
+    ) -> gapic_prediction_service_types.GenerateContentResponse.UsageMetadata:
+        return self._raw_response.usage_metadata
+
 
 class Candidate:
     """A response candidate generated by the model."""