mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
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:
parent
27b6f4f2ba
commit
b55622e8ba
3 changed files with 47 additions and 45 deletions
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue