scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

fix(guardrails): keep chat headers as before and check escaped systemone entries

Browse source 
Vendor guardrails with their own hooks list themselves in the applied
guardrails header as before on chat, Responses and Messages, and only
when they scanned text on extraction and decision routes

Model Armor lists itself only once its API answered, so a call it let
through with fail_on_error false is no longer listed

A systemone entry whose JSON escaping splits a value is scanned in its
decoded form too, and a rewrite where the two forms disagree is rejected
This commit is contained in:
mateo-berri 2026-10-29 14:29:56 -07:00
parent 48643021d4
commit 744dadabd8
20 changed files with 299 additions and 68 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/integrations/custom_guardrail.py
View file

@ -14,6 +14,7 @@ from litellm._internal_context import with_service_target
from litellm._logging import verbose_logger
from litellm.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
get_or_create_metadata_bucket,
@ -85,6 +86,11 @@ APPLY_GUARDRAIL_ONLY_CALL_TYPES: Final[frozenset[str]] = frozenset(
)
def is_apply_guardrail_only_route(route: str | None) -> bool:
call_types: Final = get_call_types_for_route(route) if route else None
return bool(call_types) and call_types[0] in APPLY_GUARDRAIL_ONLY_CALL_TYPES
def runs_through_apply_guardrail(guardrail: "CustomGuardrail", call_type: str | None) -> bool:
return "apply_guardrail" in type(guardrail).__dict__ and (
not guardrail.use_native_lifecycle_hooks or call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES
@ -245,6 +251,7 @@ class CustomGuardrail(CustomLogger):
use_native_lifecycle_hooks: ClassVar[bool] = False
records_own_guardrail_information: ClassVar[bool] = False
adds_itself_to_applied_guardrails: ClassVar[bool] = False
logging_only_scope: LoggingOnlyScope | None
logging_only_continue_on_input_failure: bool

27
litellm/llms/openai/decisions/guardrail_translation/handler.py
View file

@ -84,12 +84,18 @@ def _entry_text(key: str, value: str | float | None) -> str:
return json.dumps({key: value}, ensure_ascii=False)[1:-1]
def _escaped_texts(text: str, key: str, item: JsonValue) -> tuple[str, ...]:
return tuple(form for form in (key, item) if isinstance(form, str) and form not in text)
def _entry_text_forms(key: str, item: JsonValue) -> Iterator[str]:
if isinstance(item, (list, dict)):
yield key
yield from _json_text_forms(item)
else:
yield _entry_text(key, item)
text: Final = _entry_text(key, item)
yield text
yield from _escaped_texts(text, key, item)
def _json_text_forms(value: JsonValue) -> Iterator[str]:
@ -163,6 +169,10 @@ def _rewritten_entry(rewritten: str, guardrail_name: str | None) -> tuple[str, J
return _KEY_ADAPTER.validate_json(key.group(1)), _rewritten_value(rewritten[key.end() :])
def _decoded_text(text: str, form: str, guardrailed: Mapping[str, str]) -> str:
return form if form in text else guardrailed.get(form, form)
def _guardrailed_entry(
key: str, item: JsonValue, guardrailed: Mapping[str, str], guardrail_name: str | None
) -> tuple[str, JsonValue]:
@ -170,7 +180,16 @@ def _guardrailed_entry(
return guardrailed.get(key, key), _guardrailed_json(item, guardrailed, guardrail_name)
text: Final = _entry_text(key, item)
rewritten: Final = guardrailed.get(text, text)
return (key, item) if rewritten == text else _rewritten_entry(rewritten, guardrail_name)
decoded: Final = (
_decoded_text(text, key, guardrailed),
_decoded_text(text, item, guardrailed) if isinstance(item, str) else item,
)
if rewritten == text:
return decoded
entry: Final = _rewritten_entry(rewritten, guardrail_name)
if decoded not in ((key, item), entry):
raise unappliable_request_rewrite(guardrail_name)
return entry
def _instruction_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]:
@ -260,7 +279,9 @@ async def _guardrailed_texts(
input_type="request",
logging_obj=litellm_logging_obj,
)
if not guardrail_to_apply.records_own_guardrail_information:
if not (
guardrail_to_apply.records_own_guardrail_information or guardrail_to_apply.adds_itself_to_applied_guardrails
):
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict
)

17
litellm/proxy/guardrails/guardrail_hooks/aim/aim.py
View file

@ -17,7 +17,7 @@ from websockets.asyncio.client import ClientConnection, connect
from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm._version import version as litellm_version
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_guardrail import APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -155,7 +155,9 @@ class AimGuardrail(CustomGuardrail):
if is_non_conversational_call_type(call_type) and not self.inspect_embeddings:
verbose_proxy_logger.debug("Aim: skipping non-conversational call type %s", call_type)
return data
return await self.call_aim_guardrail(data, hook="pre_call", key_alias=user_api_key_dict.key_alias)
return await self.call_aim_guardrail(
data, hook="pre_call", key_alias=user_api_key_dict.key_alias, call_type=call_type
)
async def async_moderation_hook(
self,
@ -168,14 +170,19 @@ class AimGuardrail(CustomGuardrail):
verbose_proxy_logger.debug("Aim: skipping non-conversational call type %s", call_type)
return data
await self.call_aim_guardrail(data, hook="moderation", key_alias=user_api_key_dict.key_alias)
await self.call_aim_guardrail(
data, hook="moderation", key_alias=user_api_key_dict.key_alias, call_type=call_type
)
return data
async def call_aim_guardrail(self, data: dict, hook: str, key_alias: str | None) -> dict:
async def call_aim_guardrail(
self, data: dict, hook: str, key_alias: str | None, call_type: CallTypesLiteral | None = None
) -> dict:
messages: Final = self._build_aim_inspection_messages(data) # pyright: ignore[reportUnknownMemberType] # hook request data is an untyped dict
if not messages:
return data # pyright: ignore[reportUnknownVariableType] # hook request data is an untyped dict
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
user_email: Final = data.get("metadata", {}).get("headers", {}).get("x-aim-user-email")
call_id: Final = data.get("litellm_call_id")
headers: Final = self._build_aim_headers(

8
litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py
View file

@ -22,7 +22,7 @@ from websockets.exceptions import ConnectionClosed
from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm._version import version as litellm_version
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_guardrail import APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail
from litellm.litellm_core_utils.core_helpers import iter_items
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -176,6 +176,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
hook="pre_call",
key_alias=user_api_key_dict.key_alias,
user_email=self._resolve_cato_user_email(user_api_key_dict),
call_type=call_type,
)
async def async_moderation_hook(
@ -193,6 +194,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
hook="moderation",
key_alias=user_api_key_dict.key_alias,
user_email=self._resolve_cato_user_email(user_api_key_dict),
call_type=call_type,
)
@classmethod
@ -297,11 +299,13 @@ class CatoNetworksGuardrail(CustomGuardrail):
hook: str,
key_alias: str | None,
user_email: str | None = None,
call_type: CallTypesLiteral | None = None,
) -> dict:
messages: Final = self._inspection_messages(data) # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType] # hook request data is an untyped dict
if not messages:
return data # pyright: ignore[reportUnknownVariableType] # hook request data is an untyped dict
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
call_id: Final = data.get("litellm_call_id")
headers: Final = self._build_cato_headers(
hook=hook,

11
litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py
View file

@ -15,7 +15,11 @@ import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_guardrail import (
APPLY_GUARDRAIL_ONLY_CALL_TYPES,
CustomGuardrail,
is_apply_guardrail_only_route,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -328,6 +332,7 @@ class DynamoAIGuardrails(CustomGuardrail):
error_message: Final = self._create_error_message(processed_result)
raise ValueError(error_message)
if _messages or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
return data
@ -372,6 +377,7 @@ class DynamoAIGuardrails(CustomGuardrail):
error_message: Final = self._create_error_message(processed_result)
raise ValueError(error_message)
if _messages or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
return data
@ -439,6 +445,9 @@ class DynamoAIGuardrails(CustomGuardrail):
error_message: Final = self._create_error_message(processed_result)
raise ValueError(error_message)
if isinstance(response, litellm.ModelResponse) or not is_apply_guardrail_only_route(
user_api_key_dict.request_route
):
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
async def async_post_call_streaming_iterator_hook(

11
litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py
View file

@ -15,7 +15,11 @@ import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_guardrail import (
APPLY_GUARDRAIL_ONLY_CALL_TYPES,
CustomGuardrail,
is_apply_guardrail_only_route,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -476,6 +480,7 @@ class IBMGuardrailDetector(CustomGuardrail):
error_message = self._create_error_message_orchestrator(orchestrator_result)
raise ValueError(error_message)
if contents_to_check or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
return data
@ -544,6 +549,7 @@ class IBMGuardrailDetector(CustomGuardrail):
error_message = self._create_error_message_orchestrator(orchestrator_result)
raise ValueError(error_message)
if contents_to_check or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
return data
@ -639,6 +645,9 @@ class IBMGuardrailDetector(CustomGuardrail):
error_message = self._create_error_message_orchestrator(orchestrator_result)
raise ValueError(error_message)
if isinstance(response, litellm.ModelResponse) or not is_apply_guardrail_only_route(
user_api_key_dict.request_route
):
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
async def async_post_call_streaming_iterator_hook(

4
litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py
View file

@ -19,6 +19,7 @@ from fastapi import HTTPException
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
APPLY_GUARDRAIL_ONLY_CALL_TYPES,
CustomGuardrail,
log_guardrail_information,
)
@ -292,7 +293,8 @@ class lakeraAI_Moderation(CustomGuardrail):
}
"""
self._check_response_flagged(response=response.json())
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
@log_guardrail_information
async def async_pre_call_hook(

10
litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py
View file

@ -34,6 +34,7 @@ from litellm import DualCache
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
APPLY_GUARDRAIL_ONLY_CALL_TYPES,
GUARDRAIL_SESSIONS_TARGET,
CustomGuardrail,
log_guardrail_information,
@ -207,7 +208,7 @@ class LassoGuardrail(CustomGuardrail):
# The conversation_id is being stored in the cache so it can be used by the post_call hook
self._get_or_generate_conversation_id(data, global_cache)
return await self._run_lasso_guardrail(data, global_cache, message_type="PROMPT")
return await self._run_lasso_guardrail(data, global_cache, message_type="PROMPT", call_type=call_type)
@log_guardrail_information
async def async_moderation_hook(
@ -235,7 +236,7 @@ class LassoGuardrail(CustomGuardrail):
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
return data
return await self._run_lasso_guardrail(data, cache, message_type="PROMPT")
return await self._run_lasso_guardrail(data, cache, message_type="PROMPT", call_type=call_type)
@log_guardrail_information
async def async_post_call_success_hook(
@ -375,6 +376,7 @@ class LassoGuardrail(CustomGuardrail):
data: dict,
cache: DualCache,
message_type: Literal["PROMPT", "COMPLETION"] = "PROMPT",
call_type: str | None = None,
):
"""
Run the Lasso guardrail with the specified message type.
@ -395,6 +397,7 @@ class LassoGuardrail(CustomGuardrail):
data: The request data containing messages
cache: The cache instance for storing conversation_id (optional for post-call)
message_type: Either "PROMPT" for input or "COMPLETION" for output
call_type: The request's call type, which decides whether Lasso lists itself as applied
Raises:
LassoGuardrailAPIError: If the Lasso API call fails
@ -412,7 +415,8 @@ class LassoGuardrail(CustomGuardrail):
messages: Final[Sequence[Mapping[str, object]]] = [*expanded_messages, *input_messages]
if not messages:
return data
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
# Lasso's classifix endpoint returns masked text that we copy back
# into ``data["messages"]``. For multimodal/Responses-API input we

6
litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
View file

@ -185,6 +185,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
"""
use_native_lifecycle_hooks: ClassVar[bool] = True
adds_itself_to_applied_guardrails: ClassVar[bool] = True
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
@ -1328,6 +1329,11 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
self._raise_if_fail_closed(e)
return inputs
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
)
add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name)
flagged: Final = self._should_block_content(armor_response, allow_sanitization=False)
end_time: Final = time.time()
self.add_standard_logging_guardrail_information_to_request_data(

5
litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py
View file

@ -22,6 +22,7 @@ from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm._version import version as litellm_version
from litellm.integrations.custom_guardrail import (
APPLY_GUARDRAIL_ONLY_CALL_TYPES,
CustomGuardrail,
log_guardrail_information,
)
@ -351,7 +352,7 @@ class PillarGuardrail(CustomGuardrail):
verbose_proxy_logger.debug("Pillar Guardrail: Pre-call hook")
result: Final = await self.run_pillar_guardrail(data, user_api_key_dict)
if data.get("messages"):
if data.get("messages") or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
return result
@ -394,7 +395,7 @@ class PillarGuardrail(CustomGuardrail):
verbose_proxy_logger.debug("Pillar Guardrail: During-call moderation hook")
result: Final = await self.run_pillar_guardrail(data, user_api_key_dict)
if data.get("messages"):
if data.get("messages") or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
return result

10
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -306,7 +306,11 @@ class UnifiedLLMGuardrails(CustomLogger):
litellm_logging_obj=data.get("litellm_logging_obj"),
)
if not guardrail_to_apply.records_own_guardrail_information and not translation.lists_guardrail_when_called:
if not (
guardrail_to_apply.records_own_guardrail_information
or guardrail_to_apply.adds_itself_to_applied_guardrails
or translation.lists_guardrail_when_called
):
add_guardrail_to_applied_guardrails_header(
request_data=data, guardrail_name=guardrail_to_apply.guardrail_name
)
@ -447,7 +451,9 @@ class UnifiedLLMGuardrails(CustomLogger):
if e.original_response is None:
e.original_response = response
raise
if not guardrail_to_apply.records_own_guardrail_information:
if not (
guardrail_to_apply.records_own_guardrail_information or guardrail_to_apply.adds_itself_to_applied_guardrails
):
add_guardrail_to_applied_guardrails_header(
request_data=data, guardrail_name=guardrail_to_apply.guardrail_name
)

31
tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py
View file

@ -10,6 +10,8 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
from litellm.llms import get_guardrail_translation_mapping
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
from litellm.llms.openai.decisions.guardrail_translation.handler import DecisionsHandler, SystemOneHandler
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
from litellm.types.guardrails import ContentFilterAction, ContentFilterPattern
from litellm.types.decisions import OpenAIDecisionResponse
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
@ -212,6 +214,35 @@ async def test_systemone_secret_key_and_card_number_are_masked_before_the_provid
assert result["state"] == {"<AWS_KEY>": "<X>", "card": "<CARD>", "items": 2}
@pytest.mark.asyncio
async def test_systemone_masks_a_card_number_that_json_escaping_splits() -> None:
guardrail: Final = ContentFilterGuardrail(
guardrail_name="cards",
patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="visa", action=ContentFilterAction.MASK)],
)
data: Final = {
"model": "judge",
"state": {"card": "4111\t1111\t1111\t1111", "4111\n1111\n1111\n1111": "on file", "note": "first\nsecond"},
"questions": {"q": {"type": "noul", "criteria": {}}},
}
result: Final = await SystemOneHandler().process_input_messages(data, guardrail)
assert result["state"] == {"card": "[VISA_REDACTED]", "[VISA_REDACTED]": "on file", "note": "first\nsecond"}
@pytest.mark.asyncio
async def test_systemone_rejects_an_escaped_entry_whose_two_scanned_forms_disagree() -> None:
data: Final = {
"model": "judge",
"state": {"note": "first\nsecond"},
"questions": {"q": {"type": "noul", "criteria": {}}},
}
with pytest.raises(UnappliableRequestRewrite):
await SystemOneHandler().process_input_messages(data, _TaggingGuardrail())
@pytest.mark.asyncio
async def test_systemone_keeps_every_state_entry_when_masked_keys_collide() -> None:
data: Final = {

45
tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py
View file

@ -15,6 +15,7 @@ import respx
import litellm
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import CallTypesLiteral
from litellm.proxy.guardrails.guardrail_hooks.dynamoai import DynamoAIGuardrails
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@ -209,15 +210,15 @@ async def test_post_call_hook_skips_dynamoai_when_the_reply_has_no_text(respx_mo
_INPUT_HOOKS = [
pytest.param(
"pre_call",
lambda guardrail, data: guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type="completion"
lambda guardrail, data, call_type: guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type
),
id="pre_call",
),
pytest.param(
"during_call",
lambda guardrail, data: guardrail.async_moderation_hook(
data=data, user_api_key_dict=UserAPIKeyAuth(), call_type="completion"
lambda guardrail, data, call_type: guardrail.async_moderation_hook(
data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type
),
id="during_call",
),
@ -228,17 +229,19 @@ _INPUT_HOOKS = [
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS)
@pytest.mark.parametrize(
("request_fields", "applied"),
("request_fields", "call_type", "applied"),
[
pytest.param({"messages": [{"role": "user", "content": "hi"}]}, ["test-dynamoai"], id="messages"),
pytest.param({"state": "hi"}, None, id="no-messages"),
pytest.param({"messages": [{"role": "user", "content": "hi"}]}, "completion", ["test-dynamoai"], id="messages"),
pytest.param({"state": "hi"}, "aextraction", None, id="no-messages-on-extraction"),
pytest.param({"state": "hi"}, "completion", ["test-dynamoai"], id="no-messages-on-chat"),
],
)
async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text(
respx_mock: respx.MockRouter,
event_hook: str,
run_hook: Callable[[DynamoAIGuardrails, Mapping[str, object]], Awaitable[object]],
run_hook: Callable[[DynamoAIGuardrails, Mapping[str, object], CallTypesLiteral], Awaitable[object]],
request_fields: Mapping[str, object],
call_type: CallTypesLiteral,
applied: Sequence[str] | None,
) -> None:
respx_mock.post(_ANALYZE_URL).respond(json={"finalAction": "NONE", "appliedPolicies": []})
@ -251,7 +254,7 @@ async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text(
)
metadata: Final[dict[str, object]] = {}
await run_hook(guardrail, {**request_fields, "metadata": metadata})
await run_hook(guardrail, {**request_fields, "metadata": metadata}, call_type)
assert metadata.get("applied_guardrails") == applied
@ -259,21 +262,35 @@ async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text(
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.parametrize(
("response", "applied"),
("response", "request_route", "applied"),
[
pytest.param(_assistant_reply("Hello, how are you?"), ["test-dynamoai"], id="chat-reply"),
pytest.param({"answers": {"aws_key": {"type": "text", "value": None}}}, None, id="no-choices"),
pytest.param(_assistant_reply("Hello, how are you?"), "/v1/chat/completions", ["test-dynamoai"], id="chat-reply"),
pytest.param(
{"answers": {"aws_key": {"type": "text", "value": None}}},
"/v1/extraction",
None,
id="no-choices-on-extraction",
),
pytest.param(
{"answers": {"aws_key": {"type": "text", "value": None}}},
"/v1/responses",
["test-dynamoai"],
id="no-choices-on-responses",
),
],
)
async def test_post_call_hook_lists_the_guardrail_only_when_it_scanned_text(
respx_mock: respx.MockRouter, response: litellm.ModelResponse | Mapping[str, object], applied: Sequence[str] | None
respx_mock: respx.MockRouter,
response: litellm.ModelResponse | Mapping[str, object],
request_route: str,
applied: Sequence[str] | None,
) -> None:
respx_mock.post(_ANALYZE_URL).respond(json={"finalAction": "NONE", "appliedPolicies": []})
metadata: Final[dict[str, object]] = {}
await _post_call_guardrail().async_post_call_success_hook(
data={"messages": [{"role": "user", "content": "hi"}], "metadata": metadata},
user_api_key_dict=UserAPIKeyAuth(),
user_api_key_dict=UserAPIKeyAuth(request_route=request_route),
response=response,
)

44
tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py
View file

@ -9,6 +9,7 @@ import litellm
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import CallTypesLiteral
from litellm.proxy.guardrails.guardrail_hooks.ibm_guardrails.ibm_detector import IBMGuardrailDetector
pytestmark = pytest.mark.usefixtures("httpx_transport")
@ -131,33 +132,35 @@ async def test_orchestrator_detection_blocks_the_request_and_names_the_detector(
_INPUT_HOOKS = [
pytest.param(
"pre_call",
lambda guardrail, data: guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type="completion"
lambda guardrail, data, call_type: guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type
),
id="pre_call",
),
pytest.param(
"during_call",
lambda guardrail, data: guardrail.async_moderation_hook(
data=data, user_api_key_dict=UserAPIKeyAuth(), call_type="completion"
lambda guardrail, data, call_type: guardrail.async_moderation_hook(
data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type
),
id="during_call",
),
]
_REQUEST_TEXT_CASES = [
pytest.param({"messages": [{"role": "user", "content": "hello there"}]}, ["ibm-guard"], id="messages"),
pytest.param({"state": "hello there"}, None, id="no-messages"),
pytest.param({"messages": [{"role": "user", "content": "hello there"}]}, "completion", ["ibm-guard"], id="messages"),
pytest.param({"state": "hello there"}, "aextraction", None, id="no-messages-on-extraction"),
pytest.param({"state": "hello there"}, "completion", ["ibm-guard"], id="no-messages-on-chat"),
]
@pytest.mark.asyncio
@pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS)
@pytest.mark.parametrize(("request_fields", "applied"), _REQUEST_TEXT_CASES)
@pytest.mark.parametrize(("request_fields", "call_type", "applied"), _REQUEST_TEXT_CASES)
async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text(
respx_mock: respx.MockRouter,
event_hook: str,
run_hook: Callable[[IBMGuardrailDetector, Mapping[str, object]], Awaitable[object]],
run_hook: Callable[[IBMGuardrailDetector, Mapping[str, object], CallTypesLiteral], Awaitable[object]],
request_fields: Mapping[str, object],
call_type: CallTypesLiteral,
applied: Sequence[str] | None,
) -> None:
respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[]])
@ -172,27 +175,42 @@ async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text(
)
metadata: Final[dict[str, object]] = {}
await run_hook(guardrail, {**request_fields, "metadata": metadata})
await run_hook(guardrail, {**request_fields, "metadata": metadata}, call_type)
assert metadata.get("applied_guardrails") == applied
@pytest.mark.asyncio
@pytest.mark.parametrize(
("response", "applied"),
("response", "request_route", "applied"),
[
pytest.param(
litellm.ModelResponse(
choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}]
),
"/v1/chat/completions",
["ibm-guard"],
id="chat-reply",
),
pytest.param({"answers": {"aws_key": {"type": "text", "value": None}}}, None, id="no-choices"),
pytest.param(
{"answers": {"aws_key": {"type": "text", "value": None}}},
"/v1/extraction",
None,
id="no-choices-on-extraction",
),
pytest.param(
{"answers": {"aws_key": {"type": "text", "value": None}}},
"/v1/responses",
["ibm-guard"],
id="no-choices-on-responses",
),
],
)
async def test_post_call_hook_lists_the_guardrail_only_when_it_scanned_text(
respx_mock: respx.MockRouter, response: litellm.ModelResponse | Mapping[str, object], applied: Sequence[str] | None
respx_mock: respx.MockRouter,
response: litellm.ModelResponse | Mapping[str, object],
request_route: str,
applied: Sequence[str] | None,
) -> None:
respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[]])
guardrail: Final = IBMGuardrailDetector(
@ -208,7 +226,7 @@ async def test_post_call_hook_lists_the_guardrail_only_when_it_scanned_text(
await guardrail.async_post_call_success_hook(
data={"messages": [{"role": "user", "content": "hello"}], "metadata": metadata},
user_api_key_dict=UserAPIKeyAuth(),
user_api_key_dict=UserAPIKeyAuth(request_route=request_route),
response=response,
)

1
tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py
View file

@ -829,6 +829,7 @@ async def test_aim_streaming_block_stops_after_verified_chunks(monkeypatch: pyte
"adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, 1, ["aim"], id="decisions"
),
pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"),
pytest.param("acompletion", {"messages": [{"role": "user", "content": "Refund?"}]}, 1, None, id="chat"),
],
)
async def test_aim_calls_aim_and_lists_itself_only_when_the_request_has_text(

1
tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py
View file

@ -2800,6 +2800,7 @@ async def test_cato_still_inspects_every_conversational_call_type(call_type: str
"adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, 1, ["cato"], id="decisions"
),
pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"),
pytest.param("acompletion", {"messages": [{"role": "user", "content": "Refund?"}]}, 1, None, id="chat"),
],
)
async def test_cato_calls_cato_and_lists_itself_only_when_the_request_has_text(

19
tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py
View file

@ -83,3 +83,22 @@ async def test_clean_scan_of_string_input_lists_the_guardrail() -> None:
)
assert _applied_guardrails(data) == ("lakera-v1",)
@pytest.mark.asyncio
async def test_clean_scan_of_chat_messages_keeps_the_header_as_before() -> None:
scans: Final[list[httpx.Request]] = []
def respond(request: httpx.Request) -> httpx.Response:
scans.append(request)
return httpx.Response(200, json={"results": [{"flagged": False}]})
guardrail: Final = _lakera(httpx.MockTransport(respond))
data: Final[dict[str, object]] = {"model": "m", "messages": [{"role": "user", "content": "is this a refund?"}]}
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), cache=DualCache(), data=data, call_type="completion"
)
assert len(scans) == 1
assert _applied_guardrails(data) == ()

35
tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py
View file

@ -1554,30 +1554,41 @@ async def test_api_error_handling_legacy_messages(
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.parametrize(
("call_type", "request_fields", "applied"),
("call_type", "request_fields", "lasso_calls", "applied"),
[
pytest.param(
"adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, ["lasso"], id="decisions"
"adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, 1, ["lasso"], id="decisions"
),
pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, None, id="extraction"),
pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"),
pytest.param("acompletion", {"messages": [{"role": "user", "content": "Refund?"}]}, 1, None, id="chat"),
],
)
@pytest.mark.parametrize("event_hook", ["pre_call", "during_call"])
async def test_lasso_lists_itself_only_when_it_classified_request_text(
respx_mock: respx.MockRouter, call_type: str, request_fields: dict[str, object], applied: list[str] | None
respx_mock: respx.MockRouter,
event_hook: str,
call_type: str,
request_fields: dict[str, object],
lasso_calls: int,
applied: list[str] | None,
) -> None:
respx_mock.post("https://server.lasso.security/gateway/v3/classify").respond(
classify: Final = respx_mock.post("https://server.lasso.security/gateway/v3/classify").respond(
json={"deputies": {}, "findings": {}, "violations_detected": False}
)
guardrail: Final = LassoGuardrail(
lasso_api_key="test-api-key", guardrail_name="lasso", event_hook="pre_call", default_on=True
lasso_api_key="test-api-key", guardrail_name="lasso", event_hook=event_hook, default_on=True
)
metadata: Final[dict[str, object]] = {}
data: Final = {**request_fields, "metadata": metadata}
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
cache=DualCache(),
data={**request_fields, "metadata": metadata},
call_type=call_type,
)
if event_hook == "pre_call":
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type
)
else:
await guardrail.async_moderation_hook(
data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type, cache=DualCache()
)
assert classify.call_count == lasso_calls
assert metadata.get("applied_guardrails") == applied

21
tests/unit/proxy/guardrails/test_pillar_guardrails.py
View file

@ -27,6 +27,7 @@ from httpx import Request, Response
import litellm
from litellm import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import CallTypesLiteral
from litellm.proxy.common_utils.callback_utils import get_logging_caching_headers
from litellm.proxy.guardrails.guardrail_hooks.pillar import (
PillarGuardrail,
@ -1521,15 +1522,15 @@ if __name__ == "__main__":
_INPUT_HOOKS = [
pytest.param(
"pre_call",
lambda guardrail, data: guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type="completion"
lambda guardrail, data, call_type: guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type
),
id="pre_call",
),
pytest.param(
"during_call",
lambda guardrail, data: guardrail.async_moderation_hook(
data=data, user_api_key_dict=UserAPIKeyAuth(), call_type="completion"
lambda guardrail, data, call_type: guardrail.async_moderation_hook(
data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type
),
id="during_call",
),
@ -1540,17 +1541,19 @@ _INPUT_HOOKS = [
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS)
@pytest.mark.parametrize(
("request_fields", "applied"),
("request_fields", "call_type", "applied"),
[
pytest.param({"messages": [{"role": "user", "content": "hi"}]}, ["pillar-test"], id="messages"),
pytest.param({"state": "hi"}, None, id="no-messages"),
pytest.param({"messages": [{"role": "user", "content": "hi"}]}, "completion", ["pillar-test"], id="messages"),
pytest.param({"state": "hi"}, "aextraction", None, id="no-messages-on-extraction"),
pytest.param({"state": "hi"}, "completion", ["pillar-test"], id="no-messages-on-chat"),
],
)
async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text(
respx_mock: respx.MockRouter,
event_hook: str,
run_hook: Callable[[PillarGuardrail, Mapping[str, object]], Awaitable[object]],
run_hook: Callable[[PillarGuardrail, Mapping[str, object], CallTypesLiteral], Awaitable[object]],
request_fields: Mapping[str, object],
call_type: CallTypesLiteral,
applied: Sequence[str] | None,
) -> None:
respx_mock.post(url__startswith="https://api.pillar.security/api/v1/protect").respond(
@ -1566,6 +1569,6 @@ async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text(
)
metadata: Final[dict[str, object]] = {}
await run_hook(guardrail, {**request_fields, "metadata": metadata})
await run_hook(guardrail, {**request_fields, "metadata": metadata}, call_type)
assert metadata.get("applied_guardrails") == applied

54
tests/unit/proxy/test_proxy_logging_hook_detection.py
View file

@ -2,6 +2,7 @@ from collections.abc import Mapping
from itertools import chain
from typing import Final, Literal
import httpx
import pytest
from fastapi import HTTPException
from pydantic import BaseModel, ConfigDict
@ -13,7 +14,9 @@ from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ModelArmorGuardrail
from litellm.proxy.guardrails.guardrail_hooks.tool_permission import ToolPermissionGuardrail
from litellm.proxy.utils import ProxyLogging
from litellm.types.decisions import OpenAIDecisionResponse
@ -1199,6 +1202,57 @@ async def test_pre_call_hook_keeps_request_tags_where_the_router_reads_them(
assert _applied_guardrails(result) == ("keeps_native",)
class _ModelArmorAnswering(AsyncHTTPHandler):
def __init__(self, status_code: int) -> None:
self.status_code = status_code
async def post(self, url: str, **kwargs: object) -> httpx.Response:
clean: Final = {"sanitizationResult": {"filterMatchState": "NO_MATCH_FOUND", "filterResults": {}}}
return httpx.Response(self.status_code, json=clean, request=httpx.Request("POST", url))
async def _model_armor_token() -> tuple[str, str]:
return ("test-token", "test-project")
@pytest.mark.asyncio
@pytest.mark.parametrize(("status_code", "applied"), ((200, ("model-armor",)), (503, ())))
@pytest.mark.parametrize(
("call_type", "request_fields"),
(
("aextraction", {"state": "refund order 4411", "fields": {"total": {"type": "number"}}}),
("asystemone", {"state": {"ticket": "refund order 4411"}, "questions": {"q": {"type": "noul"}}}),
("adecisions", {"input": "refund order 4411", "questions": [{"type": "predicate", "instructions": "Refund?"}]}),
),
)
async def test_pre_call_hook_lists_model_armor_only_when_it_answered(
monkeypatch: pytest.MonkeyPatch,
call_type: CallTypesLiteral,
request_fields: Mapping[str, object],
status_code: int,
applied: tuple[str, ...],
) -> None:
guardrail: Final = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
guardrail_name="model-armor",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
fail_on_error=False,
async_handler=_ModelArmorAnswering(status_code),
access_token_provider=_model_armor_token,
)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"),
data={"model": "m", **request_fields},
call_type=call_type,
)
assert _applied_guardrails(result) == applied
@pytest.mark.asyncio
async def test_pre_call_hook_does_not_list_a_guardrail_for_a_systemone_request_without_text(
monkeypatch: pytest.MonkeyPatch,