From 2c522b3544ab0ad4acb5df3b8d02d3e280b1704b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 17:09:43 -0700 Subject: [PATCH] fix(guardrails): drop vLLM token ids from non-streamed choices a post-call hook rewrote A masking post_call guardrail rewrote the answer text but left the hosted_vllm token_ids that decode back to the unmasked words, in the response and in the logged payload. Compare each choice that carries token ids before and after the post-call hooks, and set its token_ids to null when its output changed. The logging overlay that swaps in the served text drops the logged ids the same way. --- litellm/litellm_core_utils/redact_messages.py | 13 +++ .../litellm_core_utils/served_output_texts.py | 16 ++- litellm/proxy/utils.py | 32 ++++++ .../test_served_output_texts.py | 34 +++++++ .../test_post_call_success_hook.py | 99 +++++++++++++++++++ 5 files changed, 191 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 981859c620b..df50f635905 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -227,6 +227,19 @@ def _redact_provider_specific_fields(message: object) -> None: setattr(message, "provider_specific_fields", _redacted_provider_specific_fields(fields)) # noqa: B010 # typed as object +def _drop_provider_specific_token_ids(holder: object) -> None: + fields: Final = _as_mapping(getattr(holder, "provider_specific_fields", None)) + if fields is not None and fields.get("token_ids") is not None: + setattr(holder, "provider_specific_fields", {**fields, "token_ids": None}) # noqa: B010 # typed as object + + +def drop_token_ids(choice: object) -> None: + if getattr(choice, "token_ids", None) is not None: + setattr(choice, "token_ids", None) # noqa: B010 # pydantic extra, not a declared field + for holder in (choice, getattr(choice, "delta", None), getattr(choice, "message", None)): + _drop_provider_specific_token_ids(holder) + + def _redact_provider_specific_fields_dict(choice: Mapping[str, object]) -> None: for message in (choice.get("message"), choice.get("delta")): if ( diff --git a/litellm/litellm_core_utils/served_output_texts.py b/litellm/litellm_core_utils/served_output_texts.py index b19973e06c3..4fd7d5d6592 100644 --- a/litellm/litellm_core_utils/served_output_texts.py +++ b/litellm/litellm_core_utils/served_output_texts.py @@ -3,7 +3,7 @@ post-call guardrail rewrote rather than the provider response the proxy assemble from __future__ import annotations -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from itertools import chain from typing import Final, Literal @@ -191,13 +191,23 @@ def overlay_served_output_texts( } +def _without_token_ids(holder: Mapping[str, object]) -> Mapping[str, object]: + fields: Final = _as_json_object(holder.get("provider_specific_fields")) + dropped: Final = {**holder, "token_ids": None} if holder.get("token_ids") is not None else holder + if fields is None or fields.get("token_ids") is None: + return dropped + return {**dropped, "provider_specific_fields": {**fields, "token_ids": None}} + + def _choice_with_text(choice: object, text: str | None) -> object: choice_obj: Final = _as_json_object(choice) if choice_obj is None or text is None: return choice message: Final = _as_json_object(choice_obj.get("message")) if message is not None: - return choice if message.get("content") == text else {**choice_obj, "message": {**message, "content": text}} + if message.get("content") == text: + return choice + return _without_token_ids({**choice_obj, "message": _without_token_ids({**message, "content": text})}) if isinstance(choice_obj.get("text"), str): - return {**choice_obj, "text": text} + return choice if choice_obj.get("text") == text else _without_token_ids({**choice_obj, "text": text}) return choice diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0b0698faa9a..041e6011aed 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -164,6 +164,7 @@ from litellm.litellm_core_utils.duration_parser import get_budget_window_start from litellm.litellm_core_utils.env_utils import get_env_float, get_env_int from litellm.litellm_core_utils.litellm_logging import BilledRoundsDispatch, Logging, deployment_pricing_model_info from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.litellm_core_utils.redact_messages import drop_token_ids from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.served_output_texts import ( @@ -172,6 +173,7 @@ from litellm.litellm_core_utils.served_output_texts import ( ) from litellm.litellm_core_utils.token_counter import offload_token_count, usage_tokenizer_model from litellm.llms import load_guardrail_translation_mappings +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import ( AlertType, @@ -702,6 +704,34 @@ def _pipeline_step_streams_live(step: "PipelineStep") -> bool: return not unified_guardrail.resolve_streaming_flag(callback, "streaming_buffer_until_moderated", True) +def _choice_output(choice: object) -> tuple[object, object]: + message: Final = getattr(choice, "message", None) + return (message.model_dump() if isinstance(message, BaseModel) else message, getattr(choice, "text", None)) + + +def _carries_token_ids(choice: object) -> bool: + return getattr(choice, "token_ids", None) is not None or any( + stream_item_field(stream_item_field(holder, "provider_specific_fields"), "token_ids") is not None + for holder in (choice, getattr(choice, "message", None)) + ) + + +def _token_id_choice_outputs(response: object) -> Mapping[int, tuple[object, object]]: + return MappingProxyType( + { + index: _choice_output(choice) + for index, choice in enumerate(stream_item_items(response, "choices")) + if _carries_token_ids(choice) + } + ) + + +def _drop_token_ids_of_rewritten_choices(response: object, outputs_before: Mapping[int, tuple[object, object]]) -> None: + for index, choice in enumerate(stream_item_items(response, "choices")): + if index in outputs_before and _choice_output(choice) != outputs_before[index]: + drop_token_ids(choice) + + def _pipelines_stream_live(pipelines: Sequence[tuple[str, "GuardrailPipeline"]]) -> bool: return all(_pipeline_step_streams_live(step) for step in _pipeline_steps(pipelines)) @@ -4142,6 +4172,7 @@ class ProxyLogging: from litellm.proxy.proxy_server import llm_router from litellm.types.guardrails import GuardrailEventHooks + token_id_outputs: Final = _token_id_choice_outputs(response) pipeline_response: Final = await self._run_post_call_pipelines( data=data, user_api_key_dict=user_api_key_dict, @@ -4230,6 +4261,7 @@ class ProxyLogging: response = callback_response except Exception as e: raise e + _drop_token_ids_of_rewritten_choices(response, token_id_outputs) return response async def _run_parallel_post_call_guardrails( diff --git a/tests/unit/litellm_core_utils/test_served_output_texts.py b/tests/unit/litellm_core_utils/test_served_output_texts.py index e01c710dfdb..94df1055032 100644 --- a/tests/unit/litellm_core_utils/test_served_output_texts.py +++ b/tests/unit/litellm_core_utils/test_served_output_texts.py @@ -157,6 +157,40 @@ def test_overlay_replaces_logged_choice_text_with_served_text(): assert _choice_texts(logged) == (RAW, RAW) +RAW_TOKEN_IDS: Final = [7, 8, 9] + + +def _choice_token_ids(choice: dict[str, object]) -> tuple[object, object, object]: + message: Final = choice.get("message") + message_fields: Final = message.get("provider_specific_fields") if isinstance(message, dict) else None + choice_fields: Final = choice.get("provider_specific_fields") + return ( + choice.get("token_ids"), + choice_fields.get("token_ids") if isinstance(choice_fields, dict) else None, + message_fields.get("token_ids") if isinstance(message_fields, dict) else None, + ) + + +@pytest.mark.parametrize( + "logged_choice", + ( + {"index": 0, "provider_specific_fields": {"token_ids": RAW_TOKEN_IDS}, "message": {"content": RAW}}, + {"index": 0, "message": {"content": RAW, "provider_specific_fields": {"token_ids": RAW_TOKEN_IDS}}}, + {"index": 0, "text": RAW, "token_ids": RAW_TOKEN_IDS}, + ), + ids=("chat-choice-fields", "chat-message-fields", "text-completion"), +) +@pytest.mark.parametrize(("served", "keeps_token_ids"), ((MASKED, False), (RAW, True)), ids=("rewritten", "unchanged")) +def test_overlay_drops_logged_token_ids_only_when_it_rewrites_the_choice_text( + logged_choice: dict[str, object], served: str, keeps_token_ids: bool +) -> None: + overlaid: Final = overlay_served_output_texts({"id": "x", "choices": [logged_choice]}, (served,)) + assert isinstance(overlaid, dict) + choices: Final = overlaid["choices"] + assert isinstance(choices, list) + assert (RAW_TOKEN_IDS in _choice_token_ids(choices[0])) is keeps_token_ids + + def test_overlay_leaves_unreadable_inputs_untouched(): logged = _chat_dict(RAW) assert overlay_served_output_texts(logged, None) is logged diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py index 53d8948869f..ad9dac84481 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py @@ -120,3 +120,102 @@ async def test_post_call_block_names_the_blocking_guardrail_in_applied_guardrail data=data, response=MagicMock(), user_api_key_dict=make_user_api_key_auth() ) assert data["metadata"]["applied_guardrails"] == ["passer", "blocker"] + + +def _vllm_chat_response(content: str) -> litellm.ModelResponse: + from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import ( + convert_to_model_response_object, + ) + + return convert_to_model_response_object( + response_object={ + "id": "chatcmpl-vllm", + "object": "chat.completion", + "created": 1, + "model": "Qwen/Qwen3-0.6B", + "choices": [ + { + "index": 0, + "finish_reason": "stop", + "token_ids": [785, 6722, 315, 9625, 374, 12095, 13], + "message": {"role": "assistant", "content": content}, + } + ], + "usage": {"prompt_tokens": 5, "completion_tokens": 7, "total_tokens": 12}, + }, + model_response_object=litellm.ModelResponse(), + ) + + +def _vllm_text_completion_response(text: str) -> litellm.TextCompletionResponse: + return litellm.TextCompletionResponse( + id="cmpl-vllm", + object="text_completion", + created=1, + model="Qwen/Qwen3-0.6B", + choices=[{"index": 0, "text": text, "finish_reason": "length", "logprobs": None, "token_ids": [12095, 13]}], + usage={"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}, + ) + + +def _chat_token_ids(response: litellm.ModelResponse) -> object: + return response.choices[0].provider_specific_fields["token_ids"] + + +def _text_token_ids(response: litellm.TextCompletionResponse) -> object: + return response.choices[0].token_ids + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "data", "token_ids_of", "kept"), + [ + ( + _vllm_chat_response("The capital of France is Paris."), + {"model": "m", "messages": [{"role": "user", "content": "capital of France?"}]}, + _chat_token_ids, + None, + ), + ( + _vllm_chat_response("The capital of France is Lyon."), + {"model": "m", "messages": [{"role": "user", "content": "capital of France?"}]}, + _chat_token_ids, + [785, 6722, 315, 9625, 374, 12095, 13], + ), + ( + _vllm_text_completion_response(" Paris."), + {"model": "m", "prompt": "The capital of France is"}, + _text_token_ids, + None, + ), + ( + _vllm_text_completion_response(" Lyon."), + {"model": "m", "prompt": "The capital of France is"}, + _text_token_ids, + [12095, 13], + ), + ], + ids=["chat-masked", "chat-untouched", "completion-masked", "completion-untouched"], +) +async def test_post_call_success_hook_drops_token_ids_only_from_choices_a_guardrail_rewrote( + proxy_logging, make_user_api_key_auth, monkeypatch, response, data, token_ids_of, kept +): + from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( + ContentFilterGuardrail, + ) + from litellm.types.guardrails import BlockedWord, ContentFilterAction + + mask_paris = ContentFilterGuardrail( + guardrail_name="mask-paris", + blocked_words=[BlockedWord(keyword="Paris", action=ContentFilterAction.MASK)], + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + monkeypatch.setattr(litellm, "callbacks", [mask_paris]) + + served = await proxy_logging.post_call_success_hook( + data=data, response=response, user_api_key_dict=make_user_api_key_auth() + ) + + assert served is response + assert token_ids_of(served) == kept