diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/example_config.yaml b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/example_config.yaml index b5c9f0f8b69..4f732773a75 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/example_config.yaml +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/example_config.yaml @@ -1,11 +1,3 @@ -# Example LiteLLM Proxy configuration for LLM Shield Proxy -# LLM Shield Proxy is a self-hosted PII gateway: https://github.com/ninadphalak/LLM-Shield-Proxy -# -# Unlike a masking guardrail, LLM Shield Proxy's substitution is reversible. Personal data is -# replaced with placeholders before the request goes to the provider, and the original -# values are put back into the model's reply, so the end user still sees real data while -# the provider never received it. - model_list: - model_name: gpt-4o litellm_params: @@ -13,45 +5,10 @@ model_list: api_key: os.environ/OPENAI_API_KEY guardrails: - # Both modes belong on ONE entry. pre_call redacts the outbound request and post_call - # restores the reply; listing only pre_call would send placeholders back to the user. - guardrail_name: "llm_shield_proxy" litellm_params: guardrail: llm_shield_proxy mode: ["pre_call", "post_call"] default_on: true - # Your own LLM Shield Proxy deployment. Defaults to http://localhost:8000, and also reads - # LLM_SHIELD_PROXY_API_BASE from the environment. api_base: "http://localhost:8000" - # A virtual key configured on that deployment. Also reads LLM_SHIELD_PROXY_API_KEY. api_key: os.environ/LLM_SHIELD_PROXY_API_KEY - -# Usage: -# -# 1. Run LLM Shield Proxy somewhere the proxy can reach: -# pip install llm-shield-proxy -# llm-shield-proxy --port 8000 -# -# 2. Point this config at it and start the proxy: -# export LLM_SHIELD_PROXY_API_KEY="your-virtual-key" -# litellm --config example_config.yaml -# -# 3. Send a request containing personal data: -# curl http://localhost:4000/v1/chat/completions \ -# -H "Authorization: Bearer sk-1234" \ -# -H "Content-Type: application/json" \ -# -d '{"model":"gpt-4o","messages":[{"role":"user","content":"Email jane.doe@example.com the invoice"}]}' -# -# The provider receives a stand-in value in place of the address. The reply you get -# back carries the real address again. -# -# Notes: -# -# - Requests are refused if LLM Shield Proxy is unreachable or returns an error, rather than -# being forwarded. Sending them on would hand the provider exactly the data this -# guardrail exists to withhold. -# - Restoring a value requires the request and the reply to share a session. LiteLLM's -# session id is used when present; otherwise one is generated per request. -# - Streaming replies are restored as chunks arrive. A placeholder split across two -# chunks is held back until it is complete, so partial values are never emitted. -# - Only text is redacted; images and audio pass through untouched. diff --git a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py index d43e9249a45..dc0dbfda242 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py +++ b/litellm/proxy/guardrails/guardrail_hooks/llm_shield_proxy/llm_shield_proxy.py @@ -1013,13 +1013,24 @@ class LLMShieldProxyGuardrail(CustomGuardrail): request is therefore neither read from nor written to the cache. Inside the proxy this hook does not redact -- the proxy's pre-call hook already ran -- and caching is left alone, because the proxy restores after the cache write. + + A streamed request is refused once redacted. No hook restores an SDK stream, and the + stream's cache writer reads the request from before this hook, so it would also be + cached despite the bypass. """ before: Final = self._minted_session_id(kwargs) - # The parent rewrites `kwargs` in place and hands the same dict back. _ = await super().async_pre_call_deployment_hook(kwargs, call_type) session_id: Final = self._minted_session_id(kwargs) if session_id is None or session_id == before: return kwargs + if kwargs.get("stream") is True: + raise GuardrailRaisedException( + guardrail_name=self.guardrail_name, + message=( + "LLM Shield Proxy cannot restore a streamed reply for a model-level guardrail " + "outside the LiteLLM proxy; send the request through the proxy or without stream=True." + ), + ) metadata: Final = _as_object(kwargs.get("litellm_metadata")) if metadata is not None: metadata[_DEPLOYMENT_RESTORE_KEY] = session_id @@ -1251,7 +1262,6 @@ class LLMShieldProxyGuardrail(CustomGuardrail): # fix: a fragment is an arbitrary slice of a JSON document, so the code cannot tell # whether the position it writes is inside a string literal, and escaping # unconditionally would corrupt the values that are not. - @log_guardrail_information async def async_post_call_success_hook( self, data: MutableRequest, diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index 5edf055104f..ed34c07022e 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -1938,6 +1938,41 @@ class TestProxyWiring: assert reply.choices[0].message.content == "Repeat alice@example.com" assert cache.cache_dict == {}, "the redacted request's reply must not be cached" + @pytest.mark.asyncio + async def test_model_level_streaming_outside_the_proxy_is_refused(self, monkeypatch): + """No hook restores an SDK stream, and its cache writer misses the bypass, so it fails closed.""" + guardrail = _guardrail(event_hook=["pre_call", "post_call"], default_on=False) + _mock_post(guardrail, {"texts": ["Repeat [EMAIL_1]"]}) + cache = InMemoryCache() + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setattr(litellm, "cache", litellm.Cache(type="local")) + monkeypatch.setattr(litellm.cache, "cache", cache) + + with pytest.raises(GuardrailRaisedException, match="cannot restore a streamed reply"): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "Repeat alice@example.com"}], + mock_response="Repeat [EMAIL_1]", + stream=True, + guardrails=[GUARDRAIL_NAME], + ) + await asyncio.gather(*_PENDING_CACHE_WRITES) + + assert cache.cache_dict == {} + + @pytest.mark.asyncio + async def test_restored_values_are_not_recorded_as_guardrail_telemetry(self): + """Guardrail logging is exported to traces even with message logging off, so the + restored reply must not land in it.""" + guardrail, _ = _shielded({"[EMAIL_1]": "alice@example.com"}) + reply = ModelResponse(choices=[Choices(index=0, message=Message(role="assistant", content="[EMAIL_1]"))]) + data = {"messages": [], "metadata": {}} + + restored = await guardrail.async_post_call_success_hook(data=data, user_api_key_dict=None, response=reply) + + assert restored.choices[0].message.content == "alice@example.com" + assert "alice@example.com" not in json.dumps(data, default=str) + class TestStreamUsage: @pytest.mark.asyncio