Skip to content

fix(models): aggregate streaming thought parts for all LiteLlm providers - #6898

Open
Ashfaqbs wants to merge 1 commit into
google:mainfrom
Ashfaqbs:fix/litellm-streaming-thought-part-aggregation
Open

fix(models): aggregate streaming thought parts for all LiteLlm providers#6898
Ashfaqbs wants to merge 1 commit into
google:mainfrom
Ashfaqbs:fix/litellm-streaming-thought-part-aggregation

Conversation

@Ashfaqbs

Copy link
Copy Markdown

Please ensure you have read the contribution guide before creating a pull request.

Link to Issue or Description of Change

Problem:
LiteLlm streaming builds the aggregated (non-partial) LlmResponse by joining the buffered text into a single part, but passes the buffered reasoning parts straight through unjoined. So for a provider that streams reasoning token-by-token (e.g. xai/grok-4.6, OpenAI reasoning models via LiteLLM), the final response — and therefore the persisted session event — holds one types.Part(thought=True) per streamed reasoning delta instead of one part per thinking block. _aggregate_streaming_thought_parts already does this joining (splitting on thought_signature so Anthropic's per-block boundaries are preserved), but it was only wired into the Anthropic message-building path (_content_to_message_param), not into the two stream finalizers that build the response object itself.

Solution:
Call _aggregate_streaming_thought_parts(reasoning_parts) from both _finalize_tool_call_response and _finalize_text_response in src/google/adk/models/lite_llm.py, instead of list(reasoning_parts). This makes every LiteLlm provider produce the same shape the non-streaming path already produces, and matches the aggregator's own docstring ("produces clean parts for session history and outbound requests").

Testing Plan

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.

Added test_generate_content_async_stream_aggregates_reasoning_deltas in tests/unittests/models/test_litellm.py, which streams three separate reasoning_content deltas (no thought_signature, matching how a non-Anthropic provider like xAI streams) followed by a text delta, and asserts the final non-partial response contains exactly one thought=True part with the joined text, instead of three.

$ pytest tests/unittests/models/test_litellm.py -q
404 passed, 1 warning in 5.27s

Also ran the full reasoning/thought-focused subset in isolation:

$ pytest tests/unittests/models/test_litellm.py -k "reasoning or thought" -q
44 passed, 360 deselected in 8.85s

Manual End-to-End (E2E) Tests:

Not run — this is a pure data-shape bug in response aggregation, fully exercised by the unit test above (constructs the same per-token ModelResponseStream deltas a live xAI/OpenAI-reasoning stream would produce and asserts on the resulting LlmResponse.content.parts). No live provider credentials were available to additionally verify against a real streaming API call.

Checklist

  • I have read the CONTRIBUTING.md document.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

Additional context

Both call sites previously read thought_parts=list(reasoning_parts) if reasoning_parts else None; the fix replaces list(reasoning_parts) with _aggregate_streaming_thought_parts(reasoning_parts). No behavior changes for the Anthropic path, which already routed through the aggregator via a separate code path (_content_to_message_param) for outbound requests — this PR fixes the inbound/session-persisted shape for all providers, including Anthropic's own aggregated response object.

The streaming finalizers passed the raw per-delta thought parts straight
into the response, so a provider that streams reasoning token by token
(e.g. xai/grok, OpenAI reasoning models via LiteLLM) produced one
types.Part(thought=True) per delta in the aggregated response instead of
one part per thinking block. _aggregate_streaming_thought_parts already
did this joining, splitting on thought_signature to preserve Anthropic's
per-block boundaries, but it was only wired into the Anthropic message-
building path. Call it from both stream finalizers so every provider
gets the same shape the non-streaming path already produces.

Fixes google#6895
@google-cla

google-cla Bot commented Aug 25, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

@Ashfaqbs

Copy link
Copy Markdown
Author

recheck

@benclarkeio

Copy link
Copy Markdown
Contributor

I submitted a PR for this already: #6896

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

LiteLlm streaming stores one thought part per reasoning delta

3 participants