Skip to content

Commit f9569bb

Browse files
lizzijcopybara-github
authored andcommitted
fix: Enable usage metadata in LiteLLM streaming
Closes #3181 Co-authored-by: Eliza Huang <[email protected]> PiperOrigin-RevId: 825833660
1 parent 01b48c0 commit f9569bb

File tree

2 files changed

+35
-0
lines changed

2 files changed

+35
-0
lines changed

src/google/adk/models/lite_llm.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -920,6 +920,7 @@ async def generate_content_async(
920920
# Track function calls by index
921921
function_calls = {} # index -> {name, args, id}
922922
completion_args["stream"] = True
923+
completion_args["stream_options"] = {"include_usage": True}
923924
aggregated_llm_response = None
924925
aggregated_llm_response_with_tool_call = None
925926
usage_metadata = None

tests/unittests/models/test_litellm.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1606,6 +1606,40 @@ async def test_generate_content_async_stream_with_usage_metadata(
16061606
)
16071607

16081608

1609+
@pytest.mark.asyncio
1610+
async def test_generate_content_async_stream_with_usage_metadata_only(
1611+
mock_completion, lite_llm_instance
1612+
):
1613+
streaming_model_response_with_usage_metadata = [
1614+
ModelResponse(
1615+
usage={
1616+
"prompt_tokens": 10,
1617+
"completion_tokens": 5,
1618+
"total_tokens": 15,
1619+
},
1620+
choices=[
1621+
StreamingChoices(
1622+
finish_reason="stop",
1623+
delta=Delta(content=""),
1624+
)
1625+
],
1626+
),
1627+
]
1628+
mock_completion.return_value = iter(
1629+
streaming_model_response_with_usage_metadata
1630+
)
1631+
1632+
unused_responses = [
1633+
response
1634+
async for response in lite_llm_instance.generate_content_async(
1635+
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
1636+
)
1637+
]
1638+
mock_completion.assert_called_once()
1639+
_, kwargs = mock_completion.call_args
1640+
assert kwargs["stream_options"] == {"include_usage": True}
1641+
1642+
16091643
@pytest.mark.asyncio
16101644
async def test_generate_content_async_multiple_function_calls(
16111645
mock_completion, lite_llm_instance

0 commit comments

Comments
 (0)