From 0fa294b18cf251564c974e3cea853cfec1b123c6 Mon Sep 17 00:00:00 2001 From: Tim Date: Sat, 25 Apr 2026 22:25:49 -0500 Subject: [PATCH] fix(guardrails/peyeeye): rename classes to PEyeEye, guard length mismatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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. --- .../guardrail_hooks/peyeeye/__init__.py | 6 +- .../guardrail_hooks/peyeeye/peyeeye.py | 30 ++++---- litellm/types/guardrails.py | 4 +- .../guardrails/guardrail_hooks/peyeeye.py | 6 +- .../guardrail_hooks/test_peyeeye.py | 71 +++++++++++++++---- 5 files changed, 83 insertions(+), 34 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/peyeeye/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/peyeeye/__init__.py index 457b2208980..a67f85683b6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/peyeeye/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/peyeeye/__init__.py @@ -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, } diff --git a/litellm/proxy/guardrails/guardrail_hooks/peyeeye/peyeeye.py b/litellm/proxy/guardrails/guardrail_hooks/peyeeye/peyeeye.py index 42b5e53da6e..f86d0b3f22b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/peyeeye/peyeeye.py +++ b/litellm/proxy/guardrails/guardrail_hooks/peyeeye/peyeeye.py @@ -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 diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index b6192843a46..a115f696051 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -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, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/peyeeye.py b/litellm/types/proxy/guardrails/guardrail_hooks/peyeeye.py index d69068d273c..ce7f1f8ddbc 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/peyeeye.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/peyeeye.py @@ -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, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_peyeeye.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_peyeeye.py index e559fb12eaf..bc33c2218b6 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_peyeeye.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_peyeeye.py @@ -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" + )