diff --git a/litellm/proxy/guardrails/guardrail_hooks/peyeeye/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/peyeeye/__init__.py new file mode 100644 index 00000000000..457b2208980 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/peyeeye/__init__.py @@ -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, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/peyeeye/peyeeye.py b/litellm/proxy/guardrails/guardrail_hooks/peyeeye/peyeeye.py new file mode 100644 index 00000000000..8c5cf4564ec --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/peyeeye/peyeeye.py @@ -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 diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 04347aebe3b..b6192843a46 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -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, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/peyeeye.py b/litellm/types/proxy/guardrails/guardrail_hooks/peyeeye.py new file mode 100644 index 00000000000..d69068d273c --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/peyeeye.py @@ -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" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_peyeeye.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_peyeeye.py new file mode 100644 index 00000000000..60bcfb82216 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_peyeeye.py @@ -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", + )