diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 42eb77f5f9..2d6bdb5298 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -221,6 +221,10 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): def event_stream_media_type() -> str: return _EVENT_STREAM_MEDIA_TYPE + @staticmethod + def event_stream_endpoint_is_de_anonymizable(endpoint: str) -> bool: + return _is_converse_endpoint(endpoint) + @staticmethod async def de_anonymize_event_stream( # noqa: PLR0915 body_bytes: bytes, diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index f9c93d5be0..db8d519d9b 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -300,11 +300,16 @@ class LlmPassthroughRouteHandler(BaseTranslation): return getattr(handler_cls, "de_anonymize_event_stream", None) @staticmethod - def supports_event_stream_de_anonymization(provider: Optional[str]) -> bool: - return ( - LlmPassthroughRouteHandler._resolve_event_stream_de_anonymizer(provider) - is not None + def supports_event_stream_de_anonymization( + provider: Optional[str], endpoint: Optional[str] + ) -> bool: + handler_cls = _get_provider_handlers().get(provider or "") + endpoint_check = getattr( + handler_cls, "event_stream_endpoint_is_de_anonymizable", None ) + if endpoint_check is None: + return False + return endpoint_check(endpoint or "") @staticmethod async def de_anonymize_event_stream( diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 68cf270b4e..db39a2d89c 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1340,7 +1340,7 @@ class ProxyBaseLLMRequestProcessing: if ( self._has_post_call_guardrails_for_passthrough() - and self._passthrough_provider_has_stream_guardrail_handler() + and self._passthrough_endpoint_has_stream_guardrail_handler() ): body_bytes = b"".join( [chunk async for chunk in generator] # type: ignore[union-attr] @@ -1781,20 +1781,23 @@ class ProxyBaseLLMRequestProcessing: return True return False - def _passthrough_provider_has_stream_guardrail_handler(self) -> bool: + def _passthrough_endpoint_has_stream_guardrail_handler(self) -> bool: """ - True when the resolved passthrough provider has an event-stream guardrail - handler that can rewrite buffered frames. Only such providers may have - their stream buffered for post-call guardrails; every other provider must - keep streaming so the response is not silently turned into a non-streaming - body when an unrelated post-call guardrail is registered. + True when the resolved passthrough provider AND endpoint have an + event-stream guardrail handler that can rewrite buffered frames. Only such + endpoints may have their stream buffered for post-call guardrails; every + other endpoint must keep streaming so the response is not silently turned + into a non-streaming body when no content modification would occur (e.g. + Bedrock invoke-with-response-stream, whose frames the Converse handler + leaves untouched). """ from litellm.llms.pass_through.guardrail_translation.handler import ( LlmPassthroughRouteHandler, ) return LlmPassthroughRouteHandler.supports_event_stream_de_anonymization( - self.data.get("custom_llm_provider") + self.data.get("custom_llm_provider"), + self.data.get("endpoint"), ) def _passthrough_event_stream_media_type(self) -> Optional[str]: diff --git a/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py b/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py index 1f6cafe5f3..f8bd83fc7d 100644 --- a/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/test_litellm/llms/pass_through/guardrail_translation/test_handler.py @@ -190,22 +190,35 @@ class TestDeAnonymizeEventStream: class TestSupportsEventStreamDeAnonymization: - def test_bedrock_is_supported(self): + def test_bedrock_converse_stream_is_supported(self): assert ( - LlmPassthroughRouteHandler.supports_event_stream_de_anonymization("bedrock") + LlmPassthroughRouteHandler.supports_event_stream_de_anonymization( + "bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream" + ) is True ) + def test_bedrock_invoke_stream_is_not_supported(self): + assert ( + LlmPassthroughRouteHandler.supports_event_stream_de_anonymization( + "bedrock", + "model/us.amazon.nova-lite-v1:0/invoke-with-response-stream", + ) + is False + ) + def test_unknown_provider_is_not_supported(self): assert ( LlmPassthroughRouteHandler.supports_event_stream_de_anonymization( - "anthropic" + "anthropic", "model/foo/converse-stream" ) is False ) def test_missing_provider_is_not_supported(self): assert ( - LlmPassthroughRouteHandler.supports_event_stream_de_anonymization(None) + LlmPassthroughRouteHandler.supports_event_stream_de_anonymization( + None, "model/foo/converse-stream" + ) is False ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index f7b67018e6..0b3a31d2de 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -2800,20 +2800,26 @@ class TestEventStreamAllmPassthroughRoute: class TestAllmPassthroughStreamingProviderGate: """ Regression: the streaming-buffer gate for allm_passthrough_route must only - fire for providers that have an event-stream guardrail handler (Bedrock). + fire for provider+endpoint pairs that have an event-stream guardrail handler + able to rewrite frames (Bedrock converse-stream). A non-Bedrock streaming passthrough response must keep streaming even when a post-call guardrail is registered globally, instead of being silently - buffered into a non-streaming Response. Bedrock must still be buffered so the - converse-stream de-anonymization handler can rewrite frames. + buffered into a non-streaming Response. A Bedrock endpoint the Converse + handler cannot rewrite (e.g. invoke-with-response-stream) must also keep + streaming. Only converse-stream is buffered so its frames can be + de-anonymized. """ - def _build_processing_obj(self, custom_llm_provider: str) -> ProxyBaseLLMRequestProcessing: + def _build_processing_obj( + self, custom_llm_provider: str, endpoint: str = "" + ) -> ProxyBaseLLMRequestProcessing: logging_obj = MagicMock() logging_obj.litellm_call_id = "call-123" logging_obj.cost_breakdown = None data = { "custom_llm_provider": custom_llm_provider, + "endpoint": endpoint, "litellm_logging_obj": logging_obj, } return ProxyBaseLLMRequestProcessing(data=data) @@ -2874,8 +2880,12 @@ class TestAllmPassthroughStreamingProviderGate: assert streamed == chunks @pytest.mark.asyncio - async def test_bedrock_stream_is_buffered_through_handler(self, monkeypatch): - processing_obj = self._build_processing_obj("bedrock") + async def test_bedrock_converse_stream_is_buffered_through_handler( + self, monkeypatch + ): + processing_obj = self._build_processing_obj( + "bedrock", "model/us.amazon.nova-lite-v1:0/converse-stream" + ) chunks = [b"raw-1", b"raw-2"] with patch.object( @@ -2898,3 +2908,30 @@ class TestAllmPassthroughStreamingProviderGate: assert result.body == b"modified-body" assert result.headers["content-type"] == "application/vnd.amazon.eventstream" mock_handler.assert_awaited_once() + + @pytest.mark.asyncio + async def test_bedrock_invoke_stream_is_not_buffered(self, monkeypatch): + processing_obj = self._build_processing_obj( + "bedrock", "model/us.amazon.nova-lite-v1:0/invoke-with-response-stream" + ) + chunks = [b"raw-1", b"raw-2"] + + with patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails", + return_value=False, + ), patch.object( + ProxyBaseLLMRequestProcessing, + "_has_post_call_guardrails_for_passthrough", + return_value=True, + ), patch( + "litellm.llms.bedrock.passthrough.guardrail_translation.handler." + "BedrockPassthroughGuardrailHandler.de_anonymize_event_stream", + new=AsyncMock(return_value=b"modified-body"), + ) as mock_handler: + result = await self._run(processing_obj, monkeypatch, chunks) + + assert isinstance(result, StreamingResponse) + streamed = [chunk async for chunk in result.body_iterator] + assert streamed == chunks + mock_handler.assert_not_awaited()