From 7ed9be55b1a1adc357c3894c3a833487ecb9ef9f Mon Sep 17 00:00:00 2001 From: bbarwik Date: Sat, 14 Mar 2026 20:43:00 +0000 Subject: [PATCH 1/2] fix: merge annotations from all streaming chunks in stream_chunk_builder Previously, stream_chunk_builder only took annotations from the first chunk that contained them, losing any annotations from later chunks. This is a problem because providers like Gemini/Vertex AI send grounding metadata (converted to annotations) in the final streaming chunk, while other providers may spread annotations across multiple chunks. Changes: - Collect and merge annotations from ALL annotation-bearing chunks instead of only using the first one --- litellm/main.py | 11 +- .../test_stream_chunk_builder_annotations.py | 191 ++++++++++++++++++ 2 files changed, 200 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/test_stream_chunk_builder_annotations.py diff --git a/litellm/main.py b/litellm/main.py index 781a940ca7..81319bc432 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7528,8 +7528,15 @@ def stream_chunk_builder( # noqa: PLR0915 ] if len(annotation_chunks) > 0: - annotations = annotation_chunks[0]["choices"][0]["delta"]["annotations"] - response["choices"][0]["message"]["annotations"] = annotations + # Merge annotations from ALL chunks — providers may spread + # them across multiple streaming chunks or send them only in + # the final chunk. + all_annotations: list = [] + for ac in annotation_chunks: + all_annotations.extend( + ac["choices"][0]["delta"]["annotations"] + ) + response["choices"][0]["message"]["annotations"] = all_annotations audio_chunks = [ chunk diff --git a/tests/test_litellm/test_stream_chunk_builder_annotations.py b/tests/test_litellm/test_stream_chunk_builder_annotations.py new file mode 100644 index 0000000000..9c7ad4126b --- /dev/null +++ b/tests/test_litellm/test_stream_chunk_builder_annotations.py @@ -0,0 +1,191 @@ +""" +Tests for stream_chunk_builder annotation merging. + +Previously, stream_chunk_builder only took annotations from the FIRST +annotation chunk, losing any annotations that arrived in later chunks. +This fix merges annotations from ALL chunks. +""" + +from litellm import stream_chunk_builder +from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + +def test_stream_chunk_builder_merges_annotations_from_multiple_chunks(): + """ + stream_chunk_builder must merge annotations from ALL streaming chunks, + not just take them from the first annotation chunk. + + Providers may spread annotations across multiple chunks (e.g. Gemini + sends grounding metadata in the final chunk, while intermediate chunks + may carry different annotations). + """ + annotation_a = { + "type": "url_citation", + "url_citation": { + "url": "https://example.com/a", + "title": "Source A", + "start_index": 0, + "end_index": 10, + }, + } + annotation_b = { + "type": "url_citation", + "url_citation": { + "url": "https://example.com/b", + "title": "Source B", + "start_index": 20, + "end_index": 30, + }, + } + + chunks = [ + ModelResponseStream( + id="chatcmpl-test", + created=1700000000, + model="test-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta( + content="Part one. ", + role="assistant", + annotations=[annotation_a], + ), + ) + ], + ), + ModelResponseStream( + id="chatcmpl-test", + created=1700000000, + model="test-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="Part two."), + ) + ], + ), + ModelResponseStream( + id="chatcmpl-test", + created=1700000000, + model="test-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta( + content=None, + annotations=[annotation_b], + ), + ) + ], + ), + ] + + response = stream_chunk_builder(chunks=chunks) + assert response is not None + + message = response["choices"][0]["message"] + assert message.annotations is not None + assert len(message.annotations) == 2 + assert message.annotations[0] == annotation_a + assert message.annotations[1] == annotation_b + + +def test_stream_chunk_builder_single_annotation_chunk_still_works(): + """ + When annotations come from a single chunk (most common case), + stream_chunk_builder must still work correctly (no regression). + """ + annotation = { + "type": "url_citation", + "url_citation": { + "url": "https://example.com/only", + "title": "Only Source", + "start_index": 0, + "end_index": 5, + }, + } + + chunks = [ + ModelResponseStream( + id="chatcmpl-test", + created=1700000000, + model="test-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="Hello", role="assistant"), + ) + ], + ), + ModelResponseStream( + id="chatcmpl-test", + created=1700000000, + model="test-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=None, annotations=[annotation]), + ) + ], + ), + ] + + response = stream_chunk_builder(chunks=chunks) + assert response is not None + + message = response["choices"][0]["message"] + assert message.annotations is not None + assert len(message.annotations) == 1 + assert message.annotations[0] == annotation + + +def test_stream_chunk_builder_no_annotations(): + """ + When no chunks contain annotations, the message should not have + an annotations key (no regression). + """ + chunks = [ + ModelResponseStream( + id="chatcmpl-test", + created=1700000000, + model="test-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="Hello", role="assistant"), + ) + ], + ), + ModelResponseStream( + id="chatcmpl-test", + created=1700000000, + model="test-model", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=None), + ) + ], + ), + ] + + response = stream_chunk_builder(chunks=chunks) + assert response is not None + + message = response["choices"][0]["message"] + assert not hasattr(message, "annotations") or message.annotations is None From dd1ea3d39ef8d2584a31454e35e6130e7a7d4b79 Mon Sep 17 00:00:00 2001 From: brtydse100 <92057527+brtydse100@users.noreply.github.com> Date: Sun, 15 Mar 2026 10:29:59 +0200 Subject: [PATCH 2/2] Support multiple headers mapped to the customer user role (#23664) * added the header mapping feature * added tests * final cleanup * final cleanup * added missing test and logic * fixed header sending bug * Update litellm/proxy/auth/auth_utils.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * added back init file in responses + fixed test_auth_utils.py int local_testing --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/auth/auth_utils.py | 27 +++-- tests/local_testing/test_auth_utils.py | 2 +- .../proxy/auth/test_auth_utils.py | 106 ++++++++++++++++++ 3 files changed, 127 insertions(+), 8 deletions(-) diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 9a24041faa..0d3c627446 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -662,11 +662,12 @@ def _has_user_setup_sso(): return sso_setup -def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[str]: +def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[list]: """Return the header_name mapped to CUSTOMER role, if any (dict-based).""" if not user_id_mapping: return None items = user_id_mapping if isinstance(user_id_mapping, list) else [user_id_mapping] + customer_headers_mappings = [] for item in items: if not isinstance(item, dict): continue @@ -675,7 +676,11 @@ def get_customer_user_header_from_mapping(user_id_mapping) -> Optional[str]: if role is None or not header_name: continue if str(role).lower() == str(LitellmUserRoles.CUSTOMER).lower(): - return header_name + customer_headers_mappings.append(header_name.lower()) + + if customer_headers_mappings: + return customer_headers_mappings + return None @@ -724,7 +729,7 @@ def get_end_user_id_from_request_body( # User query: "system not respecting user_header_name property" # This implies the key in general_settings is 'user_header_name'. if request_headers is not None: - custom_header_name_to_check: Optional[str] = None + custom_header_name_to_check: Optional[Union[list, str]] = None # Prefer user mappings (new behavior) user_id_mapping = general_settings.get("user_header_mappings", None) @@ -741,13 +746,21 @@ def get_end_user_id_from_request_body( custom_header_name_to_check = value # If we have a header name to check, try to read it from request headers - if isinstance(custom_header_name_to_check, str): + if isinstance(custom_header_name_to_check, list): + headers_lower = {k.lower(): v for k, v in request_headers.items()} + for expected_header in custom_header_name_to_check: + header_value = headers_lower.get(expected_header) + if header_value is not None: + user_id_str = str(header_value) + if user_id_str.strip(): + return user_id_str + + elif isinstance(custom_header_name_to_check, str): for header_name, header_value in request_headers.items(): if header_name.lower() == custom_header_name_to_check.lower(): - user_id_from_header = header_value user_id_str = ( - str(user_id_from_header) - if user_id_from_header is not None + str(header_value) + if header_value is not None else "" ) if user_id_str.strip(): diff --git a/tests/local_testing/test_auth_utils.py b/tests/local_testing/test_auth_utils.py index d36f96b1a3..bffcb40baf 100644 --- a/tests/local_testing/test_auth_utils.py +++ b/tests/local_testing/test_auth_utils.py @@ -268,7 +268,7 @@ def test_get_customer_user_header_from_mapping_returns_customer_header(): {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}, ] result = get_customer_user_header_from_mapping(mappings) - assert result == "X-OpenWebUI-User-Email" + assert result == ["x-openwebui-user-email"] def test_get_customer_user_header_from_mapping_no_customer_returns_none(): diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index 82920ce1d8..5e42b110aa 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -209,3 +209,109 @@ def test_get_model_from_request_supports_google_model_names_with_slashes(): def test_get_model_from_request_vertex_passthrough_still_works(): route = "/vertex_ai/v1/projects/p/locations/l/publishers/google/models/gemini-1.5-pro:generateContent" assert get_model_from_request(request_data={}, route=route) == "gemini-1.5-pro" + + +def test_get_customer_user_header_returns_none_when_no_customer_role(): + from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping + + mappings = [ + {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"} + ] + result = get_customer_user_header_from_mapping(mappings) + assert result is None + + +def test_get_customer_user_header_returns_none_for_single_non_customer_mapping(): + from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping + + mapping = {"header_name": "X-Only-Internal", "litellm_user_role": "internal_user"} + result = get_customer_user_header_from_mapping(mapping) + assert result is None + +def test_get_customer_user_header_from_mapping_returns_customer_header(): + from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping + + mappings = [ + {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}, + {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}, + ] + result = get_customer_user_header_from_mapping(mappings) + assert result == ["x-openwebui-user-email"] + + +def test_get_customer_user_header_returns_customers_header_in_config_order_when_multiple_exist(): + from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping + + mappings = [ + {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}, + {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}, + {"header_name": "X-User-Id", "litellm_user_role": "customer"}, + ] + result = get_customer_user_header_from_mapping(mappings) + assert result == ['x-openwebui-user-email', 'x-user-id'] + + +def test_get_end_user_id_returns_id_from_user_header_mappings(): + from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body + + mappings = [ + {"header_name": "x-openwebui-user-id", "litellm_user_role": "internal_user"}, + {"header_name": "x-openwebui-user-email", "litellm_user_role": "customer"}, + ] + general_settings = {"user_header_mappings": mappings} + headers = {"x-openwebui-user-email": "1234"} + + with patch("litellm.proxy.auth.auth_utils._get_customer_id_from_standard_headers", return_value=None), \ + patch("litellm.proxy.proxy_server.general_settings", general_settings): + result = get_end_user_id_from_request_body(request_body={}, request_headers=headers) + + assert result == "1234" + + +def test_get_end_user_id_returns_first_customer_header_when_multiple_mappings_exist(): + from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body + + mappings = [ + {"header_name": "x-openwebui-user-id", "litellm_user_role": "internal_user"}, + {"header_name": "x-user-id", "litellm_user_role": "customer"}, + {"header_name": "x-openwebui-user-email", "litellm_user_role": "customer"}, + ] + general_settings = {"user_header_mappings": mappings} + headers = { + "x-user-id": "user-456", + "x-openwebui-user-email": "user@example.com", + } + + with patch("litellm.proxy.auth.auth_utils._get_customer_id_from_standard_headers", return_value=None), \ + patch("litellm.proxy.proxy_server.general_settings", general_settings): + result = get_end_user_id_from_request_body(request_body={}, request_headers=headers) + + assert result == "user-456" + + +def test_get_end_user_id_returns_none_when_no_customer_role_in_mappings(): + from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body + + mappings = [ + {"header_name": "x-openwebui-user-id", "litellm_user_role": "internal_user"}, + ] + general_settings = {"user_header_mappings": mappings} + headers = {"x-openwebui-user-id": "user-789"} + + with patch("litellm.proxy.auth.auth_utils._get_customer_id_from_standard_headers", return_value=None), \ + patch("litellm.proxy.proxy_server.general_settings", general_settings): + result = get_end_user_id_from_request_body(request_body={}, request_headers=headers) + + assert result is None + +def test_get_end_user_id_falls_back_to_deprecated_user_header_name(): + from litellm.proxy.auth.auth_utils import get_end_user_id_from_request_body + + general_settings = {"user_header_name": "x-custom-user-id"} + headers = {"x-custom-user-id": "user-legacy"} + + with patch("litellm.proxy.auth.auth_utils._get_customer_id_from_standard_headers", return_value=None), \ + patch("litellm.proxy.proxy_server.general_settings", general_settings): + result = get_end_user_id_from_request_body(request_body={}, request_headers=headers) + + assert result == "user-legacy"