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(