From b55622e8ba550fbc8e67bb381797b7d348246396 Mon Sep 17 00:00:00 2001 From: Ninad Phalak Date: Mon, 5 Oct 2026 09:58:34 +0000 Subject: [PATCH] fix(guardrails): keep restored llm_shield_proxy values out of telemetry, refuse SDK streams - The post-call restore hook no longer goes through log_guardrail_information, which recorded its whole return value, the restored reply, as guardrail_response. That field is exported to traces even with message logging turned off. - A model-level stream outside the proxy is refused once redacted. Nothing restores an SDK stream, and its cache writer reads the request from before the deployment hook, so it also got cached despite the no-store bypass. - Drop a narrating comment, and keep example_config.yaml to config only; the how-to lives in the docs PR. --- .../llm_shield_proxy/example_config.yaml | 43 ------------------- .../llm_shield_proxy/llm_shield_proxy.py | 14 +++++- .../guardrail_hooks/test_llm_shield_proxy.py | 35 +++++++++++++++ 3 files changed, 47 insertions(+), 45 deletions(-) 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