diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py index b178efbda59..9f4b21b439f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py @@ -1,6 +1,8 @@ import re from typing import TYPE_CHECKING, Any, Dict, List, Optional +from fastapi import HTTPException + from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_last_user_message, @@ -15,6 +17,8 @@ if TYPE_CHECKING: # Azure Content Safety APIs have a 10,000 character limit per request. AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH = 10000 +AZURE_CONTENT_SAFETY_MAX_CHUNKS_PER_TEXT = 10 +AZURE_CONTENT_SAFETY_MAX_TEXTS_PER_CALL = 10 class AzureGuardrailBase: @@ -117,6 +121,43 @@ class AzureGuardrailBase: return chunks + @staticmethod + def raise_if_text_too_long(text: str) -> None: + """ + Raise HTTPException(413) if scanning ``text`` would require more than + AZURE_CONTENT_SAFETY_MAX_CHUNKS_PER_TEXT Azure Content Safety API calls. + """ + max_length = AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH * AZURE_CONTENT_SAFETY_MAX_CHUNKS_PER_TEXT + if len(text) > max_length: + raise HTTPException( + status_code=413, + detail={ + "error": "Azure Content Safety Guardrail: input text too large", + "detection_message": ( + f"Text length {len(text)} exceeds the maximum of {max_length} " + "characters allowed per guardrail scan" + ), + }, + ) + + @staticmethod + def raise_if_too_many_texts(texts: list[str]) -> None: + """ + Raise HTTPException(413) if ``texts`` has more entries than + AZURE_CONTENT_SAFETY_MAX_TEXTS_PER_CALL. + """ + if len(texts) > AZURE_CONTENT_SAFETY_MAX_TEXTS_PER_CALL: + raise HTTPException( + status_code=413, + detail={ + "error": "Azure Content Safety Guardrail: too many texts", + "detection_message": ( + f"Received {len(texts)} texts, exceeding the maximum of " + f"{AZURE_CONTENT_SAFETY_MAX_TEXTS_PER_CALL} per guardrail scan" + ), + }, + ) + def get_user_prompt(self, messages: List["AllMessageValues"]) -> Optional[str]: """ Get the last consecutive block of messages from the user. diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py index 6dad9ae8fc6..d543ee160c9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py @@ -76,6 +76,7 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai AzurePromptShieldGuardrailResponse, ) + self.raise_if_text_too_long(user_prompt) chunks = self.split_text_by_words(user_prompt, AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH) last_response: Optional[AzurePromptShieldGuardrailResponse] = None @@ -115,7 +116,9 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai # ad-hoc /apply_guardrail endpoint may pass input_type="response". # The Azure text:shieldPrompt API accepts any text, so we scan # regardless of direction rather than silently skipping. - for text in inputs.get("texts") or []: + texts = inputs.get("texts") or [] + self.raise_if_too_many_texts(texts) + for text in texts: if text: await self.async_make_request(user_prompt=text) return inputs diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py index 184285ade27..c7e980e722b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py @@ -116,6 +116,7 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr AzureTextModerationGuardrailResponse, ) + self.raise_if_text_too_long(text) chunks = self.split_text_by_words(text, AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH) last_response: Optional[AzureTextModerationGuardrailResponse] = None @@ -155,7 +156,9 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr input_type: Literal["request", "response"], logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: - for text in inputs.get("texts") or []: + texts = inputs.get("texts") or [] + self.raise_if_too_many_texts(texts) + for text in texts: if text: await self.async_make_request(text=text) return inputs diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py index 72d4dc2ecb5..a7d88e6c46d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py @@ -288,6 +288,100 @@ async def test_apply_guardrail_empty_texts_makes_no_request(): mock_post.assert_not_called() +@pytest.mark.asyncio +async def test_apply_guardrail_too_many_texts_rejected(): + """A texts list larger than the per-call cap must be rejected before any + Azure API call is made, so a caller cannot force unbounded fan-out by + submitting many texts in a single apply_guardrail call.""" + from litellm.proxy.guardrails.guardrail_hooks.azure.base import ( + AZURE_CONTENT_SAFETY_MAX_TEXTS_PER_CALL, + ) + + guardrail = _make_prompt_shield_guardrail() + texts = ["safe text"] * (AZURE_CONTENT_SAFETY_MAX_TEXTS_PER_CALL + 1) + + with patch.object(guardrail.async_handler, "post") as mock_post: + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": texts}, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 413 + mock_post.assert_not_called() + + +@pytest.mark.asyncio +async def test_apply_guardrail_max_texts_per_call_allowed(): + """Exactly the per-call cap must still be scanned normally.""" + from litellm.proxy.guardrails.guardrail_hooks.azure.base import ( + AZURE_CONTENT_SAFETY_MAX_TEXTS_PER_CALL, + ) + + guardrail = _make_prompt_shield_guardrail() + texts = ["safe text"] * AZURE_CONTENT_SAFETY_MAX_TEXTS_PER_CALL + + with patch.object( + guardrail.async_handler, + "post", + return_value=_mock_shield_response(False), + ) as mock_post: + result = await guardrail.apply_guardrail( + inputs={"texts": texts}, + request_data={}, + input_type="request", + ) + + assert result == {"texts": texts} + assert mock_post.call_count == AZURE_CONTENT_SAFETY_MAX_TEXTS_PER_CALL + + +@pytest.mark.asyncio +async def test_async_make_request_text_too_long_rejected(): + """A single text large enough to require more Azure calls than the + per-text chunk cap must be rejected before any request is made, so a + caller cannot force unbounded fan-out with one oversized prompt.""" + from litellm.proxy.guardrails.guardrail_hooks.azure.base import ( + AZURE_CONTENT_SAFETY_MAX_CHUNKS_PER_TEXT, + AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH, + ) + + guardrail = _make_prompt_shield_guardrail() + max_length = AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH * AZURE_CONTENT_SAFETY_MAX_CHUNKS_PER_TEXT + oversized_text = "a" * (max_length + 1) + + with patch.object(guardrail.async_handler, "post") as mock_post: + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_make_request(user_prompt=oversized_text) + + assert exc_info.value.status_code == 413 + mock_post.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_make_request_max_text_length_allowed(): + """Text exactly at the max length must still be scanned, split into the + expected number of chunks.""" + from litellm.proxy.guardrails.guardrail_hooks.azure.base import ( + AZURE_CONTENT_SAFETY_MAX_CHUNKS_PER_TEXT, + AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH, + ) + + guardrail = _make_prompt_shield_guardrail() + max_length = AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH * AZURE_CONTENT_SAFETY_MAX_CHUNKS_PER_TEXT + text_at_limit = "a" * max_length + + with patch.object( + guardrail.async_handler, + "post", + return_value=_mock_shield_response(False), + ) as mock_post: + await guardrail.async_make_request(user_prompt=text_at_limit) + + assert mock_post.call_count == AZURE_CONTENT_SAFETY_MAX_CHUNKS_PER_TEXT + + def test_split_text_by_words(): """Test the word-based text splitting functionality.""" guardrail = AzureContentSafetyPromptShieldGuardrail( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py index 46f8a7279d3..f66b74e4930 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py @@ -424,6 +424,116 @@ async def test_apply_guardrail_empty_texts_makes_no_request(): mock_post.assert_not_called() +@pytest.mark.asyncio +async def test_apply_guardrail_too_many_texts_rejected(): + """A texts list larger than the per-call cap must be rejected before any + Azure API call is made, so a caller cannot force unbounded fan-out by + submitting many texts in a single apply_guardrail call.""" + from litellm.proxy.guardrails.guardrail_hooks.azure.base import ( + AZURE_CONTENT_SAFETY_MAX_TEXTS_PER_CALL, + ) + + guardrail = AzureContentSafetyTextModerationGuardrail( + guardrail_name="azure_text_moderation", + api_key="azure_text_moderation_api_key", + api_base="azure_text_moderation_api_base", + ) + texts = ["safe text"] * (AZURE_CONTENT_SAFETY_MAX_TEXTS_PER_CALL + 1) + + with patch.object(guardrail.async_handler, "post") as mock_post: + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": texts}, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 413 + mock_post.assert_not_called() + + +@pytest.mark.asyncio +async def test_apply_guardrail_max_texts_per_call_allowed(): + """Exactly the per-call cap must still be scanned normally.""" + from litellm.proxy.guardrails.guardrail_hooks.azure.base import ( + AZURE_CONTENT_SAFETY_MAX_TEXTS_PER_CALL, + ) + + guardrail = AzureContentSafetyTextModerationGuardrail( + guardrail_name="azure_text_moderation", + api_key="azure_text_moderation_api_key", + api_base="azure_text_moderation_api_base", + ) + texts = ["safe text"] * AZURE_CONTENT_SAFETY_MAX_TEXTS_PER_CALL + + with patch.object( + guardrail.async_handler, + "post", + return_value=_mock_moderation_response({"Hate": 0}), + ) as mock_post: + result = await guardrail.apply_guardrail( + inputs={"texts": texts}, + request_data={}, + input_type="request", + ) + + assert result == {"texts": texts} + assert mock_post.call_count == AZURE_CONTENT_SAFETY_MAX_TEXTS_PER_CALL + + +@pytest.mark.asyncio +async def test_async_make_request_text_too_long_rejected(): + """A single text large enough to require more Azure calls than the + per-text chunk cap must be rejected before any request is made, so a + caller cannot force unbounded fan-out with one oversized text.""" + from litellm.proxy.guardrails.guardrail_hooks.azure.base import ( + AZURE_CONTENT_SAFETY_MAX_CHUNKS_PER_TEXT, + AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH, + ) + + guardrail = AzureContentSafetyTextModerationGuardrail( + guardrail_name="azure_text_moderation", + api_key="azure_text_moderation_api_key", + api_base="azure_text_moderation_api_base", + ) + max_length = AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH * AZURE_CONTENT_SAFETY_MAX_CHUNKS_PER_TEXT + oversized_text = "a" * (max_length + 1) + + with patch.object(guardrail.async_handler, "post") as mock_post: + with pytest.raises(HTTPException) as exc_info: + await guardrail.async_make_request(text=oversized_text) + + assert exc_info.value.status_code == 413 + mock_post.assert_not_called() + + +@pytest.mark.asyncio +async def test_async_make_request_max_text_length_allowed(): + """Text exactly at the max length must still be scanned, split into the + expected number of chunks.""" + from litellm.proxy.guardrails.guardrail_hooks.azure.base import ( + AZURE_CONTENT_SAFETY_MAX_CHUNKS_PER_TEXT, + AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH, + ) + + guardrail = AzureContentSafetyTextModerationGuardrail( + guardrail_name="azure_text_moderation", + api_key="azure_text_moderation_api_key", + api_base="azure_text_moderation_api_base", + ) + max_length = AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH * AZURE_CONTENT_SAFETY_MAX_CHUNKS_PER_TEXT + text_at_limit = "a" * max_length + + with patch.object( + guardrail.async_handler, + "post", + return_value=_mock_moderation_response({"Hate": 0}), + ) as mock_post: + await guardrail.async_make_request(text=text_at_limit) + + assert mock_post.call_count == AZURE_CONTENT_SAFETY_MAX_CHUNKS_PER_TEXT + + def test_split_text_by_words(): """Test the word-based text splitting functionality.""" guardrail = AzureContentSafetyTextModerationGuardrail(