From f07a549e137f0dc2b2e767f40704f9d80b5911fc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 21:28:28 -0700 Subject: [PATCH] fix(guardrails): replace blocked text completion streams like chat, keep Gray Swan's vendor payload, scan text streams in the content filter --- .../guardrail_translation/handler.py | 42 +++++- .../guardrail_hooks/grayswan/grayswan.py | 21 +-- .../litellm_content_filter/content_filter.py | 132 ++++++++++-------- .../test_text_completion_guardrail_handler.py | 24 ++++ .../content_filter/test_content_filter.py | 54 ++++++- .../guardrail_hooks/test_grayswan.py | 19 ++- .../proxy_logging/test_guardrail_pipeline.py | 39 ++++++ 7 files changed, 253 insertions(+), 78 deletions(-) diff --git a/litellm/llms/openai/completion/guardrail_translation/handler.py b/litellm/llms/openai/completion/guardrail_translation/handler.py index faa2776e3c3..5d86ddec014 100644 --- a/litellm/llms/openai/completion/guardrail_translation/handler.py +++ b/litellm/llms/openai/completion/guardrail_translation/handler.py @@ -16,13 +16,17 @@ from litellm.llms.base_llm.guardrail_translation.base_translation import ( StreamingScanKey, StreamTransformSink, ) -from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items -from litellm.types.utils import GenericGuardrailAPIInputs, TextChoices, TextCompletionResponse +from litellm.llms.base_llm.guardrail_translation.utils import ( + blocked_chat_stream_usage, + stream_item_field, + stream_item_items, +) +from litellm.types.utils import GenericGuardrailAPIInputs, TextChoices, TextCompletionResponse, Usage if TYPE_CHECKING: from fastapi import HTTPException - from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.integrations.custom_guardrail import CustomGuardrail, ModifyResponseException from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth @@ -258,6 +262,38 @@ class OpenAITextCompletionHandler(BaseTranslation): ) return (*responses_so_far, terminator) + def build_block_sse_chunks( + self, + exc: "ModifyResponseException", + stream_started: bool = False, + responses_so_far: Sequence[object] | None = None, + ) -> Sequence[bytes]: + streamed: Final = _text_completion_chunks(responses_so_far or ()) + source: Final = streamed[0] if streamed else TextCompletionResponse(model=exc.model) + prompt_tokens, completion_tokens = blocked_chat_stream_usage(exc.original_response) + message_chunk: Final = TextCompletionResponse( + id=source.id, + created=source.created, + model=source.model or exc.model, + object="text_completion", + choices=[TextChoices(index=0, text=exc.message)], + ) + final_chunk: Final = TextCompletionResponse( + id=source.id, + created=source.created, + model=source.model or exc.model, + object="text_completion", + choices=[TextChoices(index=0, text="", finish_reason="content_filter")], + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), + ) + return tuple( + f"data: {chunk.model_dump_json(exclude_none=True)}\n\n".encode() for chunk in (message_chunk, final_chunk) + ) + def build_stream_error_items( self, exc: "HTTPException", diff --git a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py index 456d9052a97..e57d6e78378 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py +++ b/litellm/proxy/guardrails/guardrail_hooks/grayswan/grayswan.py @@ -89,12 +89,13 @@ def _str_keyed(value: object) -> Mapping[str, object]: return _STR_KEYED_MAPPING.validate_python(value) if isinstance(value, Mapping) else {} -def _proxy_metadata(request_data: Mapping[str, object]) -> Mapping[str, object]: +def _request_route(request_data: Mapping[str, object]) -> str | None: litellm_metadata: Final = request_data.get("litellm_metadata") - if litellm_metadata is not None: - return _str_keyed(litellm_metadata) - metadata: Final = _str_keyed(request_data.get("metadata")) - return {key: value for key, value in metadata.items() if key.startswith("user_api_key_")} + metadata_bucket: Final = _str_keyed( + litellm_metadata if litellm_metadata is not None else request_data.get("metadata") + ) + route: Final = metadata_bucket.get("user_api_key_request_route") + return route if isinstance(route, str) else None class GraySwanGuardrailMissingSecrets(Exception): @@ -640,8 +641,8 @@ class GraySwanGuardrail(CustomGuardrail): from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.llms import load_guardrail_translation_mappings - request_route: Final = _proxy_metadata(request_data).get("user_api_key_request_route") - route_call_types: Final = get_call_types_for_route(request_route) if isinstance(request_route, str) else None + request_route: Final = _request_route(request_data) + route_call_types: Final = get_call_types_for_route(request_route) if request_route is not None else None call_type: Final = ( (route_call_types[0].value if route_call_types else None) or (logging_obj.call_type if logging_obj is not None else None) @@ -710,11 +711,11 @@ class GraySwanGuardrail(CustomGuardrail): inbound_headers: Final = self._extract_inbound_headers(request_data, logging_obj) - proxy_metadata: Final = _proxy_metadata(request_data) + litellm_metadata: Final = _str_keyed(request_data.get("litellm_metadata")) cleaned_litellm_metadata: Final = ( - {**proxy_metadata, "headers": {**_str_keyed(proxy_metadata.get("headers")), **inbound_headers}} + {**litellm_metadata, "headers": {**_str_keyed(litellm_metadata.get("headers")), **inbound_headers}} if inbound_headers - else proxy_metadata + else litellm_metadata ) if cleaned_litellm_metadata: sanitized: Final = _DECODED_JSON.validate_python( diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index f58c8705815..76edb2cf59c 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -47,6 +47,9 @@ from litellm.types.utils import ( Message, ModelResponse, ModelResponseStream, + StreamingChoices, + TextChoices, + TextCompletionResponse, ) if TYPE_CHECKING: @@ -139,6 +142,22 @@ class _CategoryConfigView(TypedDict): category_file: str | None +def _streamed_choice_content(choice: StreamingChoices | TextChoices) -> str | None: + if isinstance(choice, TextChoices): + return choice.text if isinstance(choice.text, str) else "" + delta: Final = getattr(choice, "delta", None) + if not delta: + return None + content: Final = getattr(delta, "content", None) + return content if isinstance(content, str) else "" + + +def _streamed_choices(item: object) -> Sequence[StreamingChoices | TextChoices]: + if isinstance(item, (ModelResponseStream, TextCompletionResponse)): + return item.choices + return () + + @dataclass(frozen=True, slots=True) class _StreamedChoiceState: buffered_text: str = "" @@ -2105,71 +2124,70 @@ class ContentFilterGuardrail(CustomGuardrail): try: async for item in response: - if isinstance(item, ModelResponseStream) and item.choices: - for choice in item.choices: - if not (hasattr(choice, "delta") and choice.delta): - continue + for choice in _streamed_choices(item): + new_content = _streamed_choice_content(choice) + if new_content is None: + continue - choice_index = getattr(choice, "index", 0) - if not isinstance(choice_index, int): - choice_index = 0 + choice_index = getattr(choice, "index", 0) + if not isinstance(choice_index, int): + choice_index = 0 - content = getattr(choice.delta, "content", None) - is_final = bool(getattr(choice, "finish_reason", None)) - new_content = content if isinstance(content, str) else "" - if not new_content and not is_final: - continue + is_final = bool(getattr(choice, "finish_reason", None)) + if not new_content and not is_final: + continue - previous_state = state_by_choice.get(choice_index, _StreamedChoiceState()) - buffered_text = previous_state.buffered_text + new_content - if not buffered_text: - continue + previous_state = state_by_choice.get(choice_index, _StreamedChoiceState()) + buffered_text = previous_state.buffered_text + new_content + if not buffered_text: + continue - # Add a space at the end if it's the final chunk to trigger word boundaries (\b) - text_to_scan = buffered_text + (" " if is_final else "") - choice_detections: list[ContentFilterDetection] = [] - scan_started = time.perf_counter() + # Add a space at the end if it's the final chunk to trigger word boundaries (\b) + text_to_scan = buffered_text + (" " if is_final else "") + choice_detections: list[ContentFilterDetection] = [] + scan_started = time.perf_counter() - try: - masked_text = self._filter_single_text(text_to_scan, detections=choice_detections) - if is_final and masked_text.endswith(" "): - masked_text = masked_text[:-1] - latest_detections = tuple(choice_detections) - except HTTPException: - state_by_choice[choice_index] = replace( - previous_state, latest_detections=tuple(choice_detections) - ) - raise - except Exception as e: - verbose_proxy_logger.error("ContentFilterGuardrail: Error in masking: %s", e) - masked_text = text_to_scan # Fallback to current text - latest_detections = previous_state.latest_detections - finally: - scan_seconds += time.perf_counter() - scan_started - - safe_to_yield_len = max( - previous_state.yielded_masked_text_len, - len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS), + try: + masked_text = self._filter_single_text(text_to_scan, detections=choice_detections) + if is_final and masked_text.endswith(" "): + masked_text = masked_text[:-1] + latest_detections = tuple(choice_detections) + except HTTPException: + state_by_choice[choice_index] = replace( + previous_state, latest_detections=tuple(choice_detections) ) - choice.delta.content = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len] - next_state = replace( - previous_state, - buffered_text=buffered_text, - yielded_masked_text_len=safe_to_yield_len, - latest_detections=latest_detections, - ) - if is_final: - state_by_choice[choice_index] = next_state - continue + raise + except Exception as e: + verbose_proxy_logger.error("ContentFilterGuardrail: Error in masking: %s", e) + masked_text = text_to_scan # Fallback to current text + latest_detections = previous_state.latest_detections + finally: + scan_seconds += time.perf_counter() - scan_started - trim_started = time.perf_counter() - state_by_choice[choice_index] = self._trim_streamed_choice_buffer(next_state, masked_text, plan) - scan_seconds += time.perf_counter() - trim_started + safe_to_yield_len = max( + previous_state.yielded_masked_text_len, + len(masked_text) - (0 if is_final else CONTENT_FILTER_STREAMING_HOLDBACK_CHARS), + ) + released_text = masked_text[previous_state.yielded_masked_text_len : safe_to_yield_len] + if isinstance(choice, TextChoices): + choice.text = released_text + else: + choice.delta.content = released_text + next_state = replace( + previous_state, + buffered_text=buffered_text, + yielded_masked_text_len=safe_to_yield_len, + latest_detections=latest_detections, + ) + if is_final: + state_by_choice[choice_index] = next_state + continue - yield item - else: - # Not a ModelResponseStream or no choices - yield as is - yield item + trim_started = time.perf_counter() + state_by_choice[choice_index] = self._trim_streamed_choice_buffer(next_state, masked_text, plan) + scan_seconds += time.perf_counter() - trim_started + + yield item except HTTPException: status = "guardrail_intervened" raise diff --git a/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py b/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py index 33e6a9887c2..cc213cf35f5 100644 --- a/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py +++ b/tests/unit/llms/openai/completion/test_text_completion_guardrail_handler.py @@ -2,6 +2,7 @@ Unit tests for OpenAI Text Completion Guardrail Translation Handler """ +import json from typing import List, Optional, Tuple from unittest.mock import MagicMock @@ -367,6 +368,13 @@ class _RecordingKeywordFilter(ContentFilterGuardrail): return await super().apply_guardrail(inputs, request_data, input_type, logging_obj) +class _PassthroughReplaceGuardrail(CustomGuardrail): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + if any("kumquat" in text for text in inputs.get("texts") or []): + self.raise_passthrough_exception(violation_message="replaced kumquat", request_data=request_data) + return inputs + + def _text_chunk(index: int, text: str | None, finish_reason: str | None = None) -> TextCompletionResponse: return TextCompletionResponse( id="cmpl-stream", @@ -486,3 +494,19 @@ class TestStreamingOutputProcessing: assert isinstance(error_frame, bytes) assert error_frame.startswith(b'data: {"error": ') assert b"kumquat" in error_frame + + @pytest.mark.asyncio + async def test_flat_stream_passthrough_block_appends_the_replacement_and_ends_with_content_filter(self): + guardrail = _PassthroughReplaceGuardrail( + guardrail_name="passthrough", event_hook=GuardrailEventHooks.post_call, default_on=True + ) + chunks = [_text_chunk(0, "A kum"), _text_chunk(0, "quat tart"), _text_chunk(0, "", "stop"), _usage_chunk()] + + delivered = await _drive_flat_stream(guardrail, chunks) + + frames = [json.loads(item.decode().removeprefix("data: ")) for item in delivered if isinstance(item, bytes)] + assert [(frame["choices"][0]["text"], frame["choices"][0].get("finish_reason")) for frame in frames] == [ + ("replaced kumquat", None), + ("", "content_filter"), + ] + assert frames[-1]["usage"]["total_tokens"] == 8 diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py index f7fb7e1fed0..169cd8c0651 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/content_filter/test_content_filter.py @@ -30,7 +30,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter impor ContentFilterCategoryConfig, ContentFilterDetection, ) -from litellm.types.utils import StandardLoggingGuardrailInformation +from litellm.types.utils import StandardLoggingGuardrailInformation, TextChoices, TextCompletionResponse class TestContentFilterGuardrail: @@ -3523,3 +3523,55 @@ async def test_image_whose_description_is_clean_passes_through_unchanged(): ) assert result == {"images": ["https://example.com/desk.png"], "texts": []} + + +def _fruit_word_filter() -> ContentFilterGuardrail: + return ContentFilterGuardrail( + guardrail_name="fruit-filter", + blocked_words=[ + BlockedWord(keyword="persimmon", action=ContentFilterAction.MASK), + BlockedWord(keyword="kumquat", action=ContentFilterAction.BLOCK), + ], + event_hook=GuardrailEventHooks.post_call, + ) + + +async def _text_completion_stream(*parts: tuple[str | None, str | None]): + for text, finish_reason in parts: + yield TextCompletionResponse( + id="cmpl-stream", + model="gpt-3.5-turbo-instruct", + object="text_completion", + choices=[TextChoices(index=0, text=text, finish_reason=finish_reason)], + ) + + +@pytest.mark.asyncio +async def test_streaming_hook_masks_a_word_split_across_text_completion_chunks(): + request_data: Final[dict] = {} + + delivered: Final = [ + chunk + async for chunk in _fruit_word_filter().async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), + response=_text_completion_stream(("A pers", None), ("immon tart", None), (None, "stop")), + request_data=request_data, + ) + ] + + assert "".join(chunk.choices[0].text or "" for chunk in delivered) == "A [KEYWORD_REDACTED] tart" + assert request_data["metadata"]["standard_logging_guardrail_information"][0]["guardrail_status"] == "success" + + +@pytest.mark.asyncio +async def test_streaming_hook_blocks_a_word_split_across_text_completion_chunks(): + with pytest.raises(HTTPException) as exc_info: + async for _ in _fruit_word_filter().async_post_call_streaming_iterator_hook( + user_api_key_dict=MagicMock(), + response=_text_completion_stream(("A kum", None), ("quat tart", None), (None, "stop")), + request_data={}, + ): + pass + + assert exc_info.value.status_code == 400 + assert "kumquat" in str(exc_info.value.detail) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py index f064f7fa5d5..a5ba9dab0bd 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_grayswan.py @@ -889,7 +889,7 @@ async def test_pre_call_payload_unchanged() -> None: @pytest.mark.asyncio -async def test_post_call_on_a_text_completion_forwards_only_the_key_and_team_metadata() -> None: +async def test_post_call_never_forwards_the_proxy_key_metadata_of_the_request_metadata() -> None: guardrail = _post_call_guardrail() client = _CapturingClient() guardrail.async_handler = client @@ -899,11 +899,16 @@ async def test_post_call_on_a_text_completion_forwards_only_the_key_and_team_met id="cmpl-1", model="gpt-4.1-mini", choices=[TextChoices(index=0, text="hello", finish_reason="stop")] ), guardrail_to_apply=guardrail, - user_api_key_dict=UserAPIKeyAuth(team_id="team-456", key_alias="ci-key"), - request_data={"model": "gpt-4.1-mini", "prompt": "say hello", "metadata": {"requester_note": "private"}}, + user_api_key_dict=UserAPIKeyAuth(team_id="team-456", key_alias="ci-key", user_email="dev@example.com"), + request_data={ + "model": "gpt-4.1-mini", + "prompt": "say hello", + "metadata": { + "user_api_key_hash": "hashed-key", + "user_api_key_user_email": "dev@example.com", + "user_api_key_request_route": "/v1/completions", + }, + }, ) - forwarded = client.calls[0]["json"]["litellm_metadata"] - assert forwarded["user_api_key_team_id"] == "team-456" - assert forwarded["user_api_key_key_alias"] == "ci-key" - assert "requester_note" not in forwarded + assert "litellm_metadata" not in client.calls[0]["json"] diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 52c3d71dd6a..cfd4fe197be 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -2730,6 +2730,45 @@ async def test_streaming_iterator_hook_pipeline_modify_response_emits_translated assert not any(item is chunk for item in delivered for chunk in chunks) +@pytest.mark.asyncio +async def test_streaming_iterator_hook_pipeline_modify_response_replaces_a_text_completion_stream( + proxy_logging, make_user_api_key_auth, monkeypatch +): + seen: Final[Dict[str, Any]] = {} + monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen, block=True)]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + pipeline: Final = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep( + guardrail="gr-post", + on_pass="allow", + on_fail="modify_response", + modify_response_message="content policy block", + ) + ], + ) + data: Final = {**_post_call_pipeline_data(stream=True), "prompt": "hi"} + data["metadata"]["_guardrail_pipelines"] = [("response-governance", pipeline)] + chunks: Final = [_text_completion_chunk("hello "), _text_completion_chunk("world"), _text_completion_chunk("", "stop")] + auth: Final = make_user_api_key_auth(request_route="/v1/completions") + + await proxy_logging.pre_call_hook(user_api_key_dict=auth, data=data, call_type="text_completion", guardrails_only=True) + delivered: Final = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(chunks), request_data=data + ) + ] + + frames: Final = [json.loads(item.decode().removeprefix("data: ")) for item in delivered] + assert seen["count"] == 1 + assert [(frame["choices"][0]["text"], frame["choices"][0].get("finish_reason")) for frame in frames] == [ + ("content policy block", None), + ("", "content_filter"), + ] + + @pytest.mark.asyncio async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthropic_sse( proxy_logging, make_user_api_key_auth, monkeypatch