From 44d9737609b026a453e680c9e88dbbb3702296af Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 16 Jul 2026 01:27:44 +0300 Subject: [PATCH] fix(llm_guard): apply sanitized prompt returned by moderation API to request (#33331) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../enterprise_callbacks/llm_guard.py | 141 ++++++++++++------ tests/local_testing/test_llm_guard.py | 79 ++++++++-- 2 files changed, 159 insertions(+), 61 deletions(-) diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py index ad8aabf77b6..d10b5a2ab09 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py @@ -7,7 +7,8 @@ # Thank you users! We ❤️ you! - Krrish & Ishaan ## This provides an LLM Guard Integration for content moderation on the proxy -from typing import Literal, Optional +import asyncio +from typing import Optional import aiohttp from fastapi import HTTPException @@ -18,7 +19,6 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.secret_managers.main import get_secret_str from litellm.types.utils import CallTypesLiteral -from litellm.utils import get_formatted_prompt class _ENTERPRISE_LLMGuard(CustomLogger): @@ -46,45 +46,44 @@ class _ENTERPRISE_LLMGuard(CustomLogger): except Exception: pass - async def moderation_check(self, text: str): + async def moderation_check(self, text: str) -> str: """ + Runs the LLM Guard moderation check on ``text``. + + Raises an HTTPException when the content violates the safety policy; + otherwise returns the sanitized prompt from LLM Guard, falling back to + the original text when the API does not provide one. + [TODO] make this more performant for high-throughput scenario """ try: - async with aiohttp.ClientSession() as session: - if self.mock_redacted_text is not None: - redacted_text = self.mock_redacted_text - else: - # Make the first request to /analyze - analyze_url = f"{self.llm_guard_api_base}analyze/prompt" - verbose_proxy_logger.debug("Making request to: %s", analyze_url) - analyze_payload = {"prompt": text} - redacted_text = None + if self.mock_redacted_text is not None: + redacted_text = self.mock_redacted_text + else: + analyze_url = f"{self.llm_guard_api_base}analyze/prompt" + verbose_proxy_logger.debug("Making request to: %s", analyze_url) + async with aiohttp.ClientSession() as session: async with session.post( - analyze_url, json=analyze_payload + analyze_url, json={"prompt": text} ) as response: redacted_text = await response.json() - verbose_proxy_logger.debug( - f"LLM Guard: Received response - {redacted_text}" + verbose_proxy_logger.debug( + f"LLM Guard: Received response - {redacted_text}" + ) + if redacted_text is None: + raise HTTPException( + status_code=500, + detail={ + "error": f"Invalid content moderation response: {redacted_text}" + }, ) - if redacted_text is not None: - if ( - redacted_text.get("is_valid", None) is not None - and redacted_text["is_valid"] is False - ): - raise HTTPException( - status_code=400, - detail={"error": "Violated content safety policy"}, - ) - else: - pass - else: - raise HTTPException( - status_code=500, - detail={ - "error": f"Invalid content moderation response: {redacted_text}" - }, - ) + if redacted_text.get("is_valid", None) is False: + raise HTTPException( + status_code=400, + detail={"error": "Violated content safety policy"}, + ) + sanitized_prompt = redacted_text.get("sanitized_prompt") + return sanitized_prompt if isinstance(sanitized_prompt, str) else text except Exception as e: verbose_proxy_logger.exception( "litellm.enterprise.enterprise_hooks.llm_guard::moderation_check - Exception occurred - {}".format( @@ -138,23 +137,75 @@ class _ENTERPRISE_LLMGuard(CustomLogger): return self.print_verbose("Makes LLM Guard Check") - try: - assert call_type in [ - "completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - ] - except Exception: + if call_type not in [ + "completion", + "embeddings", + "image_generation", + "moderation", + "audio_transcription", + ]: self.print_verbose( f"Call Type - {call_type}, not in accepted list - ['completion','embeddings','image_generation','moderation','audio_transcription']" ) return data - formatted_prompt = get_formatted_prompt(data=data, call_type=call_type) # type: ignore - self.print_verbose(f"LLM Guard, formatted_prompt: {formatted_prompt}") - return await self.moderation_check(text=formatted_prompt) + return await self._moderate_request(data=data) + + async def _moderate_request(self, data: dict) -> dict: + """ + Sanitizes the request in place using the prompt returned by LLM Guard so + the provider-bound request carries the redacted content, then returns it. + """ + messages = data.get("messages") + if messages is not None: + data["messages"] = list( + await asyncio.gather( + *(self._moderate_message(message) for message in messages) + ) + ) + return data + + input_ = data.get("input") + if input_ is not None: + data["input"] = await self._moderate_input(input_) + return data + + prompt = data.get("prompt") + if isinstance(prompt, str): + data["prompt"] = await self.moderation_check(text=prompt) + return data + + async def _moderate_message(self, message: dict) -> dict: + content = message.get("content") + if isinstance(content, str): + return {**message, "content": await self.moderation_check(text=content)} + if isinstance(content, list): + return { + **message, + "content": list( + await asyncio.gather( + *(self._moderate_content_part(part) for part in content) + ) + ), + } + return message + + async def _moderate_content_part(self, part: dict) -> dict: + if part.get("type") == "text" and isinstance(part.get("text"), str): + return {**part, "text": await self.moderation_check(text=part["text"])} + return part + + async def _moderate_input(self, input_: object) -> object: + if isinstance(input_, str): + return await self.moderation_check(text=input_) + if isinstance(input_, list): + return [ + await self.moderation_check(text=item) + if isinstance(item, str) + else item + for item in input_ + ] + return input_ async def async_post_call_streaming_hook( self, user_api_key_dict: UserAPIKeyAuth, response: str diff --git a/tests/local_testing/test_llm_guard.py b/tests/local_testing/test_llm_guard.py index da881e19d3c..78bbd1c0af8 100644 --- a/tests/local_testing/test_llm_guard.py +++ b/tests/local_testing/test_llm_guard.py @@ -28,8 +28,11 @@ from litellm.caching.caching import DualCache @pytest.mark.asyncio async def test_llm_guard_valid_response(): """ - Tests to see llm guard raises an error for a flagged response + A valid (is_valid=True) LLM Guard response must apply the returned + sanitized_prompt back onto the request data so the provider receives the + redacted content. """ + litellm.llm_guard_mode = "all" input_a_anonymizer_results = { "sanitized_prompt": "hello world", "is_valid": True, @@ -44,21 +47,65 @@ async def test_llm_guard_valid_response(): user_api_key_dict = UserAPIKeyAuth(api_key=_api_key) local_cache = DualCache() - try: - await llm_guard.async_moderation_hook( - data={ - "messages": [ - { - "role": "user", - "content": "hello world, my name is Jane Doe. My number is: 23r323r23r2wwkl", - } - ] - }, - user_api_key_dict=user_api_key_dict, - call_type="completion", - ) - except Exception as e: - pytest.fail(f"An exception occurred - {str(e)}") + data = { + "messages": [ + { + "role": "user", + "content": "hello world, my name is Jane Doe. My number is: 23r323r23r2wwkl", + } + ] + } + + result = await llm_guard.async_moderation_hook( + data=data, + user_api_key_dict=user_api_key_dict, + call_type="completion", + ) + + assert result is data + assert data["messages"][0]["content"] == "hello world" + + +@pytest.mark.asyncio +async def test_llm_guard_sanitizes_multimodal_and_input(): + """ + Sanitization must reach text parts of multimodal message content and the + ``input`` field (embeddings/moderation) while leaving non-text parts intact. + """ + litellm.llm_guard_mode = "all" + llm_guard = _ENTERPRISE_LLMGuard( + mock_testing=True, + mock_redacted_text={ + "sanitized_prompt": "email: [REDACTED]", + "is_valid": True, + "scanners": {"Regex": 0.0}, + }, + ) + user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("sk-12345")) + + image_part = {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}} + data = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "email: person@example.com"}, + image_part, + ], + } + ] + } + result = await llm_guard.async_moderation_hook( + data=data, user_api_key_dict=user_api_key_dict, call_type="completion" + ) + assert result["messages"][0]["content"][0]["text"] == "email: [REDACTED]" + assert result["messages"][0]["content"][1] == image_part + + input_data = {"input": ["email: person@example.com", "another prompt"]} + input_result = await llm_guard.async_moderation_hook( + data=input_data, user_api_key_dict=user_api_key_dict, call_type="embeddings" + ) + assert input_result["input"] == ["email: [REDACTED]", "email: [REDACTED]"] @pytest.mark.asyncio