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

fix(guardrails/peyeeye): rename classes to PEyeEye, guard length mismatch

Browse source 
- P1: `_redact_batch` returning fewer texts than were sent would silently
  truncate via `zip(...)` and forward un-redacted messages to the LLM.
  Raise `PEyeEyeGuardrailAPIError` on any length mismatch so a partial
  response can never leak PII downstream.
- Rename Python classes to the project's preferred camelization:
  `PeyeeyeGuardrail` → `PEyeEyeGuardrail` and friends. The integration id,
  env vars (`PEYEEYE_API_KEY`, …) and config keys (`peyeeye_*`) are
  unchanged so existing configs keep working.
- Add tests for the length-guard and the previously-fixed unexpected-
  shape guard so both PII-passthrough paths are now covered.
This commit is contained in:
Tim 2026-04-25 22:25:49 -05:00
parent 3eda6087d1
commit 0fa294b18c
5 changed files with 83 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/guardrails/guardrail_hooks/peyeeye/__init__.py
View file

@ -2,7 +2,7 @@ from typing import TYPE_CHECKING
from litellm.types.guardrails import SupportedGuardrailIntegrations
from .peyeeye import PeyeeyeGuardrail
from .peyeeye import PEyeEyeGuardrail
if TYPE_CHECKING:
from litellm.types.guardrails import Guardrail, LitellmParams
@ -11,7 +11,7 @@ if TYPE_CHECKING:
def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"):
import litellm
_peyeeye_callback = PeyeeyeGuardrail(
_peyeeye_callback = PEyeEyeGuardrail(
guardrail_name=guardrail.get("guardrail_name", ""),
api_key=litellm_params.api_key,
api_base=litellm_params.api_base,
@ -32,5 +32,5 @@ guardrail_initializer_registry = {
guardrail_class_registry = {
SupportedGuardrailIntegrations.PEYEEYE.value: PeyeeyeGuardrail,
SupportedGuardrailIntegrations.PEYEEYE.value: PEyeEyeGuardrail,
}

30
litellm/proxy/guardrails/guardrail_hooks/peyeeye/peyeeye.py
View file

@ -50,15 +50,15 @@ DEFAULT_API_BASE = "https://api.peyeeye.ai"
SESSION_CACHE_TTL_SECONDS = 3600
class PeyeeyeGuardrailMissingSecrets(Exception):
class PEyeEyeGuardrailMissingSecrets(Exception):
"""Raised when the peyeeye API key is missing."""
class PeyeeyeGuardrailAPIError(Exception):
class PEyeEyeGuardrailAPIError(Exception):
"""Raised when the peyeeye API returns an error."""
class PeyeeyeGuardrail(CustomGuardrail):
class PEyeEyeGuardrail(CustomGuardrail):
"""Peyeeye PII redaction + rehydration guardrail.
Pre-call hook redacts PII from each message's ``content`` and stores the
@ -90,7 +90,7 @@ class PeyeeyeGuardrail(CustomGuardrail):
peyeeye_api_key or api_key or os.environ.get("PEYEEYE_API_KEY")
)
if self.peyeeye_api_key is None:
raise PeyeeyeGuardrailMissingSecrets(
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."
)
@ -147,6 +147,12 @@ class PeyeeyeGuardrail(CustomGuardrail):
redacted_texts, session_id = await self._redact_batch(
[t for _, _, t in text_parts]
)
if len(redacted_texts) != len(text_parts):
raise PEyeEyeGuardrailAPIError(
f"peyeeye /v1/redact returned {len(redacted_texts)} texts for "
f"{len(text_parts)} inputs; refusing to forward partially-"
"redacted data"
)
for (msg_idx, part_path, _), redacted in zip(text_parts, redacted_texts):
_set_message_text(messages[msg_idx], part_path, redacted)
@ -248,7 +254,7 @@ class PeyeeyeGuardrail(CustomGuardrail):
elif isinstance(out_text, list):
redacted = [str(x) for x in out_text]
else:
raise PeyeeyeGuardrailAPIError(
raise PEyeEyeGuardrailAPIError(
"peyeeye /v1/redact returned unexpected response shape; "
"refusing to forward unredacted text"
)
@ -296,31 +302,31 @@ class PeyeeyeGuardrail(CustomGuardrail):
def _reraise_api_error(error: Exception, path: str) -> NoReturn:
if HTTPX_AVAILABLE and httpx is not None:
if isinstance(error, httpx.TimeoutException):
raise PeyeeyeGuardrailAPIError(f"peyeeye {path} timed out") from error
raise PEyeEyeGuardrailAPIError(f"peyeeye {path} timed out") from error
if isinstance(error, httpx.HTTPStatusError):
status = error.response.status_code
if status == 401:
raise PeyeeyeGuardrailMissingSecrets(
raise PEyeEyeGuardrailMissingSecrets(
"Invalid peyeeye API key"
) from error
if status == 429:
raise PeyeeyeGuardrailAPIError(
raise PEyeEyeGuardrailAPIError(
"peyeeye rate limit exceeded"
) from error
raise PeyeeyeGuardrailAPIError(
raise PEyeEyeGuardrailAPIError(
f"peyeeye {path} returned {status}"
) from error
raise PeyeeyeGuardrailAPIError(
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,
PEyeEyeGuardrailConfigModel,
)
return PeyeeyeGuardrailConfigModel
return PEyeEyeGuardrailConfigModel
# ---------------------------------------------------------------- text helpers

4
litellm/types/guardrails.py
View file

@ -447,7 +447,7 @@ class LassoGuardrailConfigModel(BaseModel):
)
class PeyeeyeGuardrailConfigModel(BaseModel):
class PEyeEyeGuardrailConfigModel(BaseModel):
"""Configuration parameters for the peyeeye PII redaction & rehydration guardrail"""
peyeeye_locale: Optional[str] = Field(
@ -780,7 +780,7 @@ class LitellmParams(
BedrockGuardrailConfigModel,
LakeraV2GuardrailConfigModel,
LassoGuardrailConfigModel,
PeyeeyeGuardrailConfigModel,
PEyeEyeGuardrailConfigModel,
PillarGuardrailConfigModel,
GraySwanGuardrailConfigModel,
NomaGuardrailConfigModel,

6
litellm/types/proxy/guardrails/guardrail_hooks/peyeeye.py
View file

@ -5,7 +5,7 @@ from pydantic import BaseModel, Field
from .base import GuardrailConfigModel
class PeyeeyeGuardrailConfigModelOptionalParams(BaseModel):
class PEyeEyeGuardrailConfigModelOptionalParams(BaseModel):
peyeeye_locale: Optional[str] = Field(
default="auto",
description="BCP-47 language tag for PII detection. 'auto' lets peyeeye detect.",
@ -20,8 +20,8 @@ class PeyeeyeGuardrailConfigModelOptionalParams(BaseModel):
)
class PeyeeyeGuardrailConfigModel(
GuardrailConfigModel[PeyeeyeGuardrailConfigModelOptionalParams]
class PEyeEyeGuardrailConfigModel(
GuardrailConfigModel[PEyeEyeGuardrailConfigModelOptionalParams]
):
api_key: Optional[str] = Field(
default=None,

71
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_peyeeye.py
View file

@ -10,9 +10,9 @@ import litellm
from litellm import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.peyeeye.peyeeye import (
PeyeeyeGuardrail,
PeyeeyeGuardrailAPIError,
PeyeeyeGuardrailMissingSecrets,
PEyeEyeGuardrail,
PEyeEyeGuardrailAPIError,
PEyeEyeGuardrailMissingSecrets,
)
from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
@ -27,14 +27,14 @@ def _ok(json_payload: dict):
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")
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")
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"
@ -43,7 +43,7 @@ def test_peyeeye_init_reads_env():
def test_peyeeye_init_explicit_args():
g = PeyeeyeGuardrail(
g = PEyeEyeGuardrail(
peyeeye_api_key="pk_x",
api_base="https://api.example/",
peyeeye_locale="en",
@ -81,7 +81,7 @@ def test_peyeeye_guardrail_config_via_init():
@pytest.mark.asyncio
async def test_pre_call_redacts_messages_and_caches_session():
g = PeyeeyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g = PEyeEyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g.async_handler = MagicMock()
g.async_handler.post = AsyncMock(
return_value=_ok(
@ -107,7 +107,7 @@ async def test_pre_call_redacts_messages_and_caches_session():
@pytest.mark.asyncio
async def test_pre_call_stateless_returns_skey():
g = PeyeeyeGuardrail(
g = PEyeEyeGuardrail(
peyeeye_api_key="pk",
peyeeye_session_mode="stateless",
guardrail_name="t",
@ -137,7 +137,7 @@ async def test_pre_call_stateless_returns_skey():
@pytest.mark.asyncio
async def test_pre_call_skips_when_no_messages():
g = PeyeeyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g = PEyeEyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g.async_handler = MagicMock()
g.async_handler.post = AsyncMock()
@ -157,7 +157,7 @@ async def test_pre_and_post_call_roundtrip_uses_shared_cache():
global_cache,
)
g = PeyeeyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g = PEyeEyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g.async_handler = MagicMock()
g.async_handler.post = AsyncMock(
side_effect=[
@ -195,7 +195,7 @@ async def test_pre_and_post_call_roundtrip_uses_shared_cache():
@pytest.mark.asyncio
async def test_post_call_noop_without_session():
g = PeyeeyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g = PEyeEyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g.async_handler = MagicMock()
g.async_handler.post = AsyncMock()
@ -216,7 +216,7 @@ async def test_post_call_noop_without_session():
@pytest.mark.asyncio
async def test_redact_api_error_raises_typed():
g = PeyeeyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g = PEyeEyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g.async_handler = MagicMock()
bad = MagicMock()
@ -224,7 +224,7 @@ async def test_redact_api_error_raises_typed():
g.async_handler.post = AsyncMock(return_value=bad)
cache = DualCache()
with pytest.raises(PeyeeyeGuardrailAPIError):
with pytest.raises(PEyeEyeGuardrailAPIError):
await g.async_pre_call_hook(
UserAPIKeyAuth(api_key="x"),
cache,
@ -234,3 +234,46 @@ async def test_redact_api_error_raises_typed():
},
"completion",
)
@pytest.mark.asyncio
async def test_pre_call_raises_on_length_mismatch():
"""If /v1/redact returns fewer texts than sent, refuse to forward."""
g = PEyeEyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g.async_handler = MagicMock()
g.async_handler.post = AsyncMock(
return_value=_ok({"text": ["[EMAIL_1]"], "session_id": "ses_x"})
)
cache = DualCache()
data = {
"messages": [
{"role": "user", "content": "hi alice@acme.com"},
{"role": "user", "content": "hi bob@acme.com"},
],
"litellm_call_id": "len-1",
}
with pytest.raises(PEyeEyeGuardrailAPIError, match="partially-redacted"):
await g.async_pre_call_hook(
UserAPIKeyAuth(api_key="x"), cache, data, "completion"
)
@pytest.mark.asyncio
async def test_pre_call_raises_on_unexpected_response_shape():
"""If /v1/redact returns neither str nor list for `text`, refuse to forward."""
g = PEyeEyeGuardrail(peyeeye_api_key="pk", guardrail_name="t")
g.async_handler = MagicMock()
g.async_handler.post = AsyncMock(
return_value=_ok({"text": 42, "session_id": "ses_x"})
)
cache = DualCache()
data = {
"messages": [{"role": "user", "content": "hi alice@acme.com"}],
"litellm_call_id": "shape-1",
}
with pytest.raises(PEyeEyeGuardrailAPIError, match="unexpected response shape"):
await g.async_pre_call_hook(
UserAPIKeyAuth(api_key="x"), cache, data, "completion"
)