NeoAnd/adk-python
1
0
Fork 0
mirror of https://github.com/EvolutionAPI/adk-python.git synced 2025-07-16 04:02:55 -06:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add usage_metadata for sse response with finish_reason

Browse Source 
This is needed for covering the generation what ends with finish reason during streaming.

PiperOrigin-RevId: 761144245
This commit is contained in:
Selcuk Gun 2025-05-20 10:46:30 -07:00 committed by Copybara-Service
parent 7d809eb9d9
commit 00798af881
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/google/adk/models/google_llm.py
View File

@ -98,6 +98,7 @@ class Gemini(BaseLlm):
) )
response = None response = None
text = '' text = ''
usage_metadata = None
# for sse, similar as bidi (see receive method in gemini_llm_connecton.py), # for sse, similar as bidi (see receive method in gemini_llm_connecton.py),
# we need to mark those text content as partial and after all partial # we need to mark those text content as partial and after all partial
# contents are sent, we send an accumulated event which contains all the # contents are sent, we send an accumulated event which contains all the
@ -106,6 +107,7 @@ class Gemini(BaseLlm):
async for response in responses: async for response in responses:
logger.info(_build_response_log(response)) logger.info(_build_response_log(response))
llm_response = LlmResponse.create(response) llm_response = LlmResponse.create(response)
usage_metadata = llm_response.usage_metadata
if ( if (
llm_response.content llm_response.content
and llm_response.content.parts and llm_response.content.parts
@ -123,7 +125,7 @@ class Gemini(BaseLlm):
content=types.ModelContent( content=types.ModelContent(
parts=[types.Part.from_text(text=text)], parts=[types.Part.from_text(text=text)],
), ),
usage_metadata=llm_response.usage_metadata, usage_metadata=usage_metadata,
) )
text = '' text = ''
yield llm_response yield llm_response
@ -137,6 +139,7 @@ class Gemini(BaseLlm):
content=types.ModelContent( content=types.ModelContent(
parts=[types.Part.from_text(text=text)], parts=[types.Part.from_text(text=text)],
), ),
usage_metadata=usage_metadata,
) )
else: else: