Skip to content

fix(litellm): stream partial function-call arguments - #6632

Open
chelsealong wants to merge 1 commit into
google:mainfrom
chelsealong:fix-litellm-partial-function-call-streaming
Open

fix(litellm): stream partial function-call arguments#6632
chelsealong wants to merge 1 commit into
google:mainfrom
chelsealong:fix-litellm-partial-function-call-streaming

Conversation

@chelsealong

Copy link
Copy Markdown

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

Link to Issue or Description of Change

1. Link to an existing issue (if applicable):

This PR addresses the LiteLLM portion of #6630. The issue asks for partial
function-call argument streaming across several adapters (LiteLLM, OpenAI
Chat Completions, OpenAI Responses/Azure Responses, Anthropic, Apigee). This
PR covers the LiteLLM adapter only, to keep the change small and reviewable;
the other adapters are natural follow-ups.

Problem:
When stream=True, LiteLlm.generate_content_async buffers tool-call
argument deltas internally (function_calls[index]["args_parts"]) and only
emits a FunctionCall once the model has finished generating the whole
call. Callers see token-by-token text updates but no signal at all while a
large tool call is being streamed, which can look like several seconds of
inactivity for big schemas.

Solution:
Emit an LlmResponse(partial=True) for each raw argument fragment as it
arrives from the provider, carrying the fragment via
FunctionCall.partial_args (a list of PartialArg(string_value=...)) with
will_continue=True, and the function-call id/name accumulated so far.
The existing aggregation logic is unchanged: once the tool call finishes,
the same final aggregated FunctionCall with parsed args is emitted as
before.

This mirrors the pattern already used by the Interactions adapter
(src/google/adk/models/interactions_utils.py::_handle_arguments_delta),
so partial function-call events follow an established shape in the
codebase rather than inventing a new one.

Partial updates are informational only and never trigger tool execution:
base_llm_flow.py already skips function-call execution whenever the
model-response event has partial=True.

Testing Plan

Unit Tests:

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

Added test_streaming_tool_call_emits_partial_argument_deltas in
tests/unittests/models/test_litellm.py, which streams a tool call across
two argument fragments and asserts:

  • one partial LlmResponse is emitted per fragment, each with
    partial=True, will_continue=True, args=None, and
    partial_args=[PartialArg(string_value=<fragment>)],
  • the final response still carries the normal, fully-parsed
    FunctionCall.args with partial_args=None.

I verified this test fails without the fix (reverted only the source file
with git stash, keeping the new test): assert len(responses) == len(fragments) + 1 fails with assert 1 == 3 since no partial events are emitted.

This change also increases the response count for several existing
streaming tests that assert exact counts (e.g.
test_streaming_tool_call_truncated_by_max_tokens,
test_streaming_tool_call_args_assembled_from_many_fragments,
test_streaming_buffers_hold_fragments_instead_of_growing_copies); I
updated those to account for the new partial events.

pytest results (local run, Python 3.12, minimal venv with project
extras installed via pip):

$ python -m pytest tests/unittests/models/test_litellm.py tests/unittests/models/test_lite_llm_gemma_tool_role.py tests/unittests/models/test_gemma_llm.py tests/unittests/models/test_litellm_import.py -q
398 passed in 10.39s

$ python -m pytest tests/unittests/models -q
1007 passed, 34 warnings in 18.30s

$ python -m pytest tests/unittests/flows -q --ignore=tests/unittests/flows/llm_flows/test_audio_transcriber.py
523 passed, 1 skipped, 77 warnings in 5.74s

(test_audio_transcriber.py was excluded because it fails to import in
this environment due to a missing unrelated optional dependency
(google.cloud.speech), independent of this change.)

Also ran formatting/type checks on the touched files:

$ isort --settings-path pyproject.toml --check-only --diff src/google/adk/models/lite_llm.py tests/unittests/models/test_litellm.py
(no output — clean)

$ pyink --diff src/google/adk/models/lite_llm.py tests/unittests/models/test_litellm.py
All done! (2 files would be left unchanged.)

$ mypy src/google/adk/models/lite_llm.py --ignore-missing-imports
Success: no issues found in 1 source file

Manual End-to-End (E2E) Tests:

Not performed — this is a streaming-shape change validated via unit tests
against a mocked LiteLLM client; there's no live LiteLLM-backed agent
available in this environment to drive manually.

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

This PR was prepared with AI assistance (Claude Code), with all changes
reviewed and verified by running the test suite, formatter, and type
checker locally before submission.

LiteLlm buffered tool-call argument deltas internally and only emitted
a FunctionCall once the model finished generating it, so callers saw
no progress while a large tool call was being streamed.

Emit an LlmResponse(partial=True) for each argument fragment as it
arrives, carrying the raw delta via FunctionCall.partial_args with
will_continue=True, mirroring the pattern already used by the
interactions adapter. These are informational only: base_llm_flow
already skips function execution for partial events.

Addresses the LiteLLM adapter portion of google#6630.
@adk-bot adk-bot added the models [Component] This issue is related to model support label Aug 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

models [Component] This issue is related to model support

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants