fix: merge annotations from all streaming chunks in stream_chunk_builder

[bbarwik]

Relevant issues

Fixes stream_chunk_builder losing annotations from all but the first annotation-bearing streaming chunk.

Pre-Submission checklist

• I have Added testing in the tests/test_litellm/ directory, Adding at least 1 test is a hard requirement
• My PR passes all unit tests on make test-unit
• My PR's scope is as isolated as possible, it only solves 1 specific problem

Type

Bug Fix

Changes

Problem

In litellm/main.py, stream_chunk_builder() collects all streaming chunks that contain annotations, but then only uses annotations from the first chunk:

annotations = annotation_chunks[0]["choices"][0]["delta"]["annotations"]

This means if annotations arrive across multiple streaming chunks (e.g. Gemini sends grounding metadata in the final chunk while other providers may include citations in intermediate chunks), all but the first set of annotations are silently lost.

Fix

Merge annotations from ALL annotation-bearing chunks by extending a list:

all_annotations: list = []
for ac in annotation_chunks:
    all_annotations.extend(ac["choices"][0]["delta"]["annotations"])
response["choices"][0]["message"]["annotations"] = all_annotations

Tests added

tests/test_litellm/test_stream_chunk_builder_annotations.py:

• test_stream_chunk_builder_merges_annotations_from_multiple_chunks — verifies annotations from chunk 1 and chunk 3 are both present in the assembled response
• test_stream_chunk_builder_single_annotation_chunk_still_works — verifies existing single-chunk behavior (no regression)
• test_stream_chunk_builder_no_annotations — verifies no annotations attribute when none are present (no regression)

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
litellm	Ready	Preview, Comment	Mar 14, 2026 8:57pm

[CLAassistant]

All committers have signed the CLA.

[greptile-apps]

Greptile Summary

This PR fixes a bug in stream_chunk_builder where annotations from all but the first annotation-bearing streaming chunk were silently dropped. The fix replaces a single direct index access (annotation_chunks[0]) with a loop that extends a merged list across all annotation chunks, then assigns the full list to the assembled response message.

Key changes:

• litellm/main.py: Replaces annotations = annotation_chunks[0][...]["annotations"] with a loop that calls all_annotations.extend(...) for every annotation chunk, ensuring no annotations are lost regardless of which chunk they arrive in.
• tests/test_litellm/test_stream_chunk_builder_annotations.py: Adds three focused mock unit tests — multi-chunk merge correctness, single-chunk regression guard, and no-annotation regression guard — all using locally constructed ModelResponseStream objects with no network calls.

Confidence Score: 5/5

• This PR is safe to merge — it is a minimal, well-tested bug fix with no backwards-incompatible changes.
• The change is a straightforward one-to-many extension of an existing list-merge pattern already used elsewhere in the same function. The logic is correct, the three new tests cover the fixed case plus two regression scenarios, tests make no real network calls, and no existing behavior for callers without annotations is affected.
• No files require special attention.

Important Files Changed

Filename	Overview
litellm/main.py	Fixes stream_chunk_builder to merge annotations from all streaming chunks instead of only the first — minimal, targeted, and correct change.
tests/test_litellm/test_stream_chunk_builder_annotations.py	New test file with three pure unit tests (no network calls): multi-chunk merge, single-chunk regression, and no-annotation regression — all compliant with the mock-only rule for this directory.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[stream_chunk_builder called with chunks] --> B[Filter annotation_chunks\nchunks where delta.annotations is not None]
    B --> C{len annotation_chunks > 0?}
    C -- No --> D[Skip annotation assignment]
    C -- Yes --> E["all_annotations = []"]
    E --> F[For each ac in annotation_chunks]
    F --> G["all_annotations.extend(\n  ac[choices][0][delta][annotations]\n)"]
    G --> F
    F -- done --> H["response[choices][0][message][annotations]\n= all_annotations"]
    H --> I[Continue with audio/image chunks...]

Loading

_{Last reviewed commit: e38bbec}