mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
fix: Enable usage metadata in LiteLLM streaming
Closes #3181 Co-authored-by: Eliza Huang <heliza@google.com> PiperOrigin-RevId: 825833660
This commit is contained in:
committed by
Copybara-Service
parent
01b48c09ad
commit
f9569bbb1a
@@ -920,6 +920,7 @@ class LiteLlm(BaseLlm):
|
|||||||
# Track function calls by index
|
# Track function calls by index
|
||||||
function_calls = {} # index -> {name, args, id}
|
function_calls = {} # index -> {name, args, id}
|
||||||
completion_args["stream"] = True
|
completion_args["stream"] = True
|
||||||
|
completion_args["stream_options"] = {"include_usage": True}
|
||||||
aggregated_llm_response = None
|
aggregated_llm_response = None
|
||||||
aggregated_llm_response_with_tool_call = None
|
aggregated_llm_response_with_tool_call = None
|
||||||
usage_metadata = None
|
usage_metadata = None
|
||||||
|
|||||||
@@ -1606,6 +1606,40 @@ async def test_generate_content_async_stream_with_usage_metadata(
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_generate_content_async_stream_with_usage_metadata_only(
|
||||||
|
mock_completion, lite_llm_instance
|
||||||
|
):
|
||||||
|
streaming_model_response_with_usage_metadata = [
|
||||||
|
ModelResponse(
|
||||||
|
usage={
|
||||||
|
"prompt_tokens": 10,
|
||||||
|
"completion_tokens": 5,
|
||||||
|
"total_tokens": 15,
|
||||||
|
},
|
||||||
|
choices=[
|
||||||
|
StreamingChoices(
|
||||||
|
finish_reason="stop",
|
||||||
|
delta=Delta(content=""),
|
||||||
|
)
|
||||||
|
],
|
||||||
|
),
|
||||||
|
]
|
||||||
|
mock_completion.return_value = iter(
|
||||||
|
streaming_model_response_with_usage_metadata
|
||||||
|
)
|
||||||
|
|
||||||
|
unused_responses = [
|
||||||
|
response
|
||||||
|
async for response in lite_llm_instance.generate_content_async(
|
||||||
|
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
|
||||||
|
)
|
||||||
|
]
|
||||||
|
mock_completion.assert_called_once()
|
||||||
|
_, kwargs = mock_completion.call_args
|
||||||
|
assert kwargs["stream_options"] == {"include_usage": True}
|
||||||
|
|
||||||
|
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_generate_content_async_multiple_function_calls(
|
async def test_generate_content_async_multiple_function_calls(
|
||||||
mock_completion, lite_llm_instance
|
mock_completion, lite_llm_instance
|
||||||
|
|||||||
Reference in New Issue
Block a user