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
This commit is contained in:
avidelatm
2025-06-23 10:30:12 -07:00
committed by Copybara-Service
parent ea69c9093a
commit bd67e8480f
2 changed files with 21 additions and 4 deletions
+1 -1
View File
@@ -679,7 +679,7 @@ class LiteLlm(BaseLlm):
aggregated_llm_response_with_tool_call = None aggregated_llm_response_with_tool_call = None
usage_metadata = None usage_metadata = None
fallback_index = 0 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): for chunk, finish_reason in _model_response_to_chunk(part):
if isinstance(chunk, FunctionChunk): if isinstance(chunk, FunctionChunk):
index = chunk.index or fallback_index index = chunk.index or fallback_index
+20 -3
View File
@@ -416,9 +416,26 @@ class MockLLMClient(LiteLLMClient):
self.completion_mock = completion_mock self.completion_mock = completion_mock
async def acompletion(self, model, messages, tools, **kwargs): async def acompletion(self, model, messages, tools, **kwargs):
return await self.acompletion_mock( if kwargs.get("stream", False):
model=model, messages=messages, tools=tools, **kwargs 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): def completion(self, model, messages, tools, stream, **kwargs):
return self.completion_mock( return self.completion_mock(