Skip to content

fix(llm_guard): apply sanitized prompt returned by moderation API to request #33331

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yassin-berriai merged 1 commit into from
Jul 15, 2026
+159 −61
Merged

fix(llm_guard): apply sanitized prompt returned by moderation API to request #33331

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
141 changes: 96 additions & 45 deletions enterprise/litellm_enterprise/enterprise_callbacks/llm_guard.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -18,7 +19,6 @@
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):
Expand Down Expand Up @@ -46,45 +46,44 @@ def print_verbose(self, print_statement):
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(
Expand Down Expand Up @@ -138,23 +137,75 @@ async def async_moderation_hook(
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_
]
Comment on lines +201 to +207

@greptile-apps greptile-apps Bot Jul 15, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Sequential awaits in _moderate_input list branch serialize LLM Guard API calls. For a batch embedding request with N string inputs, this method will issue N back-to-back HTTP calls to the LLM Guard server, while _moderate_message and _moderate_request both use asyncio.gather for concurrency. In production with large batches (e.g. 50 embeddings), this could add 50× the round-trip latency of a single guard call.

Suggested change
if isinstance(input_, list):
return [
await self.moderation_check(text=item)
if isinstance(item, str)
else item
for item in input_
]
if isinstance(input_, list):
async def _check_or_pass(item: object) -> object:
return await self.moderation_check(text=item) if isinstance(item, str) else item
return list(await asyncio.gather(*(_check_or_pass(item) for item in input_)))

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

return input_

async def async_post_call_streaming_hook(
self, user_api_key_dict: UserAPIKeyAuth, response: str
Expand Down
79 changes: 63 additions & 16 deletions tests/local_testing/test_llm_guard.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,11 @@
@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,
Expand All @@ -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
Expand Down
Loading