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 <igal@pillar.security>

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:

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,

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,

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,

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,

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,

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,

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,

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,

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,

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

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:

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

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"""