Copybara import of the project:

--
e4aceb8428b8d2c321e7b44a16869997dd5f0dff by ajdaling <ajdaling@gmail.com>:

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
This commit is contained in:
Alec Daling 2025-05-13 17:09:53 -07:00 committed by Copybara-Service
parent 66553b18c6
commit 1f0fd7bfce

View File

@ -85,6 +85,9 @@ class LlmResponse(BaseModel):
NOTE: the entire dict must be JSON serializable. NOTE: the entire dict must be JSON serializable.
""" """
usage_metadata: Optional[types.GenerateContentResponseUsageMetadata] = None
"""The usage metadata of the LlmResponse"""
@staticmethod @staticmethod
def create( def create(
generate_content_response: types.GenerateContentResponse, generate_content_response: types.GenerateContentResponse,
@ -98,18 +101,20 @@ class LlmResponse(BaseModel):
Returns: Returns:
The LlmResponse. The LlmResponse.
""" """
usage_metadata = generate_content_response.usage_metadata
if generate_content_response.candidates: if generate_content_response.candidates:
candidate = generate_content_response.candidates[0] candidate = generate_content_response.candidates[0]
if candidate.content and candidate.content.parts: if candidate.content and candidate.content.parts:
return LlmResponse( return LlmResponse(
content=candidate.content, content=candidate.content,
grounding_metadata=candidate.grounding_metadata, grounding_metadata=candidate.grounding_metadata,
usage_metadata=usage_metadata,
) )
else: else:
return LlmResponse( return LlmResponse(
error_code=candidate.finish_reason, error_code=candidate.finish_reason,
error_message=candidate.finish_message, error_message=candidate.finish_message,
usage_metadata=usage_metadata,
) )
else: else:
if generate_content_response.prompt_feedback: if generate_content_response.prompt_feedback:
@ -117,9 +122,11 @@ class LlmResponse(BaseModel):
return LlmResponse( return LlmResponse(
error_code=prompt_feedback.block_reason, error_code=prompt_feedback.block_reason,
error_message=prompt_feedback.block_reason_message, error_message=prompt_feedback.block_reason_message,
usage_metadata=usage_metadata,
) )
else: else:
return LlmResponse( return LlmResponse(
error_code='UNKNOWN_ERROR', error_code='UNKNOWN_ERROR',
error_message='Unknown error.', error_message='Unknown error.',
usage_metadata=usage_metadata,
) )