From f43757f71b68ff3a014d5d0781d4ae317b9ca41d Mon Sep 17 00:00:00 2001 From: jquinter Date: Sat, 24 Jan 2026 01:48:42 -0300 Subject: [PATCH] Feature/guardrail model argument (#19619) * [Feat] Add model parameter to Generic Guardrail API Add model information to guardrail requests, allowing guardrails to make model-specific security decisions. Changes: - Add `model` field to GenericGuardrailAPIInputs TypedDict - Add `model` field to GenericGuardrailAPIRequest Pydantic model - Update OpenAI and Anthropic handlers to pass model from request/response - Add unit tests for model parameter handling * [Feat] Add model parameter to all guardrail_translation handlers Extend model parameter support to all guardrail handlers for consistent implementation across all endpoint types: - OpenAI Responses API (input/output + streaming) - OpenAI Image Generation (input only) - OpenAI Text Completion (input/output) - OpenAI Text-to-Speech (input only) - OpenAI Audio Transcription (output only) - Cohere Rerank (input only) - Pass-through Endpoints (input/output) - MCP Server (input only) This addresses the review feedback requesting consistent model parameter handling across all guardrail_translation/handler.py files. --------- Co-authored-by: Igal Boxerman --- .../chat/guardrail_translation/handler.py | 14 ++++- .../rerank/guardrail_translation/handler.py | 8 ++- .../chat/guardrail_translation/handler.py | 14 +++++ .../guardrail_translation/handler.py | 21 ++++++- .../guardrail_translation/handler.py | 8 ++- .../guardrail_translation/handler.py | 43 +++++++++++--- .../speech/guardrail_translation/handler.py | 8 ++- .../guardrail_translation/handler.py | 7 ++- .../guardrail_translation/handler.py | 15 ++++- .../guardrail_translation/handler.py | 7 ++- .../generic_guardrail_api.py | 2 + .../guardrail_hooks/generic_guardrail_api.py | 17 +++--- litellm/types/utils.py | 1 + .../test_generic_guardrail_api.py | 56 +++++++++++++++++++ 14 files changed, 195 insertions(+), 26 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 9d50cc4d92..71d74121a3 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -110,6 +110,10 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["tools"] = tools_to_check if structured_messages: inputs["structured_messages"] = structured_messages + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=data, @@ -309,6 +313,14 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["images"] = images_to_check if tool_calls_to_check: inputs["tool_calls"] = tool_calls_to_check + # Include model information from the response if available + response_model = None + if isinstance(response, dict): + response_model = response.get("model") + elif hasattr(response, "model"): + response_model = getattr(response, "model", None) + if response_model: + inputs["model"] = response_model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -552,7 +564,7 @@ class AnthropicMessagesHandler(BaseTranslation): response_content = response.get("content", []) else: response_content = getattr(response, "content", None) or [] - + if not response_content: return False for content_block in response_content: diff --git a/litellm/llms/cohere/rerank/guardrail_translation/handler.py b/litellm/llms/cohere/rerank/guardrail_translation/handler.py index 6893a5991c..b8133c59f7 100644 --- a/litellm/llms/cohere/rerank/guardrail_translation/handler.py +++ b/litellm/llms/cohere/rerank/guardrail_translation/handler.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional 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 @@ -49,8 +50,13 @@ class CohereRerankHandler(BaseTranslation): # Process query only query = data.get("query") if query is not None and isinstance(query, str): + inputs = GenericGuardrailAPIInputs(texts=[query]) + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [query]}, + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index d0ed3f165c..fb00aa28f4 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -87,6 +87,10 @@ class OpenAIChatCompletionsHandler(BaseTranslation): tools = data.get("tools") if tools: inputs["tools"] = tools + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -297,6 +301,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs["images"] = images_to_check if tool_calls_to_check: inputs["tool_calls"] = tool_calls_to_check # type: ignore + # Include model information from the response if available + if hasattr(response, "model") and response.model: + inputs["model"] = response.model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -417,6 +424,13 @@ class OpenAIChatCompletionsHandler(BaseTranslation): inputs = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: inputs["images"] = images_to_check + # Include model information from the first response if available + if ( + responses_so_far + and hasattr(responses_so_far[0], "model") + and responses_so_far[0].model + ): + inputs["model"] = responses_so_far[0].model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=request_data, diff --git a/litellm/llms/openai/completion/guardrail_translation/handler.py b/litellm/llms/openai/completion/guardrail_translation/handler.py index 73d08cfead..1f8c6159da 100644 --- a/litellm/llms/openai/completion/guardrail_translation/handler.py +++ b/litellm/llms/openai/completion/guardrail_translation/handler.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional 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 @@ -53,8 +54,13 @@ class OpenAITextCompletionHandler(BaseTranslation): if isinstance(prompt, str): # Single string prompt + inputs = GenericGuardrailAPIInputs(texts=[prompt]) + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [prompt]}, + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, @@ -80,8 +86,13 @@ class OpenAITextCompletionHandler(BaseTranslation): text_indices.append(idx) if texts_to_check: + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": texts_to_check}, + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, @@ -154,8 +165,12 @@ class OpenAITextCompletionHandler(BaseTranslation): if user_metadata: request_data["litellm_metadata"] = user_metadata + inputs = GenericGuardrailAPIInputs(texts=texts_to_check) + # Include model information from the response if available + if hasattr(response, "model") and response.model: + inputs["model"] = response.model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": texts_to_check}, + inputs=inputs, request_data=request_data, input_type="response", logging_obj=litellm_logging_obj, diff --git a/litellm/llms/openai/image_generation/guardrail_translation/handler.py b/litellm/llms/openai/image_generation/guardrail_translation/handler.py index 842a64b187..e6340ba470 100644 --- a/litellm/llms/openai/image_generation/guardrail_translation/handler.py +++ b/litellm/llms/openai/image_generation/guardrail_translation/handler.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional 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 @@ -52,8 +53,13 @@ class OpenAIImageGenerationHandler(BaseTranslation): # Apply guardrail to the prompt if isinstance(prompt, str): + inputs = GenericGuardrailAPIInputs(texts=[prompt]) + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [prompt]}, + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 9b8f15c762..d943662f9e 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -105,6 +105,10 @@ class OpenAIResponsesHandler(BaseTranslation): inputs["tools"] = tools_to_check if structured_messages: inputs["structured_messages"] = structured_messages # type: ignore + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -150,6 +154,10 @@ class OpenAIResponsesHandler(BaseTranslation): inputs["tools"] = tools_to_check if structured_messages: inputs["structured_messages"] = structured_messages # type: ignore + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=data, @@ -344,6 +352,14 @@ class OpenAIResponsesHandler(BaseTranslation): inputs["images"] = images_to_check if tool_calls_to_check: inputs["tool_calls"] = tool_calls_to_check + # Include model information from the response if available + response_model = None + if isinstance(response, dict): + response_model = response.get("model") + elif hasattr(response, "model"): + response_model = getattr(response, "model", None) + if response_model: + inputs["model"] = response_model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=inputs, @@ -388,12 +404,15 @@ class OpenAIResponsesHandler(BaseTranslation): tool_calls = model_response_stream.choices[0].delta.tool_calls if tool_calls: + inputs = GenericGuardrailAPIInputs() + inputs["tool_calls"] = cast( + List[ChatCompletionToolCallChunk], tool_calls + ) + # Include model information if available + if hasattr(model_response_stream, "model") and model_response_stream.model: + inputs["model"] = model_response_stream.model _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={ - "tool_calls": cast( - List[ChatCompletionToolCallChunk], tool_calls - ) - }, + inputs=inputs, request_data={}, input_type="response", logging_obj=litellm_logging_obj, @@ -417,7 +436,11 @@ class OpenAIResponsesHandler(BaseTranslation): guardrail_inputs["tool_calls"] = cast( List[ChatCompletionToolCallChunk], tool_calls ) - if tool_calls: + # Include model information from the response if available + response_model = final_chunk.get("response", {}).get("model") + if response_model: + guardrail_inputs["model"] = response_model + if tool_calls or text: _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( inputs=guardrail_inputs, request_data={}, @@ -429,8 +452,14 @@ class OpenAIResponsesHandler(BaseTranslation): # tool_calls = model_response_stream.choices[0].tool_calls # convert openai response to model response string_so_far = self.get_streaming_string_so_far(responses_so_far) + inputs = GenericGuardrailAPIInputs(texts=[string_so_far]) + # Try to get model from the final chunk if available + if isinstance(final_chunk, dict): + response_model = final_chunk.get("response", {}).get("model") if isinstance(final_chunk.get("response"), dict) else None + if response_model: + inputs["model"] = response_model _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [string_so_far]}, + inputs=inputs, request_data={}, input_type="response", logging_obj=litellm_logging_obj, diff --git a/litellm/llms/openai/speech/guardrail_translation/handler.py b/litellm/llms/openai/speech/guardrail_translation/handler.py index 4c2f71477b..e6796fbac2 100644 --- a/litellm/llms/openai/speech/guardrail_translation/handler.py +++ b/litellm/llms/openai/speech/guardrail_translation/handler.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional 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 @@ -50,8 +51,13 @@ class OpenAITextToSpeechHandler(BaseTranslation): return data if isinstance(input_text, str): + inputs = GenericGuardrailAPIInputs(texts=[input_text]) + # Include model information if available (voice model) + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [input_text]}, + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, diff --git a/litellm/llms/openai/transcriptions/guardrail_translation/handler.py b/litellm/llms/openai/transcriptions/guardrail_translation/handler.py index ac416f42c8..3d76a21c38 100644 --- a/litellm/llms/openai/transcriptions/guardrail_translation/handler.py +++ b/litellm/llms/openai/transcriptions/guardrail_translation/handler.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Any, Optional 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 @@ -88,8 +89,12 @@ class OpenAIAudioTranscriptionHandler(BaseTranslation): if user_metadata: request_data["litellm_metadata"] = user_metadata + inputs = GenericGuardrailAPIInputs(texts=[original_text]) + # Include model information from the response if available + if hasattr(response, "model") and response.model: + inputs["model"] = response.model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [original_text]}, + inputs=inputs, request_data=request_data, input_type="response", logging_obj=litellm_logging_obj, diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index c0979e37e6..40433d5341 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -11,6 +11,7 @@ from typing import TYPE_CHECKING, Any, List, Optional from litellm._logging import verbose_proxy_logger from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.proxy._types import PassThroughGuardrailSettings +from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.integrations.custom_guardrail import CustomGuardrail @@ -118,8 +119,13 @@ class PassThroughEndpointHandler(BaseTranslation): return data # Apply guardrail (pass-through doesn't modify the text, just checks it) + inputs = GenericGuardrailAPIInputs(texts=[text_to_check]) + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [text_to_check]}, + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, @@ -178,8 +184,13 @@ class PassThroughEndpointHandler(BaseTranslation): request_data["litellm_metadata"] = user_metadata # Apply guardrail (pass-through doesn't modify the text, just checks it) + inputs = GenericGuardrailAPIInputs(texts=[text_to_check]) + # Include model information from the response if available + response_model = response.get("model") if isinstance(response, dict) else None + if response_model: + inputs["model"] = response_model _guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs={"texts": [text_to_check]}, + inputs=inputs, request_data=request_data, input_type="response", logging_obj=litellm_logging_obj, diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py index 8d6d236b88..4d53ae7059 100644 --- a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py +++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py @@ -46,8 +46,13 @@ class MCPGuardrailTranslationHandler(BaseTranslation): ) return data + inputs = GenericGuardrailAPIInputs(texts=[content]) + # Include model information if available + model = data.get("model") + if model: + inputs["model"] = model guardrailed_inputs = await guardrail_to_apply.apply_guardrail( - inputs=GenericGuardrailAPIInputs(texts=[content]), + inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index d5a6266146..b37074e25e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -185,6 +185,7 @@ class GenericGuardrailAPI(CustomGuardrail): tools = inputs.get("tools") structured_messages = inputs.get("structured_messages") tool_calls = inputs.get("tool_calls") + model = inputs.get("model") # Use provided request_data or create an empty dict if request_data is None: @@ -215,6 +216,7 @@ class GenericGuardrailAPI(CustomGuardrail): tool_calls=tool_calls, additional_provider_specific_params=additional_params, input_type=input_type, + model=model, ) # Prepare headers diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py index cbca58e651..96d78cf882 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/generic_guardrail_api.py @@ -51,19 +51,20 @@ class GenericGuardrailAPIRequest(BaseModel): """Request model for the Generic Guardrail API""" input_type: Literal["request", "response"] - litellm_call_id: Optional[str] # the call id of the individual LLM call + litellm_call_id: Optional[str] = None # the call id of the individual LLM call litellm_trace_id: Optional[ str - ] # the trace id of the LLM call - useful if there are multiple LLM calls for the same conversation - structured_messages: Optional[List[AllMessageValues]] - images: Optional[List[str]] - tools: Optional[List[ChatCompletionToolParam]] - texts: Optional[List[str]] + ] = None # the trace id of the LLM call - useful if there are multiple LLM calls for the same conversation + structured_messages: Optional[List[AllMessageValues]] = None + images: Optional[List[str]] = None + tools: Optional[List[ChatCompletionToolParam]] = None + texts: Optional[List[str]] = None request_data: GenericGuardrailAPIMetadata - additional_provider_specific_params: Optional[Dict[str, Any]] + additional_provider_specific_params: Optional[Dict[str, Any]] = None tool_calls: Optional[ Union[List[ChatCompletionToolCallChunk], List[ChatCompletionMessageToolCall]] - ] + ] = None + model: Optional[str] = None # the model being used for the LLM call class GenericGuardrailAPIResponse: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index cd797dd1e5..0ec79e90a1 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3428,3 +3428,4 @@ class GenericGuardrailAPIInputs(TypedDict, total=False): structured_messages: List[ AllMessageValues ] # structured messages sent to the LLM - indicates if text is from system or user + model: Optional[str] # the model being used for the LLM call diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index 61d44e46da..5c03914192 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -549,6 +549,62 @@ class TestAdditionalParams: ) +class TestModelParameter: + """Test model parameter handling in guardrail requests""" + + @pytest.mark.asyncio + async def test_model_passed_from_inputs( + self, generic_guardrail, mock_request_data_input + ): + """Test that model is passed to the API when provided in inputs""" + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "NONE", + "texts": ["test"], + } + mock_response.raise_for_status = MagicMock() + + with patch.object( + generic_guardrail.async_handler, "post", return_value=mock_response + ) as mock_post: + await generic_guardrail.apply_guardrail( + inputs={"texts": ["test"], "model": "gpt-4"}, + request_data=mock_request_data_input, + input_type="request", + ) + + # Verify API was called with model + call_args = mock_post.call_args + json_payload = call_args.kwargs["json"] + assert json_payload["model"] == "gpt-4" + + @pytest.mark.asyncio + async def test_model_none_when_not_provided( + self, generic_guardrail, mock_request_data_input + ): + """Test that model is None when not provided in inputs""" + mock_response = MagicMock() + mock_response.json.return_value = { + "action": "NONE", + "texts": ["test"], + } + mock_response.raise_for_status = MagicMock() + + with patch.object( + generic_guardrail.async_handler, "post", return_value=mock_response + ) as mock_post: + await generic_guardrail.apply_guardrail( + inputs={"texts": ["test"]}, # No model in inputs + request_data=mock_request_data_input, + input_type="request", + ) + + # Verify API was called with model=None + call_args = mock_post.call_args + json_payload = call_args.kwargs["json"] + assert json_payload["model"] is None + + class TestErrorHandling: """Test error handling scenarios"""