scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

feat(guardrails): add peyeeye PII redaction & rehydration guardrail

Browse source 
Adds a built-in `peyeeye` guardrail that:

- Pre-call: redacts PII from messages[].content via the peyeeye API. The
  redaction session id is cached against `litellm_call_id` so the
  post-call hook can rehydrate without a second client call.
- Post-call: swaps placeholders in `response.choices[].message.content`
  back to the original values, then cleans up the session.

Supports both stateful peyeeye sessions (`ses_…`) and stateless
AEAD-sealed rehydration keys (`skey_…`) for zero-retention setups.

- New: `litellm/proxy/guardrails/guardrail_hooks/peyeeye/{__init__,peyeeye}.py`
- New: `litellm/types/proxy/guardrails/guardrail_hooks/peyeeye.py`
  (UI config schema)
- Added: `SupportedGuardrailIntegrations.PEYEEYE` enum + a
  `PeyeeyeGuardrailConfigModel` mixin on `LitellmParams`
- Tests: `tests/test_litellm/proxy/guardrails/guardrail_hooks/test_peyeeye.py`

Auto-discovery in `guardrail_registry.get_guardrail_initializer_from_hooks`
picks up the new subdirectory; no edits needed to the explicit registry
dict.

Signed-off-by: Tim <info@airskip.app>
This commit is contained in:
Tim 2026-04-25 21:41:23 -05:00
parent fa81017e12
commit a636014380
5 changed files with 672 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

36
litellm/proxy/guardrails/guardrail_hooks/peyeeye/__init__.py Normal file
View file

@ -0,0 +1,36 @@
from typing import TYPE_CHECKING
from litellm.types.guardrails import SupportedGuardrailIntegrations
from .peyeeye import PeyeeyeGuardrail
if TYPE_CHECKING:
from litellm.types.guardrails import Guardrail, LitellmParams
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
import litellm
_peyeeye_callback = PeyeeyeGuardrail(
guardrail_name=guardrail.get("guardrail_name", ""),
api_key=litellm_params.api_key,
api_base=litellm_params.api_base,
peyeeye_locale=litellm_params.peyeeye_locale,
peyeeye_entities=litellm_params.peyeeye_entities,
peyeeye_session_mode=litellm_params.peyeeye_session_mode,
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
)
litellm.logging_callback_manager.add_litellm_callback(_peyeeye_callback)
return _peyeeye_callback
guardrail_initializer_registry = {
SupportedGuardrailIntegrations.PEYEEYE.value: initialize_guardrail,
}
guardrail_class_registry = {
SupportedGuardrailIntegrations.PEYEEYE.value: PeyeeyeGuardrail,
}

358
litellm/proxy/guardrails/guardrail_hooks/peyeeye/peyeeye.py Normal file
View file

@ -0,0 +1,358 @@
# +-------------------------------------------------------------+
#
# Use Peyeeye PII redaction & rehydration for LLM calls
# https://peyeeye.ai
#
# +-------------------------------------------------------------+
import os
from typing import (
TYPE_CHECKING,
Any,
Dict,
List,
Literal,
Optional,
Type,
Union,
)
try:
import httpx
HTTPX_AVAILABLE = True
except ImportError:
httpx = None # type: ignore
HTTPX_AVAILABLE = False
from fastapi import HTTPException
import litellm
from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
if TYPE_CHECKING:
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
DEFAULT_API_BASE = "https://api.peyeeye.ai"
SESSION_CACHE_TTL_SECONDS = 3600
class PeyeeyeGuardrailMissingSecrets(Exception):
"""Raised when the peyeeye API key is missing."""
class PeyeeyeGuardrailAPIError(Exception):
"""Raised when the peyeeye API returns an error."""
class PeyeeyeGuardrail(CustomGuardrail):
"""Peyeeye PII redaction + rehydration guardrail.
Pre-call hook redacts PII from each message's ``content`` and stores the
redaction session id under the request's ``litellm_call_id`` so the
post-call hook can rehydrate the model's response with the original
values.
Two session modes:
* ``stateful`` (default): peyeeye stores the token→value mapping under
a ``ses_…`` id; rehydrate references the id.
* ``stateless``: peyeeye returns a sealed ``skey_…`` blob; nothing is
retained server-side.
"""
def __init__(
self,
peyeeye_api_key: Optional[str] = None,
api_key: Optional[str] = None,
api_base: Optional[str] = None,
peyeeye_locale: Optional[str] = None,
peyeeye_entities: Optional[List[str]] = None,
peyeeye_session_mode: Optional[Literal["stateful", "stateless"]] = None,
**kwargs,
):
self.async_handler = get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback
)
self.peyeeye_api_key = (
peyeeye_api_key or api_key or os.environ.get("PEYEEYE_API_KEY")
)
if self.peyeeye_api_key is None:
raise PeyeeyeGuardrailMissingSecrets(
"Couldn't get peyeeye api key, either set the `PEYEEYE_API_KEY` "
"environment variable or pass it as `api_key` in the guardrail config."
)
self.api_base = (
api_base or os.environ.get("PEYEEYE_API_BASE") or DEFAULT_API_BASE
).rstrip("/")
self.peyeeye_locale = peyeeye_locale or "auto"
self.peyeeye_entities = peyeeye_entities
self.peyeeye_session_mode: Literal["stateful", "stateless"] = (
peyeeye_session_mode or "stateful"
)
verbose_proxy_logger.debug(
"Peyeeye guardrail initialized: name=%s, mode=%s, session_mode=%s",
kwargs.get("guardrail_name", "unknown"),
kwargs.get("event_hook", "unknown"),
self.peyeeye_session_mode,
)
super().__init__(**kwargs)
# ------------------------------------------------------------------ hooks
@log_guardrail_information
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict,
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
"anthropic_messages",
],
) -> Union[Exception, str, dict, None]:
"""Redact every text message before it reaches the model."""
event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
return data
messages: List[Dict[str, Any]] = data.get("messages") or []
text_parts = list(_iter_message_text(messages))
if not text_parts:
return data
redacted_texts, session_id = await self._redact_batch(
[t for _, _, t in text_parts]
)
for (msg_idx, part_path, _), redacted in zip(text_parts, redacted_texts):
_set_message_text(messages[msg_idx], part_path, redacted)
if session_id:
cache_key = self._cache_key(data)
try:
cache.set_cache(
cache_key, session_id, ttl=SESSION_CACHE_TTL_SECONDS
)
except Exception as e:
verbose_proxy_logger.warning(
"peyeeye: failed to cache session id: %s", e
)
return data
@log_guardrail_information
async def async_post_call_success_hook(
self,
data: dict,
user_api_key_dict: UserAPIKeyAuth,
response,
):
"""Rehydrate model output by swapping placeholders back to original PII."""
event_type: GuardrailEventHooks = GuardrailEventHooks.post_call
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
return response
cache_key = self._cache_key(data)
try:
session_id = litellm.cache.get_cache(cache_key) if litellm.cache else None
except Exception:
session_id = None
if not session_id:
# No redaction happened (or session expired); nothing to do.
return response
if isinstance(response, litellm.ModelResponse):
for choice in response.choices:
message = getattr(choice, "message", None)
if message is None:
continue
content = getattr(message, "content", None)
if isinstance(content, str) and content:
new = await self._rehydrate(content, session_id)
message.content = new
elif isinstance(content, list):
new_parts: List[Any] = []
for part in content:
if isinstance(part, dict) and part.get("type") == "text":
text = part.get("text", "")
new_parts.append(
{**part, "text": await self._rehydrate(text, session_id)}
)
else:
new_parts.append(part)
message.content = new_parts
# Clean up: drop the session server-side and from cache.
if self.peyeeye_session_mode == "stateful" and session_id.startswith("ses_"):
try:
await self._delete_session(session_id)
except Exception as e:
verbose_proxy_logger.debug(
"peyeeye: best-effort session cleanup failed: %s", e
)
try:
if litellm.cache:
litellm.cache.delete_cache(cache_key)
except Exception:
pass
return response
# --------------------------------------------------------------- internals
@staticmethod
def _cache_key(data: dict) -> str:
return f"peyeeye_session:{data.get('litellm_call_id') or id(data)}"
async def _redact_batch(self, texts: List[str]) -> tuple[List[str], Optional[str]]:
"""Redact a batch of texts in a single peyeeye session.
Returns the redacted strings (in order) and the session id (or sealed
``skey_…`` blob) so the post-call hook can rehydrate.
"""
body: Dict[str, Any] = {
"text": texts,
"locale": self.peyeeye_locale,
}
if self.peyeeye_entities:
body["entities"] = list(self.peyeeye_entities)
if self.peyeeye_session_mode == "stateless":
body["session"] = "stateless"
payload = await self._post("/v1/redact", body)
out_text = payload.get("text")
if isinstance(out_text, str):
redacted = [out_text]
elif isinstance(out_text, list):
redacted = [str(x) for x in out_text]
else:
redacted = list(texts) # fallback
if self.peyeeye_session_mode == "stateless":
session_id = payload.get("rehydration_key")
else:
session_id = payload.get("session_id") or payload.get("session")
return redacted, session_id
async def _rehydrate(self, text: str, session_id: str) -> str:
if not text:
return text
body = {"text": text, "session": session_id}
try:
payload = await self._post("/v1/rehydrate", body)
except Exception as e:
verbose_proxy_logger.warning("peyeeye: rehydrate failed: %s", e)
return text
return payload.get("text", text)
async def _delete_session(self, session_id: str) -> None:
url = f"{self.api_base}/v1/sessions/{session_id}"
await self.async_handler.delete(url=url, headers=self._headers(), timeout=10.0)
async def _post(self, path: str, body: Dict[str, Any]) -> Dict[str, Any]:
url = f"{self.api_base}{path}"
try:
response = await self.async_handler.post(
url=url, headers=self._headers(), json=body, timeout=15.0
)
response.raise_for_status()
except Exception as e:
self._reraise_api_error(e, path)
return response.json()
def _headers(self) -> Dict[str, str]:
return {
"Authorization": f"Bearer {self.peyeeye_api_key}",
"Content-Type": "application/json",
}
@staticmethod
def _reraise_api_error(error: Exception, path: str) -> None:
if isinstance(error, HTTPException):
raise error
if HTTPX_AVAILABLE and httpx is not None:
if isinstance(error, httpx.TimeoutException):
raise PeyeeyeGuardrailAPIError(f"peyeeye {path} timed out") from error
if isinstance(error, httpx.HTTPStatusError):
status = error.response.status_code
if status == 401:
raise PeyeeyeGuardrailMissingSecrets(
"Invalid peyeeye API key"
) from error
if status == 429:
raise PeyeeyeGuardrailAPIError(
"peyeeye rate limit exceeded"
) from error
raise PeyeeyeGuardrailAPIError(
f"peyeeye {path} returned {status}"
) from error
raise PeyeeyeGuardrailAPIError(
f"peyeeye {path} failed: {error}"
) from error
@staticmethod
def get_config_model() -> Optional[Type["GuardrailConfigModel"]]:
from litellm.types.proxy.guardrails.guardrail_hooks.peyeeye import (
PeyeeyeGuardrailConfigModel,
)
return PeyeeyeGuardrailConfigModel
# ---------------------------------------------------------------- text helpers
def _iter_message_text(messages: List[Dict[str, Any]]):
"""Yield (message_index, part_path, text) for every text-bearing chunk.
``part_path`` is either ``"content"`` for a plain string message or an
int index into the multimodal content list.
"""
for i, msg in enumerate(messages):
if not isinstance(msg, dict):
continue
content = msg.get("content")
if isinstance(content, str):
if content:
yield i, "content", content
elif isinstance(content, list):
for j, part in enumerate(content):
if isinstance(part, dict) and part.get("type") == "text":
text = part.get("text", "")
if text:
yield i, j, text
def _set_message_text(message: Dict[str, Any], part_path, value: str) -> None:
if part_path == "content":
message["content"] = value
return
parts = message.get("content")
if isinstance(parts, list) and isinstance(part_path, int) and part_path < len(parts):
part = parts[part_path]
if isinstance(part, dict):
part["text"] = value

19
litellm/types/guardrails.py
View file

@ -100,6 +100,7 @@ class SupportedGuardrailIntegrations(Enum):
MCP_JWT_SIGNER = "mcp_jwt_signer"
LLM_AS_A_JUDGE = "llm_as_a_judge"
QOSTODIAN_NEXUS = "qostodian_nexus"
PEYEEYE = "peyeeye"
class Role(Enum):
@ -446,6 +447,23 @@ class LassoGuardrailConfigModel(BaseModel):
)
class PeyeeyeGuardrailConfigModel(BaseModel):
"""Configuration parameters for the peyeeye PII redaction & rehydration guardrail"""
peyeeye_locale: Optional[str] = Field(
default="auto",
description="BCP-47 language tag for PII detection. 'auto' lets peyeeye detect.",
)
peyeeye_entities: Optional[List[str]] = Field(
default=None,
description="Restrict detection to these entity IDs (e.g. ['EMAIL', 'CARD']). Omit to detect all 60+ built-in entities.",
)
peyeeye_session_mode: Optional[Literal["stateful", "stateless"]] = Field(
default="stateful",
description="'stateful' uses peyeeye sessions; 'stateless' returns a sealed rehydration key (no PII retained server-side).",
)
class PillarGuardrailConfigModel(BaseModel):
"""Configuration parameters for the Pillar Security guardrail"""
@ -762,6 +780,7 @@ class LitellmParams(
BedrockGuardrailConfigModel,
LakeraV2GuardrailConfigModel,
LassoGuardrailConfigModel,
PeyeeyeGuardrailConfigModel,
PillarGuardrailConfigModel,
GraySwanGuardrailConfigModel,
NomaGuardrailConfigModel,

37
litellm/types/proxy/guardrails/guardrail_hooks/peyeeye.py Normal file
View file

@ -0,0 +1,37 @@
from typing import List, Literal, Optional
from pydantic import BaseModel, Field
from .base import GuardrailConfigModel
class PeyeeyeGuardrailConfigModelOptionalParams(BaseModel):
peyeeye_locale: Optional[str] = Field(
default="auto",
description="BCP-47 language tag for PII detection. 'auto' lets peyeeye detect.",
)
peyeeye_entities: Optional[List[str]] = Field(
default=None,
description="Restrict detection to these entity IDs (e.g. ['EMAIL', 'CARD']). Omit to detect all 60+ built-in entities.",
)
peyeeye_session_mode: Optional[Literal["stateful", "stateless"]] = Field(
default="stateful",
description="'stateful' stores token→value mappings under a peyeeye session id. 'stateless' returns a sealed rehydration key — no PII retained server-side.",
)
class PeyeeyeGuardrailConfigModel(
GuardrailConfigModel[PeyeeyeGuardrailConfigModelOptionalParams]
):
api_key: Optional[str] = Field(
default=None,
description="Peyeeye API key. Falls back to the `PEYEEYE_API_KEY` environment variable.",
)
api_base: Optional[str] = Field(
default=None,
description="Peyeeye API base URL. Defaults to https://api.peyeeye.ai. Also reads `PEYEEYE_API_BASE`.",
)
@staticmethod
def ui_friendly_name() -> str:
return "Peyeeye PII Redaction & Rehydration"

222
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_peyeeye.py Normal file
View file

@ -0,0 +1,222 @@
import os
import sys
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
sys.path.insert(0, os.path.abspath("../.."))
import litellm
from litellm import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.peyeeye.peyeeye import (
PeyeeyeGuardrail,
PeyeeyeGuardrailAPIError,
PeyeeyeGuardrailMissingSecrets,
)
from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
def _ok(json_payload: dict):
resp = MagicMock()
resp.json.return_value = json_payload
resp.raise_for_status = MagicMock()
return resp
def test_peyeeye_init_requires_api_key():
for var in ("PEYEEYE_API_KEY", "PEYEEYE_API_BASE"):
os.environ.pop(var, None)
with pytest.raises(PeyeeyeGuardrailMissingSecrets):
PeyeeyeGuardrail(guardrail_name="t")
def test_peyeeye_init_reads_env():
os.environ["PEYEEYE_API_KEY"] = "pk_test"
try:
g = PeyeeyeGuardrail(guardrail_name="t")
assert g.peyeeye_api_key == "pk_test"
assert g.api_base == "https://api.peyeeye.ai"
assert g.peyeeye_session_mode == "stateful"
finally:
del os.environ["PEYEEYE_API_KEY"]
def test_peyeeye_init_explicit_args():
g = PeyeeyeGuardrail(
peyeeye_api_key="pk_x",
api_base="https://api.example/",
peyeeye_locale="en",
peyeeye_entities=["EMAIL"],
peyeeye_session_mode="stateless",
guardrail_name="t",
)
assert g.peyeeye_api_key == "pk_x"
assert g.api_base == "https://api.example"
assert g.peyeeye_locale == "en"
assert g.peyeeye_entities == ["EMAIL"]
assert g.peyeeye_session_mode == "stateless"
def test_peyeeye_guardrail_config_via_init():
litellm.guardrail_name_config_map = {}
os.environ["PEYEEYE_API_KEY"] = "pk_test"
try:
init_guardrails_v2(
all_guardrails=[
{
"guardrail_name": "peyeeye-pre",
"litellm_params": {
"guardrail": "peyeeye",
"mode": "pre_call",
"default_on": True,
},
}
],
config_file_path="",
)
finally:
del os.environ["PEYEEYE_API_KEY"]
@pytest.mark.asyncio
async def test_pre_call_redacts_messages_and_caches_session():
g = PeyeeyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g.async_handler = MagicMock()
g.async_handler.post = AsyncMock(
return_value=_ok(
{"text": ["hi [EMAIL_1]"], "session_id": "ses_abc"}
)
)
cache = DualCache()
data = {
"messages": [{"role": "user", "content": "hi alice@acme.com"}],
"litellm_call_id": "call-1",
}
user = UserAPIKeyAuth(api_key="x")
out = await g.async_pre_call_hook(user, cache, data, "completion")
assert out["messages"][0]["content"] == "hi [EMAIL_1]"
cached = cache.get_cache("peyeeye_session:call-1")
assert cached == "ses_abc"
@pytest.mark.asyncio
async def test_pre_call_stateless_returns_skey():
g = PeyeeyeGuardrail(
peyeeye_api_key="pk",
peyeeye_session_mode="stateless",
guardrail_name="t",
)
g.async_handler = MagicMock()
g.async_handler.post = AsyncMock(
return_value=_ok(
{"text": ["[EMAIL_1]"], "rehydration_key": "skey_xyz"}
)
)
cache = DualCache()
data = {
"messages": [{"role": "user", "content": "alice@acme.com"}],
"litellm_call_id": "call-2",
}
await g.async_pre_call_hook(UserAPIKeyAuth(api_key="x"), cache, data, "completion")
sent_body = g.async_handler.post.call_args.kwargs["json"]
assert sent_body["session"] == "stateless"
assert cache.get_cache("peyeeye_session:call-2") == "skey_xyz"
@pytest.mark.asyncio
async def test_pre_call_skips_when_no_messages():
g = PeyeeyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g.async_handler = MagicMock()
g.async_handler.post = AsyncMock()
cache = DualCache()
data = {"messages": [], "litellm_call_id": "x"}
out = await g.async_pre_call_hook(
UserAPIKeyAuth(api_key="x"), cache, data, "completion"
)
assert out is data
g.async_handler.post.assert_not_called()
@pytest.mark.asyncio
async def test_post_call_rehydrates_response():
g = PeyeeyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g.async_handler = MagicMock()
g.async_handler.post = AsyncMock(
return_value=_ok({"text": "Reply to alice@acme.com", "replaced": 1})
)
g.async_handler.delete = AsyncMock()
# Seed the session id into litellm.cache so the post-call hook finds it.
litellm.cache = MagicMock()
litellm.cache.get_cache = MagicMock(return_value="ses_abc")
litellm.cache.delete_cache = MagicMock()
response = litellm.ModelResponse()
response.choices = [
litellm.utils.Choices(
finish_reason="stop",
index=0,
message=litellm.utils.Message(
content="Reply to [EMAIL_1]", role="assistant"
),
)
]
data = {"litellm_call_id": "call-1"}
out = await g.async_post_call_success_hook(
data, UserAPIKeyAuth(api_key="x"), response
)
assert out.choices[0].message.content == "Reply to alice@acme.com"
g.async_handler.delete.assert_awaited()
litellm.cache.delete_cache.assert_called_once_with("peyeeye_session:call-1")
litellm.cache = None
@pytest.mark.asyncio
async def test_post_call_noop_without_session():
g = PeyeeyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g.async_handler = MagicMock()
g.async_handler.post = AsyncMock()
litellm.cache = None
response = litellm.ModelResponse()
response.choices = [
litellm.utils.Choices(
finish_reason="stop",
index=0,
message=litellm.utils.Message(content="hello", role="assistant"),
)
]
out = await g.async_post_call_success_hook(
{"litellm_call_id": "no-session"}, UserAPIKeyAuth(api_key="x"), response
)
assert out.choices[0].message.content == "hello"
g.async_handler.post.assert_not_called()
@pytest.mark.asyncio
async def test_redact_api_error_raises_typed():
g = PeyeeyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g.async_handler = MagicMock()
bad = MagicMock()
bad.raise_for_status.side_effect = RuntimeError("boom")
g.async_handler.post = AsyncMock(return_value=bad)
cache = DualCache()
with pytest.raises(PeyeeyeGuardrailAPIError):
await g.async_pre_call_hook(
UserAPIKeyAuth(api_key="x"),
cache,
{
"messages": [{"role": "user", "content": "hi"}],
"litellm_call_id": "c",
},
"completion",
)