From 1f0fd7bfce622459ef9dbc05ec96ffe6b613e4d1 Mon Sep 17 00:00:00 2001 From: Alec Daling Date: Tue, 13 May 2025 17:09:53 -0700 Subject: [PATCH] Copybara import of the project: -- e4aceb8428b8d2c321e7b44a16869997dd5f0dff by ajdaling : Fixes issue #311: Add usage_metadata to llm_response COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/694 from chapelhillnorth:fix-issue-311 2735bc01a6a57a5302b58af33266bfbeafc8e184 PiperOrigin-RevId: 758435162 --- src/google/adk/models/llm_response.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/google/adk/models/llm_response.py b/src/google/adk/models/llm_response.py index 6fcb4e2..311b821 100644 --- a/src/google/adk/models/llm_response.py +++ b/src/google/adk/models/llm_response.py @@ -85,6 +85,9 @@ class LlmResponse(BaseModel): NOTE: the entire dict must be JSON serializable. """ + usage_metadata: Optional[types.GenerateContentResponseUsageMetadata] = None + """The usage metadata of the LlmResponse""" + @staticmethod def create( generate_content_response: types.GenerateContentResponse, @@ -98,18 +101,20 @@ class LlmResponse(BaseModel): Returns: The LlmResponse. """ - + usage_metadata = generate_content_response.usage_metadata if generate_content_response.candidates: candidate = generate_content_response.candidates[0] if candidate.content and candidate.content.parts: return LlmResponse( content=candidate.content, grounding_metadata=candidate.grounding_metadata, + usage_metadata=usage_metadata, ) else: return LlmResponse( error_code=candidate.finish_reason, error_message=candidate.finish_message, + usage_metadata=usage_metadata, ) else: if generate_content_response.prompt_feedback: @@ -117,9 +122,11 @@ class LlmResponse(BaseModel): return LlmResponse( error_code=prompt_feedback.block_reason, error_message=prompt_feedback.block_reason_message, + usage_metadata=usage_metadata, ) else: return LlmResponse( error_code='UNKNOWN_ERROR', error_message='Unknown error.', + usage_metadata=usage_metadata, )