From bd67e8480f6e8b4b0f8c22b94f15a8cda1336339 Mon Sep 17 00:00:00 2001
From: avidelatm <avidelatm@users.noreply.github.com>
Date: Mon, 23 Jun 2025 10:29:43 -0700
Subject: [PATCH] fix: make LiteLLM streaming truly asynchronous

Merge https://github.com/google/adk-python/pull/1451

## Description
Fixes https://github.com/google/adk-python/issues/1306 by using `async for` with `await self.llm_client.acompletion()` instead of synchronous `for` loop.

## Changes
- Updated test mocks to properly handle async streaming by creating an async generator
- Ensured proper parameter handling to avoid duplicate stream parameter

## Testing Plan
- All unit tests now pass with the async streaming implementation
- Verified with `pytest tests/unittests/models/test_litellm.py` that all streaming tests pass
- Manually tested with a sample agent using LiteLLM to confirm streaming works properly

# Test Evidence:
https://youtu.be/hSp3otI79DM

Let me know if you need anything else from me for this PR

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/1451 from avidelatm:fix/litellm-async-streaming d35b9dc90b2fd6fad44c3869de0fda2514e50055
PiperOrigin-RevId: 774835130
---
 src/google/adk/models/lite_llm.py      |  2 +-
 tests/unittests/models/test_litellm.py | 23 ++++++++++++++++++++---
 2 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/src/google/adk/models/lite_llm.py b/src/google/adk/models/lite_llm.py
index dce5ed7c..acc88ed1 100644
--- a/src/google/adk/models/lite_llm.py
+++ b/src/google/adk/models/lite_llm.py
@@ -679,7 +679,7 @@ class LiteLlm(BaseLlm):
       aggregated_llm_response_with_tool_call = None
       usage_metadata = None
       fallback_index = 0
-      for part in self.llm_client.completion(**completion_args):
+      async for part in await self.llm_client.acompletion(**completion_args):
         for chunk, finish_reason in _model_response_to_chunk(part):
           if isinstance(chunk, FunctionChunk):
             index = chunk.index or fallback_index
diff --git a/tests/unittests/models/test_litellm.py b/tests/unittests/models/test_litellm.py
index 8b43cc48..d058aa44 100644
--- a/tests/unittests/models/test_litellm.py
+++ b/tests/unittests/models/test_litellm.py
@@ -416,9 +416,26 @@ class MockLLMClient(LiteLLMClient):
     self.completion_mock = completion_mock
 
   async def acompletion(self, model, messages, tools, **kwargs):
-    return await self.acompletion_mock(
-        model=model, messages=messages, tools=tools, **kwargs
-    )
+    if kwargs.get("stream", False):
+      kwargs_copy = dict(kwargs)
+      kwargs_copy.pop("stream", None)
+
+      async def stream_generator():
+        stream_data = self.completion_mock(
+            model=model,
+            messages=messages,
+            tools=tools,
+            stream=True,
+            **kwargs_copy,
+        )
+        for item in stream_data:
+          yield item
+
+      return stream_generator()
+    else:
+      return await self.acompletion_mock(
+          model=model, messages=messages, tools=tools, **kwargs
+      )
 
   def completion(self, model, messages, tools, stream, **kwargs):
     return self.completion_mock(