fix: Enable usage metadata in LiteLLM streaming

Closes #3181

Co-authored-by: Eliza Huang <heliza@google.com>
PiperOrigin-RevId: 825833660

tests/unittests/models/test_litellm.py

@@ -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
 async def test_generate_content_async_multiple_function_calls(
     mock_completion, lite_llm_instance