google / google/adk-python

LiteLlm streaming keeps only the last tool call when one chunk carries N calls plus a finish_reason (latent)

Offen
#7,004 3 Kommentare 0 Reaktionen 1 zugewiesene Person Beansprucht von @surajksharma07 Auf GitHub ansehen
models
Vorherrschende Sprache
Python
Sterne
21.5k
Forks
4k
Ø Merge
1 T. 14 Std.
Gemergte PRs (30 T.)
37

Beschreibung

## 🔴 Required Information

**Describe the Bug:**

This issue is a latent robustness defect that is present in
`LiteLlm`'s streaming aggregation, not a live data-loss bug with a stock provider. Through
`litellm.acompletion(stream=True)` — what the default `LiteLLMClient` uses — litellm's
`CustomStreamWrapper` normalizes every provider's stream so that `finish_reason` only ever
arrives on a trailing empty-delta chunk, and that shape is handled correctly. The defect
fires only when a chunk carries complete tool calls *and* a finish reason in the same
delta, which today reaches `generate_content_async` only via the public `LiteLlm.llm_client`
field with a custom client, or if litellm's normalization ever changes. I'm filing it
because the code path is real, silent when hit, contradicts `BaseLlm`'s documented contract,
is exercised by ADK's own unit tests with raw chunks, and has a small fix that passes the
existing suite. Details follow.

In `LiteLlm.generate_content_async(..., stream=True)`, the "finalize the tool-call response"
check runs *inside* the per-chunk inner loop, and every finalization *replaces* the
aggregated response after clearing the buffers (`main` d637d1b4 lines 3399–3410; installed
2.8.0 lines 3352–3363):

```python
for chunk, finish_reason in _model_response_to_chunk(part):
...
if function_calls and (
finish_reason == "tool_calls"
or finish_reason == "length"
or (finish_reason == "stop" and chunk is None)
):
aggregated_llm_response_with_tool_call = _finalize_tool_call_response(...)
_reset_stream_buffers()
```

`_model_response_to_chunk` (head 2349–2369) yields one `FunctionChunk` **per tool call** in a
delta, each paired with the same choice-level `finish_reason`. So when one streamed chunk
carries two complete tool calls *and* `finish_reason="tool_calls"`:

1. `FunctionChunk(call_1)` + `"tool_calls"` → `function_calls={0: f1}` → finalize `{f1}` →
buffers reset.
2. `FunctionChunk(call_2)` + `"tool_calls"` → `function_calls={1: f2}` → the condition matches
again → finalize `{f2}`, **overwriting** the `{f1}` response.

`_finalize_tool_call_response` rebuilds from `function_calls` only, and the end-of-stream
fallback (`if function_calls and not aggregated_llm_response_with_tool_call`) does not fire
because the buffers are empty and the aggregate is already set. Only `f2` is yielded, with a
well-formed `STOP` response and no error or log. Three calls → only `f3`. The `"length"` arm
behaves the same way. The `"stop"` arm is protected by `chunk is None`; the `"tool_calls"`
and `"length"` arms are not. Any text or `reasoning_content` carried in that same chunk is
folded into the first, overwritten response and lost too.

**Steps to Reproduce:**

1. `pip install google-adk litellm` (reproduced on google-adk 2.8.0 with litellm 1.99.0, and
on `main` at `d637d1b4`).
2. Run the first script below. It feeds real `litellm.types.utils.ModelResponseStream` chunks
to `LiteLlm` through a custom `llm_client`, with tool declarations `f1`/`f2`/`f3`.
3. Compare the surviving function calls in the final non-partial `LlmResponse` across the
stream shapes.
4. Run the second script to see why the shape does not arrive through
`litellm.acompletion`: the same two-call response, passed through litellm's real
`CustomStreamWrapper`, is normalized to `[calls | finish=None]` + `[empty | finish]` and
ADK keeps both calls.

**Expected Behavior:**

Every tool call in the delta is present in the final non-partial response, for every shape.
`BaseLlm.generate_content_async`'s docstring (base_llm.py 153–190) says the final
`partial=False` chunk is identical to the `stream=False` output, and the non-streaming branch
returns both calls for the same input.

**Observed Behavior:**

```text
A_two_calls_one_chunk_finish_tool_calls -> survivors=['f2'] BUG
A3_three_calls_one_chunk_finish_tool_calls -> survivors=['f3'] BUG
L_two_calls_one_chunk_finish_length -> survivors=['f2'] BUG
B_separate_chunks_then_empty_finish -> survivors=['f1', 'f2'] OK
C_two_calls_one_chunk_finish_stop -> survivors=['f1', 'f2'] OK
D_two_calls_one_chunk_then_empty_finish -> survivors=['f1', 'f2'] OK
```

Instrumenting `_message_to_generate_content_response` in case A shows it invoked twice — once
with `['call_1']`, then with `['call_2']` — versus once with `['call_1', 'call_2']` in the
control cases.

**Environment Details:**

- ADK Library Version (pip show google-adk): 2.8.0 (also `main` @ d637d1b4)
- Desktop OS: Windows 11
- Python Version (python -V): 3.12.10

**Model Information:**

- Are you using LiteLLM: Yes (1.99.0)
- Which model is being used: a stub `LiteLLMClient` replaying `ModelResponseStream` chunks
for `openai/gpt-4o`; the chunk shape is what matters, not the model.

---

## 🟡 Optional Information

**Regression:**

No. The check has sat inside the inner per-chunk loop since the initial public commit
`982782014` (2025-04-08), when the aggregator tracked a single `function_id`, and was
inherited unchanged by `05f48347` (PR #759, index-keyed `function_calls` dict), `e8019b1b`
(#4225, the `chunk is None` guard on the `"stop"` arm only), `4c6096baa` (#4482, the
`"length"` arm), `36fd2c8e`, and `eaed0aa8` (2026-08-31, `last_finish_reason` tracking).
`eaed0aa8` is not in 2.8.0 (the CHANGELOG dates 2.8.0 to 2026-08-25) — head's end-of-stream
fallback reports the real finish reason where 2.8.0 hard-codes `tool_calls` — but the
placement of the check is identical in both. No commit in that history mentions more than one
tool call per chunk; the placement is an inherited artefact of the single-call era, not a
design choice.

**Logs:**

N/A — nothing is logged; the second finalization silently replaces the first.

**Screenshots / Video:**

N/A

**Additional Context:**

**Why this is latent today.** `litellm.acompletion(stream=True)` always returns a
`CustomStreamWrapper`, and ADK's default `LiteLLMClient.acompletion` (lite_llm.py 866–895)
uses it. `return_processed_chunk_logic` pops `finish_reason` from every non-empty chunk
(`litellm_core_utils/streaming_handler.py` ~1048 in 1.99.0, present in 1.84.0 at ~976 —
the comment there says it exists "for mistral etc. which return a value in their last chunk")
and re-emits it via `received_finish_reason` on a trailing empty-delta chunk (~1091–1130);
the custom-provider branch strips it explicitly (~1203, "so it appears only on the trailing
empty-delta chunk (OpenAI spec)"). This applies to fake-streamed and natively-streamed
providers alike — verified end-to-end for openai, azure, azure_ai, bedrock, vertex_ai,
gemini, ollama_chat, anthropic, openrouter, groq, together_ai, hosted_vllm, custom providers
and cached-response replay: ADK receives `[tool_calls chunk, finish=None]` then
`[empty chunk, finish="tool_calls"|"length"]` and keeps every call. So on litellm 1.84–1.99
(ADK's supported range) no stock provider path delivers the failing shape. It reaches
`generate_content_async` only through a user-supplied `llm_client` — a public `LiteLlm`
field — or a future change to litellm's normalization. Streaming is also opt-in
(`RunConfig.streaming_mode` defaults to `NONE`; `/run_sse` defaults `streaming=false`;
adk web's token-streaming toggle defaults off).

**Why it is still worth fixing.** (1) It violates the documented `BaseLlm` contract that
the final `partial=False` chunk equals the `stream=False` output. (2) It is silent — a
well-formed `STOP` response with the wrong number of calls. (3) ADK's own unit tests drive
this code with raw `ModelResponseStream` chunks through a stub client (exactly the vulnerable
path), so a one-fixture regression test would pin it; no existing fixture has more than one
tool call per delta or a tool-call delta sharing a chunk with `finish_reason`
`tool_calls`/`length`, which is why `MULTIPLE_FUNCTION_CALLS_STREAM` passes. (4) The fix is a
few lines and keeps the existing suite green (see below). (5) When it does fire, only the
surviving call executes; the persisted history is self-consistent but lossy, and a model
that insists on both calls re-runs the survivor each turn until `max_llm_calls` (default 500)
raises `LlmCallsLimitExceededError`.

**Prior art, all covering the *separate-chunk* shape (finish reason on its own empty chunk),
not this same-chunk shape:** #484 / #1038, fixed by PR #759, which created this aggregation
loop and the index-keyed dict; #187 / #153 (PR #172, message conversion); #4225 (the `stop`
guard); and #4482 (closed 2026-03-10), where a tool call was dropped entirely because
`"length"` was missing from the yield condition — its fix *added* the `"length"` arm to the
same `if` this report concerns. That fixed "nothing is yielded"; this is "only the last of N
is yielded".

**Related but separate, not claimed here:** text arriving *after* a mid-stream *text*
finalization (`"length"`, or an empty `"stop"` delta) is also single-slot-overwritten, but an
existing test (`test_streaming_text_buffer_is_reset_between_aggregated_responses`, from
`36fd2c8e`) pins last-segment-wins there, and moving the tool-call check does not change it.

**Suggested fix:** move the finalize decision to once per `part`, after the inner loop —
record the part's `finish_reason` and whether the finishing chunk was `None` inside the loop,
then finalize once. Prototyped: cases A/A3/L (and the text/reasoning carried in the same
chunk) flip to all-calls-preserved, with `tests/unittests/models/test_litellm.py` at 417/417
and the two other litellm test files at 11/11. "Merge instead of replace" is *not* the right
shape: the `"length"` arm calls `_parse_tool_call_arguments` and can return an error
`LlmResponse`, so merging would need error/normal reconciliation. Regression fixtures to add:
two `ChatCompletionDeltaToolCall` in one delta with `finish_reason="tool_calls"`; the same
with `"length"`; and the same with `content`/`reasoning_content` in the chunk.

**Minimal Reproduction Code:**

Script 1 — the mechanism, via a custom `llm_client` replaying raw chunks:

```python
"""LiteLlm streaming keeps only the last tool call when one chunk carries N complete calls + a finish reason."""
import asyncio
from google.adk.models.lite_llm import LiteLlm, LiteLLMClient
from google.adk.models.llm_request import LlmRequest
from google.genai import types
from litellm.types.utils import (
ChatCompletionDeltaToolCall, Delta, Function, ModelResponseStream, StreamingChoices,
)

def tc(id_, name, args, index):
return ChatCompletionDeltaToolCall(type="function", id=id_, function=Function(name=name, arguments=args), index=index)

def chunk(tool_calls, finish_reason):
return ModelResponseStream(model="openai/gpt-4o", choices=[
StreamingChoices(finish_reason=finish_reason, delta=Delta(role="assistant", tool_calls=tool_calls or None))])

F1 = ("call_1", "f1", '{"a": 1}', 0)
F2 = ("call_2", "f2", '{"b": 2}', 1)
F3 = ("call_3", "f3", '{"c": 3}', 2)
CASES = {
"A_two_calls_one_chunk_finish_tool_calls": [chunk([tc(*F1), tc(*F2)], "tool_calls")],
"A3_three_calls_one_chunk_finish_tool_calls": [chunk([tc(*F1), tc(*F2), tc(*F3)], "tool_calls")],
"L_two_calls_one_chunk_finish_length": [chunk([tc(*F1), tc(*F2)], "length")],
"B_separate_chunks_then_empty_finish": [chunk([tc(*F1)], None), chunk([tc(*F2)], None), chunk([], "tool_calls")],
"C_two_calls_one_chunk_finish_stop": [chunk([tc(*F1), tc(*F2)], "stop")],
"D_two_calls_one_chunk_then_empty_finish": [chunk([tc(*F1), tc(*F2)], None), chunk([], "tool_calls")],
}

class FakeClient(LiteLLMClient):
def __init__(self, chunks):
self._chunks = chunks
async def acompletion(self, model, messages, tools, **kwargs):
async def gen():
for c in self._chunks:
yield c
return gen()
def completion(self, *a, **k):
raise NotImplementedError

REQ = LlmRequest(
contents=[types.Content(role="user", parts=[types.Part.from_text(text="go")])],
config=types.GenerateContentConfig(tools=[types.Tool(function_declarations=[
types.FunctionDeclaration(name=n, description=n, parameters=types.Schema(
type=types.Type.OBJECT, properties={k: types.Schema(type=types.Type.INTEGER)}))
for n, k in (("f1", "a"), ("f2", "b"), ("f3", "c"))])]),
)

async def run(name, chunks):
expected = ["f1", "f2", "f3"] if "three" in name else ["f1", "f2"]
llm = LiteLlm(model="openai/gpt-4o", llm_client=FakeClient(chunks))
finals = [r async for r in llm.generate_content_async(REQ, stream=True) if not r.partial]
survivors = [p.function_call.name for r in finals for p in (r.content.parts if r.content else []) if p.function_call]
print(f"{name:44s} -> survivors={survivors} {'OK' if sorted(survivors) == expected else 'BUG'}")

async def main():
for name, chunks in CASES.items():
await run(name, chunks)

asyncio.run(main())
```

Script 2 — why the shape does not arrive through litellm.acompletion (real CustomStreamWrapper)

```python
"""The same two-call response through litellm's real streaming wrapper: finish_reason is moved to a trailing empty chunk and ADK keeps both calls."""
import asyncio, litellm
from litellm.types.utils import ModelResponse, Choices, Message, ChatCompletionMessageToolCall, Function
from google.adk.models.lite_llm import LiteLlm, LiteLLMClient
from google.adk.models.llm_request import LlmRequest
from google.genai import types

full = ModelResponse(model="openai/gpt-4o", choices=[Choices(finish_reason="tool_calls", index=0, message=Message(
role="assistant", content=None, tool_calls=[
ChatCompletionMessageToolCall(id="call_1", type="function", function=Function(name="f1", arguments='{"a": 1}')),
ChatCompletionMessageToolCall(id="call_2", type="function", function=Function(name="f2", arguments='{"b": 2}')),
]))])

async def main():
wrapper = await litellm.acompletion(model="openai/gpt-4o", messages=[{"role": "user", "content": "go"}],
stream=True, mock_response=full)
print("litellm.acompletion(stream=True) returned:", type(wrapper).__name__)
chunks = []
async for c in wrapper:
ch = c.choices[0]; tcs = ch.delta.tool_calls or []
print(f" wrapper chunk: finish_reason={ch.finish_reason!r:13} tool_calls={[t.function.name for t in tcs]}")
chunks.append(c)

class Replay(LiteLLMClient):
async def acompletion(self, model, messages, tools, **kw):
async def gen():
for c in chunks: yield c
return gen()
def completion(self, *a, **k): raise NotImplementedError

req = LlmRequest(contents=[types.Content(role="user", parts=[types.Part.from_text(text="go")])],
config=types.GenerateContentConfig(tools=[types.Tool(function_declarations=[
types.FunctionDeclaration(name=n, description=n, parameters=types.Schema(type=types.Type.OBJECT,
properties={k: types.Schema(type=types.Type.INTEGER)})) for n, k in (("f1","a"),("f2","b"))])]))
llm = LiteLlm(model="openai/gpt-4o", llm_client=Replay())
finals = [r async for r in llm.generate_content_async(req, stream=True) if not r.partial]
surv = [p.function_call.name for r in finals for p in (r.content.parts if r.content else []) if p.function_call]
print("ADK survivors via the REAL litellm stream path:", surv, "->", "both kept (bug NOT reachable this way)" if sorted(surv)==["f1","f2"] else "DROPPED")
asyncio.run(main())
```

Output on google-adk 2.8.0 / litellm 1.99.0:

```text
litellm.acompletion(stream=True) returned: CustomStreamWrapper
wrapper chunk: finish_reason=None tool_calls=['f1', 'f2']
wrapper chunk: finish_reason='tool_calls' tool_calls=[]
ADK survivors via the REAL litellm stream path: ['f1', 'f2'] -> both kept (bug NOT reachable this way)
```

**How often has this issue occurred?:**

- Always (100%) for the one-chunk shape, when delivered via a custom `llm_client`; never
observed through `litellm.acompletion` on litellm 1.84–1.99.

Beitragsleitfaden

Beitragsleitfaden öffnen

Bewertung

Dieses Issue wurde noch nicht bewertet.

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.