google / google/adk-python

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

Đang mở
#7,004 3 bình luận 0 reaction 1 người được giao Được @surajksharma07 nhận Xem trên GitHub
models
Ngôn ngữ chính
Python
Star
21.5k
Fork
4k
Merge trung bình
1 ngày 14 giờ
Pull request đã merge (30 ngày)
37

Mô tả

## 🔴 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.

Hướng dẫn đóng góp

Mở hướng dẫn đóng góp

Đánh giá

Issue này chưa được đánh giá.

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.