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.
This commit is contained in:
Ninad Phalak 2026-10-05 09:58:34 +00:00
parent 27b6f4f2ba
commit b55622e8ba
No known key found for this signature in database
3 changed files with 47 additions and 45 deletions

View file

@ -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.

View file

@ -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,

View file

@ -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