Skip to content

Commit 8a92fd1

Browse files
xuanyang15copybara-github
authored andcommitted
fix: ignore empty function chunk in LiteLlm streaming response
Fixes #1532 PiperOrigin-RevId: 808636127
1 parent c37bd27 commit 8a92fd1

File tree

2 files changed

+108
-2
lines changed

2 files changed

+108
-2
lines changed

src/google/adk/models/lite_llm.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -437,10 +437,17 @@ def _model_response_to_chunk(
437437
for tool_call in message.get("tool_calls"):
438438
# aggregate tool_call
439439
if tool_call.type == "function":
440+
func_name = tool_call.function.name
441+
func_args = tool_call.function.arguments
442+
443+
# Ignore empty chunks that don't carry any information.
444+
if not func_name and not func_args:
445+
continue
446+
440447
yield FunctionChunk(
441448
id=tool_call.id,
442-
name=tool_call.function.name,
443-
args=tool_call.function.arguments,
449+
name=func_name,
450+
args=func_args,
444451
index=tool_call.index,
445452
), finish_reason
446453

tests/unittests/models/test_litellm.py

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -267,6 +267,77 @@
267267
]
268268

269269

270+
STREAM_WITH_EMPTY_CHUNK = [
271+
ModelResponse(
272+
choices=[
273+
StreamingChoices(
274+
finish_reason=None,
275+
delta=Delta(
276+
role="assistant",
277+
tool_calls=[
278+
ChatCompletionDeltaToolCall(
279+
type="function",
280+
id="call_abc",
281+
function=Function(
282+
name="test_function",
283+
arguments='{"test_arg":',
284+
),
285+
index=0,
286+
)
287+
],
288+
),
289+
)
290+
]
291+
),
292+
ModelResponse(
293+
choices=[
294+
StreamingChoices(
295+
finish_reason=None,
296+
delta=Delta(
297+
role="assistant",
298+
tool_calls=[
299+
ChatCompletionDeltaToolCall(
300+
type="function",
301+
id=None,
302+
function=Function(
303+
name=None,
304+
arguments=' "value"}',
305+
),
306+
index=0,
307+
)
308+
],
309+
),
310+
)
311+
]
312+
),
313+
# This is the problematic empty chunk that should be ignored.
314+
ModelResponse(
315+
choices=[
316+
StreamingChoices(
317+
finish_reason=None,
318+
delta=Delta(
319+
role="assistant",
320+
tool_calls=[
321+
ChatCompletionDeltaToolCall(
322+
type="function",
323+
id=None,
324+
function=Function(
325+
name=None,
326+
arguments="",
327+
),
328+
index=0,
329+
)
330+
],
331+
),
332+
)
333+
]
334+
),
335+
ModelResponse(
336+
choices=[StreamingChoices(finish_reason="tool_calls", delta=Delta())]
337+
),
338+
]
339+
340+
270341
@pytest.fixture
271342
def mock_response():
272343
return ModelResponse(
@@ -1591,6 +1662,34 @@ async def test_generate_content_async_non_compliant_multiple_function_calls(
15911662
assert final_response.content.parts[1].function_call.args == {"arg": "value2"}
15921663

15931664

1665+
@pytest.mark.asyncio
1666+
async def test_generate_content_async_stream_with_empty_chunk(
1667+
mock_completion, lite_llm_instance
1668+
):
1669+
"""Tests that empty tool call chunks in a stream are ignored."""
1670+
mock_completion.return_value = iter(STREAM_WITH_EMPTY_CHUNK)
1671+
1672+
responses = [
1673+
response
1674+
async for response in lite_llm_instance.generate_content_async(
1675+
LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
1676+
)
1677+
]
1678+
1679+
assert len(responses) == 1
1680+
final_response = responses[0]
1681+
assert final_response.content.role == "model"
1682+
1683+
# Crucially, assert that only ONE tool call was generated,
1684+
# proving the empty chunk was ignored.
1685+
assert len(final_response.content.parts) == 1
1686+
1687+
function_call = final_response.content.parts[0].function_call
1688+
assert function_call.name == "test_function"
1689+
assert function_call.id == "call_abc"
1690+
assert function_call.args == {"test_arg": "value"}
1691+
1692+
15941693
@pytest.mark.asyncio
15951694
def test_get_completion_inputs_generation_params():
15961695
# Test that generation_params are extracted and mapped correctly

0 commit comments

Comments
 (0)