File tree Expand file tree Collapse file tree 2 files changed +35
-0
lines changed Expand file tree Collapse file tree 2 files changed +35
-0
lines changed Original file line number Diff line number Diff 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
Original file line number Diff line number Diff 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
16101644async def test_generate_content_async_multiple_function_calls (
16111645 mock_completion , lite_llm_instance
You can’t perform that action at this time.
0 commit comments