From 887a907e4224dee7a69079fbb3ee5d9c9c5f1b48 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 5 Feb 2026 14:40:15 -0800 Subject: [PATCH] [Fix] Guardrails API - Ensure OpenAI Moderations Guard works with OpenAI Embeddings (#20523) * 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 --- .../guardrail_translation/__init__.py | 13 ++ .../guardrail_translation/handler.py | 179 ++++++++++++++++ .../guardrail_hooks/openai/moderations.py | 195 ++++-------------- .../test_openai_moderations_hook.py | 70 +++---- .../test_embeddings_guardrail_handler.py | 83 ++++++++ .../openai/test_moderations.py | 108 +++++++--- 6 files changed, 431 insertions(+), 217 deletions(-) create mode 100644 litellm/llms/openai/embeddings/guardrail_translation/__init__.py create mode 100644 litellm/llms/openai/embeddings/guardrail_translation/handler.py create mode 100644 tests/test_litellm/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py diff --git a/litellm/llms/openai/embeddings/guardrail_translation/__init__.py b/litellm/llms/openai/embeddings/guardrail_translation/__init__.py new file mode 100644 index 0000000000..a60662282c --- /dev/null +++ b/litellm/llms/openai/embeddings/guardrail_translation/__init__.py @@ -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"] diff --git a/litellm/llms/openai/embeddings/guardrail_translation/handler.py b/litellm/llms/openai/embeddings/guardrail_translation/handler.py new file mode 100644 index 0000000000..7458020e10 --- /dev/null +++ b/litellm/llms/openai/embeddings/guardrail_translation/handler.py @@ -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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index 852cf01bc0..030b603681 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -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( diff --git a/tests/local_testing/test_openai_moderations_hook.py b/tests/local_testing/test_openai_moderations_hook.py index 3acd36f32f..3632976d03 100644 --- a/tests/local_testing/test_openai_moderations_hook.py +++ b/tests/local_testing/test_openai_moderations_hook.py @@ -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!") diff --git a/tests/test_litellm/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py b/tests/test_litellm/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py new file mode 100644 index 0000000000..c4afa7c6f1 --- /dev/null +++ b/tests/test_litellm/llms/openai/embeddings/guardrail_translation/test_embeddings_guardrail_handler.py @@ -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 diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index 8957b534ea..cebba2ff5e 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -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