scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-07 07:58:09 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 150

fix(azure): cap Azure Content Safety fan-out per guardrail call

Browse source 
apply_guardrail and async_make_request accepted unbounded text, so an
  authenticated caller could submit a very long prompt or many texts and
  force a single inbound request into an unbounded number of Azure API
  calls, exhausting shared Azure quota. Cap the per-text chunk count and
  the number of texts scanned per call, rejecting with 413 before any
  Azure request is made once the caps are exceeded.
This commit is contained in:
Damien Smrt 2026-08-06 07:21:54 -06:00
parent dbc3dbd920
commit aa6b020ff6
No known key found for this signature in database
GPG key ID: 479BC509806B469B
5 changed files with 253 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
litellm/proxy/guardrails/guardrail_hooks/azure/base.py
View file

@ -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.

5
litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py
View file

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

5
litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py
View file

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

94
tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py
View file

@ -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(

110
tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py
View file

@ -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(