ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Fix] Guardrails API - Ensure OpenAI Moderations Guard works with OpenAI Embeddings (#20523)

Browse Source 
* init OpenAIEmbeddingsHandler

* init apply_guardrail

* use apply guardrails for OpenAI moderations

* test_embeddings_handler_string_input

* test_openai_moderation_guardrail_apply_guardrail

* fix typing

* test_openai_moderation_responses_api_input_field

* test fixes
This commit is contained in:
Ishaan Jaff 2026-02-05 14:40:15 -08:00 committed by GitHub
parent 0649720f79
commit 887a907e42
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 431 additions and 217 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
litellm/llms/openai/embeddings/guardrail_translation/__init__.py Normal file
View File

@ -0,0 +1,13 @@
"""OpenAI Embeddings handler for Unified Guardrails."""
from litellm.llms.openai.embeddings.guardrail_translation.handler import (
OpenAIEmbeddingsHandler,
)
from litellm.types.utils import CallTypes
guardrail_translation_mappings = {
CallTypes.embedding: OpenAIEmbeddingsHandler,
CallTypes.aembedding: OpenAIEmbeddingsHandler,
}
__all__ = ["guardrail_translation_mappings", "OpenAIEmbeddingsHandler"]

179
litellm/llms/openai/embeddings/guardrail_translation/handler.py Normal file
View File

@ -0,0 +1,179 @@
"""
OpenAI Embeddings Handler for Unified Guardrails
This module provides guardrail translation support for OpenAI's embeddings endpoint.
The handler processes the 'input' parameter for guardrails.
"""
from typing import TYPE_CHECKING, Any, List, Optional, Union
from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.types.utils import EmbeddingResponse
class OpenAIEmbeddingsHandler(BaseTranslation):
"""
Handler for processing OpenAI embeddings requests with guardrails.
This class provides methods to:
1. Process input text (pre-call hook)
2. Process output response (post-call hook) - embeddings don't typically need output guardrails
The handler specifically processes the 'input' parameter which can be:
- A single string
- A list of strings (for batch embeddings)
- A list of integers (token IDs - not processed by guardrails)
- A list of lists of integers (batch token IDs - not processed by guardrails)
"""
async def process_input_messages(
self,
data: dict,
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Optional[Any] = None,
) -> Any:
"""
Process input text by applying guardrails to text content.
Args:
data: Request data dictionary containing 'input' parameter
guardrail_to_apply: The guardrail instance to apply
litellm_logging_obj: Optional logging object
Returns:
Modified data with guardrails applied to input
"""
input_data = data.get("input")
if input_data is None:
verbose_proxy_logger.debug(
"OpenAI Embeddings: No input found in request data"
)
return data
if isinstance(input_data, str):
data = await self._process_string_input(
data, input_data, guardrail_to_apply, litellm_logging_obj
)
elif isinstance(input_data, list):
data = await self._process_list_input(
data, input_data, guardrail_to_apply, litellm_logging_obj
)
else:
verbose_proxy_logger.warning(
"OpenAI Embeddings: Unexpected input type: %s. Expected string or list.",
type(input_data),
)
return data
async def _process_string_input(
self,
data: dict,
input_data: str,
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Optional[Any],
) -> dict:
"""Process a single string input through the guardrail."""
inputs = GenericGuardrailAPIInputs(texts=[input_data])
if model := data.get("model"):
inputs["model"] = model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
)
if guardrailed_texts := guardrailed_inputs.get("texts"):
data["input"] = guardrailed_texts[0]
verbose_proxy_logger.debug(
"OpenAI Embeddings: Applied guardrail to string input. "
"Original length: %d, New length: %d",
len(input_data),
len(data["input"]),
)
return data
async def _process_list_input(
self,
data: dict,
input_data: List[Union[str, int, List[int]]],
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Optional[Any],
) -> dict:
"""Process a list input through the guardrail (if it contains strings)."""
if len(input_data) == 0:
return data
first_item = input_data[0]
# Skip non-text inputs (token IDs)
if isinstance(first_item, (int, list)):
verbose_proxy_logger.debug(
"OpenAI Embeddings: Input is token IDs, skipping guardrail processing"
)
return data
if not isinstance(first_item, str):
verbose_proxy_logger.warning(
"OpenAI Embeddings: Unexpected input list item type: %s",
type(first_item),
)
return data
# List of strings - apply guardrail
inputs = GenericGuardrailAPIInputs(texts=input_data) # type: ignore
if model := data.get("model"):
inputs["model"] = model
guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
inputs=inputs,
request_data=data,
input_type="request",
logging_obj=litellm_logging_obj,
)
if guardrailed_texts := guardrailed_inputs.get("texts"):
data["input"] = guardrailed_texts
verbose_proxy_logger.debug(
"OpenAI Embeddings: Applied guardrail to %d inputs",
len(guardrailed_texts),
)
return data
async def process_output_response(
self,
response: "EmbeddingResponse",
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: Optional[Any] = None,
user_api_key_dict: Optional[Any] = None,
) -> Any:
"""
Process output response - embeddings responses contain vectors, not text.
For embeddings, the output is numerical vectors, so there's typically
no text content to apply guardrails to. This method is a no-op but
is included for interface consistency.
Args:
response: Embedding response object
guardrail_to_apply: The guardrail instance to apply
litellm_logging_obj: Optional logging object
user_api_key_dict: User API key metadata
Returns:
Unmodified response (embeddings don't have text output to guard)
"""
verbose_proxy_logger.debug(
"OpenAI Embeddings: Output response processing skipped - "
"embeddings contain vectors, not text"
)
return response

195
litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py
View File

@ -22,16 +22,18 @@ from litellm.integrations.custom_guardrail import (
CustomGuardrail,
log_guardrail_information,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.types.utils import GenericGuardrailAPIInputs
from .base import OpenAIGuardrailBase
if TYPE_CHECKING:
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.llms.openai import AllMessageValues, OpenAIModerationResponse
from litellm.types.llms.openai import OpenAIModerationResponse
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
from litellm.types.utils import ModelResponse, ModelResponseStream
@ -178,170 +180,59 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail):
},
)
def _extract_user_content_from_data(self, data: Dict[str, Any]) -> Optional[str]:
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict,
input_type: Literal["request", "response"],
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
"""
Extract user content from request data, supporting both Chat Completions and Responses API.
Apply OpenAI moderation guardrail using the unified guardrail interface.
For Chat Completions: extracts from 'messages' field
For Responses API: extracts from 'input' field
This method is called by the UnifiedLLMGuardrails system for all endpoint types
(chat completions, embeddings, responses API, etc.).
Args:
inputs: GenericGuardrailAPIInputs containing texts and/or structured_messages
request_data: The original request data
input_type: Whether this is a "request" (pre-call) or "response" (post-call)
logging_obj: Optional logging object
Returns:
The extracted user content string, or None if no content found
The inputs unchanged (moderation doesn't modify content, only blocks)
Raises:
HTTPException: If content violates moderation policy
"""
# Try to get messages first (Chat Completions API)
messages: Optional[List["AllMessageValues"]] = data.get("messages")
if messages is not None:
return self.get_user_prompt(messages)
# Extract text to moderate from inputs
text_to_moderate: Optional[str] = None
# Try to get input (Responses API)
input_data = data.get("input")
if input_data is not None:
# input can be a string or a list of message-like objects
if isinstance(input_data, str):
return input_data
elif isinstance(input_data, list):
# Treat input as messages and extract user content
return self.get_user_prompt(input_data)
# Prefer structured_messages if available (has role context)
if structured_messages := inputs.get("structured_messages"):
text_to_moderate = self.get_user_prompt(structured_messages)
return None
@log_guardrail_information
async def async_pre_call_hook(
self,
user_api_key_dict: "UserAPIKeyAuth",
cache: Any,
data: Dict[str, Any],
call_type: Literal[
"completion",
"text_completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"pass_through_endpoint",
"rerank",
"mcp_call",
],
) -> Optional[Dict[str, Any]]:
"""
Pre-call hook to scan user prompts before sending to LLM.
Raises HTTPException if content should be blocked.
"""
verbose_proxy_logger.debug(
"OpenAI Moderation: Running pre-call prompt scan, on call_type: %s",
call_type,
)
# Fall back to texts
if not text_to_moderate:
if texts := inputs.get("texts"):
# Join all texts for moderation
text_to_moderate = "\n".join(texts)
# Skip moderation calls to avoid infinite recursion
if call_type == "moderation":
return data
user_prompt = self._extract_user_content_from_data(data)
if user_prompt is None:
verbose_proxy_logger.warning(
"OpenAI Moderation: not running guardrail. No messages or input in data"
)
return data
if user_prompt:
if not text_to_moderate:
verbose_proxy_logger.debug(
f"OpenAI Moderation: User prompt: {user_prompt[:100]}..." # Log first 100 chars for debugging
"OpenAI Moderation: No text content to moderate in inputs"
)
moderation_response = await self.async_make_request(
input_text=user_prompt,
)
# Check if content is flagged and raise exception if needed
self._check_moderation_result(moderation_response)
else:
verbose_proxy_logger.warning(
"OpenAI Moderation: No user prompt found"
)
return data
@log_guardrail_information
async def async_moderation_hook(
self,
data: Dict[str, Any],
user_api_key_dict: "UserAPIKeyAuth",
call_type: Literal[
"completion",
"embeddings",
"image_generation",
"moderation",
"audio_transcription",
"responses",
"mcp_call",
],
) -> Optional[Dict[str, Any]]:
"""
Moderation hook to scan user prompts during call processing.
Raises HTTPException if content should be blocked.
"""
verbose_proxy_logger.debug(
"OpenAI Moderation: Running moderation hook, on call_type: %s",
call_type,
)
return inputs
# Make moderation request
moderation_response = await self.async_make_request(input_text=text_to_moderate)
# Skip moderation calls to avoid infinite recursion
if call_type == "moderation":
return data
# Extract user content from either messages or input field
user_prompt = self._extract_user_content_from_data(data)
# Check if content is flagged and raise exception if needed
self._check_moderation_result(moderation_response)
if user_prompt is None:
verbose_proxy_logger.warning(
"OpenAI Moderation: not running guardrail. No messages or input in data"
)
return data
# Moderation doesn't modify content, just blocks - return inputs unchanged
return inputs
if user_prompt:
moderation_response = await self.async_make_request(
input_text=user_prompt,
)
# Check if content is flagged and raise exception if needed
self._check_moderation_result(moderation_response)
return data
@log_guardrail_information
async def async_post_call_hook(
self,
data: Dict[str, Any],
user_api_key_dict: "UserAPIKeyAuth",
response: "ModelResponse",
) -> "ModelResponse":
"""
Post-call hook to scan LLM responses before returning to user.
Raises HTTPException if response should be blocked.
"""
verbose_proxy_logger.debug(
"OpenAI Moderation: Running post-call response scan"
)
# Extract response text for moderation
response_text = self._extract_response_text(response)
if response_text:
verbose_proxy_logger.debug(
f"OpenAI Moderation: Response text: {response_text[:100]}..." # Log first 100 chars
)
moderation_response = await self.async_make_request(
input_text=response_text,
)
# Check if content is flagged and raise exception if needed
self._check_moderation_result(moderation_response)
return response
@log_guardrail_information
async def async_post_call_streaming_iterator_hook(

70
tests/local_testing/test_openai_moderations_hook.py
View File

@ -90,13 +90,12 @@ async def test_openai_moderation_error_raising(monkeypatch):
@pytest.mark.asyncio
async def test_openai_moderation_responses_api_input_field():
"""
Tests that OpenAI Moderation works with Responses API input field.
Tests that OpenAI Moderation works with Responses API input field via apply_guardrail.
This test verifies the fix for the issue where moderation was skipped
for Responses API because it only checked for 'messages' field but
Responses API uses 'input' field instead.
This test verifies that the unified guardrail interface (apply_guardrail) correctly
handles different input types: plain text strings, structured messages, and lists.
"""
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import patch
from litellm.types.llms.openai import (
OpenAIModerationResponse,
OpenAIModerationResult,
@ -104,6 +103,7 @@ async def test_openai_moderation_responses_api_input_field():
from litellm.proxy.guardrails.guardrail_hooks.openai.moderations import (
OpenAIModerationGuardrail,
)
from litellm.types.utils import GenericGuardrailAPIInputs
# Initialize the open-source OpenAI Moderation guardrail
openai_mod = OpenAIModerationGuardrail(
@ -112,10 +112,6 @@ async def test_openai_moderation_responses_api_input_field():
model="omni-moderation-latest",
)
_api_key = "sk-12345"
_api_key = hash_token("sk-12345")
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key)
# Mock the async_make_request to return a flagged response
mock_moderation_response = OpenAIModerationResponse(
id="modr-123",
@ -133,53 +129,47 @@ async def test_openai_moderation_responses_api_input_field():
with patch.object(
openai_mod, "async_make_request", return_value=mock_moderation_response
):
# Test 1: Responses API with input as string
# Test 1: Responses API / Embeddings with texts (string input)
try:
await openai_mod.async_moderation_hook(
data={
"model": "gpt-4o",
"input": "I want to hurt people",
},
user_api_key_dict=user_api_key_dict,
call_type="responses",
inputs = GenericGuardrailAPIInputs(texts=["I want to hurt people"])
await openai_mod.apply_guardrail(
inputs=inputs,
request_data={"model": "gpt-4o", "input": "I want to hurt people"},
input_type="request",
)
pytest.fail("Should have raised HTTPException for flagged content")
except Exception as e:
print("Got exception for string input: ", e)
print("Got exception for texts input: ", e)
assert "Violated OpenAI moderation policy" in str(e)
# Test 2: Responses API with input as list of messages
# Test 2: Responses API with structured_messages (list of message objects)
try:
await openai_mod.async_moderation_hook(
data={
"model": "gpt-4o",
"input": [
{"role": "user", "content": "I want to hurt people"}
],
},
user_api_key_dict=user_api_key_dict,
call_type="responses",
inputs = GenericGuardrailAPIInputs(
structured_messages=[{"role": "user", "content": "I want to hurt people"}]
)
await openai_mod.apply_guardrail(
inputs=inputs,
request_data={"model": "gpt-4o", "input": [{"role": "user", "content": "I want to hurt people"}]},
input_type="request",
)
pytest.fail("Should have raised HTTPException for flagged content")
except Exception as e:
print("Got exception for list input: ", e)
print("Got exception for structured_messages input: ", e)
assert "Violated OpenAI moderation policy" in str(e)
# Test 3: Verify it still works with messages field (Chat Completions)
# Test 3: Chat Completions with structured_messages
try:
await openai_mod.async_moderation_hook(
data={
"model": "gpt-4o",
"messages": [
{"role": "user", "content": "I want to hurt people"}
],
},
user_api_key_dict=user_api_key_dict,
call_type="completion",
inputs = GenericGuardrailAPIInputs(
structured_messages=[{"role": "user", "content": "I want to hurt people"}]
)
await openai_mod.apply_guardrail(
inputs=inputs,
request_data={"model": "gpt-4o", "messages": [{"role": "user", "content": "I want to hurt people"}]},
input_type="request",
)
pytest.fail("Should have raised HTTPException for flagged content")
except Exception as e:
print("Got exception for messages field: ", e)
print("Got exception for chat completions input: ", e)
assert "Violated OpenAI moderation policy" in str(e)
print("✓ All Responses API moderation tests passed!")

83
tests/test_litellm/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py Normal file
View File

@ -0,0 +1,83 @@
"""
Test OpenAI Embeddings Guardrail Translation Handler
"""
from unittest.mock import AsyncMock, MagicMock
import pytest
from litellm.llms.openai.embeddings.guardrail_translation.handler import (
OpenAIEmbeddingsHandler,
)
from litellm.types.utils import CallTypes
@pytest.mark.asyncio
async def test_embeddings_handler_string_input():
"""Test embeddings handler with single string input"""
handler = OpenAIEmbeddingsHandler()
# Mock guardrail
mock_guardrail = MagicMock()
mock_guardrail.apply_guardrail = AsyncMock(return_value={"texts": ["processed text"]})
data = {
"input": "Hello, world!",
"model": "text-embedding-3-small"
}
result = await handler.process_input_messages(
data=data,
guardrail_to_apply=mock_guardrail,
)
# Verify guardrail was called with correct inputs
mock_guardrail.apply_guardrail.assert_called_once()
call_args = mock_guardrail.apply_guardrail.call_args
assert call_args.kwargs["inputs"]["texts"] == ["Hello, world!"]
assert call_args.kwargs["inputs"]["model"] == "text-embedding-3-small"
# Verify result
assert result["input"] == "processed text"
@pytest.mark.asyncio
async def test_embeddings_handler_list_of_strings_input():
"""Test embeddings handler with list of strings input"""
handler = OpenAIEmbeddingsHandler()
# Mock guardrail
mock_guardrail = MagicMock()
mock_guardrail.apply_guardrail = AsyncMock(
return_value={"texts": ["processed text 1", "processed text 2"]}
)
data = {
"input": ["Hello, world!", "How are you?"],
"model": "text-embedding-3-small"
}
result = await handler.process_input_messages(
data=data,
guardrail_to_apply=mock_guardrail,
)
# Verify guardrail was called with correct inputs
mock_guardrail.apply_guardrail.assert_called_once()
call_args = mock_guardrail.apply_guardrail.call_args
assert call_args.kwargs["inputs"]["texts"] == ["Hello, world!", "How are you?"]
# Verify result
assert result["input"] == ["processed text 1", "processed text 2"]
def test_embeddings_guardrail_translation_mappings():
"""Test that embeddings handler is registered for correct call types"""
from litellm.llms.openai.embeddings.guardrail_translation import (
guardrail_translation_mappings,
)
assert CallTypes.embedding in guardrail_translation_mappings
assert CallTypes.aembedding in guardrail_translation_mappings
assert guardrail_translation_mappings[CallTypes.embedding] == OpenAIEmbeddingsHandler
assert guardrail_translation_mappings[CallTypes.aembedding] == OpenAIEmbeddingsHandler

108
tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py
View File

@ -83,7 +83,9 @@ async def test_openai_moderation_guardrail_adds_to_litellm_callbacks():
@pytest.mark.asyncio
async def test_openai_moderation_guardrail_safe_content():
"""Test OpenAI moderation guardrail with safe content"""
"""Test OpenAI moderation guardrail with safe content via apply_guardrail"""
from litellm.types.utils import GenericGuardrailAPIInputs
with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}):
guardrail = OpenAIModerationGuardrail(
guardrail_name="test-openai-moderation",
@ -122,28 +124,86 @@ async def test_openai_moderation_guardrail_safe_content():
)
with patch.object(guardrail, 'async_make_request', return_value=mock_response):
# Test pre-call hook with safe content
user_api_key_dict = UserAPIKeyAuth(api_key="test")
data = {
"messages": [
# Test apply_guardrail with safe content using structured_messages
inputs = GenericGuardrailAPIInputs(
structured_messages=[
{"role": "user", "content": "Hello, how are you today?"}
]
}
result = await guardrail.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=None,
data=data,
call_type="completion"
)
# Should return the original data unchanged
assert result == data
result = await guardrail.apply_guardrail(
inputs=inputs,
request_data={"messages": [{"role": "user", "content": "Hello, how are you today?"}]},
input_type="request"
)
# Should return the original inputs unchanged
assert result == inputs
@pytest.mark.asyncio
async def test_openai_moderation_guardrail_apply_guardrail():
"""Test OpenAI moderation guardrail apply_guardrail method (unified guardrail interface)"""
from litellm.types.utils import GenericGuardrailAPIInputs
with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}):
guardrail = OpenAIModerationGuardrail(
guardrail_name="test-openai-moderation",
)
# Mock safe moderation response
mock_response = OpenAIModerationResponse(
id="modr-123",
model="omni-moderation-latest",
results=[
OpenAIModerationResult(
flagged=False,
categories={
"sexual": False,
"hate": False,
"harassment": False,
"self-harm": False,
"violence": False,
},
category_scores={
"sexual": 0.001,
"hate": 0.001,
"harassment": 0.001,
"self-harm": 0.001,
"violence": 0.001,
},
category_applied_input_types={
"sexual": [],
"hate": [],
"harassment": [],
"self-harm": [],
"violence": [],
}
)
]
)
with patch.object(guardrail, 'async_make_request', return_value=mock_response):
# Test apply_guardrail with texts (embeddings-style input)
inputs = GenericGuardrailAPIInputs(
texts=["Hello, how are you?", "What is the weather?"]
)
result = await guardrail.apply_guardrail(
inputs=inputs,
request_data={},
input_type="request",
)
# Should return inputs unchanged (moderation doesn't modify, only blocks)
assert result == inputs
@pytest.mark.asyncio
async def test_openai_moderation_guardrail_harmful_content():
"""Test OpenAI moderation guardrail with harmful content"""
"""Test OpenAI moderation guardrail with harmful content via apply_guardrail"""
from litellm.types.utils import GenericGuardrailAPIInputs
with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}):
guardrail = OpenAIModerationGuardrail(
guardrail_name="test-openai-moderation",
@ -182,22 +242,20 @@ async def test_openai_moderation_guardrail_harmful_content():
)
with patch.object(guardrail, 'async_make_request', return_value=mock_response):
# Test pre-call hook with harmful content
user_api_key_dict = UserAPIKeyAuth(api_key="test")
data = {
"messages": [
# Test apply_guardrail with harmful content using structured_messages
inputs = GenericGuardrailAPIInputs(
structured_messages=[
{"role": "user", "content": "This is hateful content"}
]
}
)
# Should raise HTTPException
from fastapi import HTTPException
with pytest.raises(HTTPException) as exc_info:
await guardrail.async_pre_call_hook(
user_api_key_dict=user_api_key_dict,
cache=None,
data=data,
call_type="completion"
await guardrail.apply_guardrail(
inputs=inputs,
request_data={"messages": [{"role": "user", "content": "This is hateful content"}]},
input_type="request"
)
assert exc_info.value.status_code == 400