mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
ea69c9093a
commit
bd67e8480f
@@ -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
|
||||||
|
|||||||
@@ -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(
|
||||||
|
|||||||
Reference in New Issue
Block a user