mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
812a2217ca
* feat(guardrails): add sensitive data routing to on-premise models When a guardrail detects sensitive data, route to an on-premise model instead of blocking or redacting. All subsequent requests in that session continue routing to the same model (sticky routing). New config options for guardrails: - on_sensitive_data: 'block' (default) or 'route' - sensitive_data_route_to_model: target model for rerouting - sticky_session_routing: persist routing for session (default: true) New exception SensitiveDataRouteException triggers rerouting when raised by guardrails. The proxy catches it, stores the routing decision in cache, and modifies the request's model field. New hook _PROXY_SensitiveDataRoutingHandler checks incoming requests against cached routing decisions and applies sticky routing. https://claude.ai/code/session_01SQd4isBa3UyouRoGVou9dK * fix: black formatting for custom_guardrail.py https://claude.ai/code/session_01SQd4isBa3UyouRoGVou9dK * test: improve test coverage for sensitive data routing feature Add additional tests for: - Cache key format and TTL constants - Session ID extraction from multiple locations - Custom guardrail initialization with routing config - Exception string representation and custom messages - Redis cache paths including fallback behavior - Edge cases in pre-call hook https://claude.ai/code/session_01SQd4isBa3UyouRoGVou9dK * fix: use correct GuardrailRaisedException parameters Replace invalid 'source' parameter with 'guardrail_name' to match the exception's actual signature. https://claude.ai/code/session_01SQd4isBa3UyouRoGVou9dK * test: move sensitive data routing tests to hooks directory Move test file to align with source code structure. https://claude.ai/code/session_01SQd4isBa3UyouRoGVou9dK * fix(guardrails): honor sticky_session_routing flag and scope session routing per API key Propagate sticky_session_routing through SensitiveDataRouteException so a guardrail configured with sticky_session_routing=False reroutes only the triggering request without persisting a session override. Scope the routing cache key to the requesting API key so sessions from different tenants cannot collide, and warn when sticky routing is requested but the hook is not registered. * refactor(guardrails): dedupe session-id extraction and drop redundant import Extract the shared session-id lookup into get_session_id_from_request_data so the sensitive-data routing hook and CustomGuardrail no longer keep two identical copies of the logic. Remove the redundant local import of GuardrailRaisedException in handle_sensitive_data_detection, and document that detection_info is surfaced in request metadata and logs so it must not carry raw sensitive values. * fix(guardrails): guard None user_api_key_dict in sensitive data route handler * fix(responses): send application/json Content-Type on responses DELETE OpenAI's responses DELETE endpoint now rejects requests that arrive without a Content-Type header, defaulting them to application/octet-stream and returning 'Unsupported content type: application/octet-stream'. The delete handler sent no body and therefore no Content-Type, so the request failed. Declare application/json on the delete request, matching the OpenAI SDK. * fix(guardrails): backfill in-memory cache after redis hit in sensitive data routing When _get_routed_model resolves a routing override from Redis it now also populates the local in-memory cache. Without the write-back, a non-writing instance that only ever reads from Redis would lose the sticky routing decision the moment Redis became unavailable, silently reverting sensitive sessions to the default model. * fix(guardrails): scope sticky sensitive-data routing to JWT principal Keyless auth (JWT and similar) has no api_key, so every such caller shared the "default" cache namespace. One authenticated user could reuse another user's session_id, trip the guardrail, and silently force the other user's subsequent requests onto the cached on-prem model for the TTL. Resolve the routing tenant from the api_key when present, otherwise from a stable principal built from the user/team/org identity, before reading or writing the session route. * fix(guardrails): require route target model when on_sensitive_data='route' * fix(guardrails): mark user_api_key_dict Optional in sensitive-data route handler * fix(guardrails): use remaining redis ttl for local backfill and str env default * fix(guardrails): graceful block when routing configured but no session_id handle_sensitive_data_detection promised to raise only SensitiveDataRouteException or GuardrailRaisedException, but when routing was configured and the request had no session_id it let a ValueError from raise_sensitive_data_route_exception propagate, surfacing as an HTTP 500 instead of a block. Fall back to a graceful block in that case so the documented contract holds. * fix(guardrails): run remaining guardrails after sensitive-data reroute Defer the SensitiveDataRouteException until every guardrail in the pre-call loop has run, so downstream security guardrails are no longer skipped when an earlier guardrail triggers routing. The first reroute wins and a later guardrail that blocks still propagates. Also normalize on_sensitive_data to lowercase like sibling on_* config fields so case-insensitive values are accepted. * fix(guardrails): classify sensitive-data reroute as guardrail intervention * fix(guardrails): record sensitive-data reroute as prometheus intervention not error * fix(guardrails): record service span for routing guardrail and move case-normalizer to base params Drop the early continue so a guardrail that signals sensitive-data routing still emits its PROXY_PRE_CALL service span like every other callback. Move the lowercase normalizer onto BaseLitellmParams so on_sensitive_data is normalized consistently when BaseLitellmParams is constructed directly, matching the cross-field route->model validator that already lives on the base.
1295 lines
50 KiB
Python
1295 lines
50 KiB
Python
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from litellm.integrations.custom_guardrail import CustomGuardrail
|
|
from litellm.proxy._types import CallTypes, UserAPIKeyAuth
|
|
from litellm.types.utils import GuardrailTracingDetail
|
|
|
|
|
|
class TestCustomGuardrailDeploymentHook:
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_pre_call_deployment_hook_no_guardrails(self):
|
|
"""Test that method returns kwargs unchanged when no guardrails are present"""
|
|
custom_guardrail = CustomGuardrail()
|
|
|
|
# Test with guardrails as None
|
|
kwargs = {
|
|
"messages": [{"role": "user", "content": "test message"}],
|
|
"model": "gpt-3.5-turbo",
|
|
"guardrails": None,
|
|
}
|
|
|
|
result = await custom_guardrail.async_pre_call_deployment_hook(
|
|
kwargs=kwargs, call_type=CallTypes.completion
|
|
)
|
|
|
|
assert result == kwargs
|
|
|
|
# Test with guardrails as non-list
|
|
kwargs["guardrails"] = "not_a_list"
|
|
|
|
result = await custom_guardrail.async_pre_call_deployment_hook(
|
|
kwargs=kwargs, call_type=CallTypes.completion
|
|
)
|
|
|
|
assert result == kwargs
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_pre_call_deployment_hook_with_guardrails_and_message_update(
|
|
self,
|
|
):
|
|
"""Test that method processes guardrails and updates messages when result contains messages"""
|
|
custom_guardrail = CustomGuardrail()
|
|
|
|
# Mock the async_pre_call_hook method
|
|
mock_result = {"messages": [{"role": "user", "content": "filtered message"}]}
|
|
custom_guardrail.async_pre_call_hook = AsyncMock(return_value=mock_result)
|
|
|
|
original_messages = [{"role": "user", "content": "original message"}]
|
|
kwargs = {
|
|
"messages": original_messages,
|
|
"model": "gpt-3.5-turbo",
|
|
"guardrails": ["some_guardrail"],
|
|
"user_api_key_user_id": "test_user",
|
|
"user_api_key_team_id": "test_team",
|
|
"user_api_key_end_user_id": "test_end_user",
|
|
"user_api_key_hash": "test_hash",
|
|
"user_api_key_request_route": "test_route",
|
|
}
|
|
|
|
result = await custom_guardrail.async_pre_call_deployment_hook(
|
|
kwargs=kwargs, call_type=CallTypes.completion
|
|
)
|
|
|
|
# Verify async_pre_call_hook was called with correct parameters
|
|
custom_guardrail.async_pre_call_hook.assert_called_once()
|
|
call_args = custom_guardrail.async_pre_call_hook.call_args
|
|
|
|
# Check that UserAPIKeyAuth was created properly
|
|
user_api_key_dict = call_args[1]["user_api_key_dict"]
|
|
assert isinstance(user_api_key_dict, UserAPIKeyAuth)
|
|
assert user_api_key_dict.user_id == "test_user"
|
|
assert user_api_key_dict.team_id == "test_team"
|
|
assert user_api_key_dict.end_user_id == "test_end_user"
|
|
assert user_api_key_dict.api_key == "test_hash"
|
|
assert user_api_key_dict.request_route == "test_route"
|
|
|
|
# Check other parameters
|
|
assert call_args[1]["data"] == kwargs
|
|
assert call_args[1]["call_type"] == "completion"
|
|
|
|
# Verify messages were updated in result
|
|
assert result["messages"] == mock_result["messages"]
|
|
assert result["messages"] != original_messages
|
|
|
|
|
|
class TestCustomGuardrailShouldRunGuardrail:
|
|
|
|
def test_should_run_guardrail_with_litellm_metadata(self):
|
|
"""Test that should_run_guardrail works with litellm_metadata pattern"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
custom_guardrail = CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
default_on=False,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# Test with guardrails in litellm_metadata
|
|
data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"litellm_metadata": {"guardrails": ["test_guardrail"]},
|
|
}
|
|
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
|
|
assert result is True
|
|
|
|
def test_should_run_guardrail_with_metadata(self):
|
|
"""Test that should_run_guardrail works with metadata pattern"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
custom_guardrail = CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
default_on=False,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# Test with guardrails in metadata
|
|
data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"metadata": {"guardrails": ["test_guardrail"]},
|
|
}
|
|
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
|
|
assert result is True
|
|
|
|
def test_should_run_guardrail_with_root_level_guardrails(self):
|
|
"""Test that should_run_guardrail works with root level guardrails"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
custom_guardrail = CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
default_on=False,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# Test with guardrails at root level
|
|
data = {"model": "gpt-3.5-turbo", "guardrails": ["test_guardrail"]}
|
|
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
|
|
assert result is True
|
|
|
|
def test_should_run_guardrail_no_matching_guardrail(self):
|
|
"""Test that should_run_guardrail returns False when guardrail name doesn't match"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
custom_guardrail = CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
default_on=False,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# Test with different guardrail name
|
|
data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"litellm_metadata": {"guardrails": ["different_guardrail"]},
|
|
}
|
|
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
|
|
assert result is False
|
|
|
|
def test_should_run_guardrail_with_disable_global_guardrail(self):
|
|
"""Test that disable_global_guardrails only works from admin metadata"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
custom_guardrail = CustomGuardrail(
|
|
guardrail_name="global_guardrail",
|
|
default_on=True,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# Test 1: Global guardrail runs by default
|
|
data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
}
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
assert result is True, "Global guardrail should run when default_on=True"
|
|
|
|
# Test 2: User-injected disable at root level is IGNORED
|
|
data_with_disable_root = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"disable_global_guardrails": True,
|
|
}
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data_with_disable_root, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
assert (
|
|
result is True
|
|
), "User-injected disable_global_guardrails should be ignored"
|
|
|
|
# Test 3: User-injected disable in metadata is IGNORED
|
|
data_with_disable_metadata = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"disable_global_guardrails": True},
|
|
}
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data_with_disable_metadata, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
assert result is True, "User-injected metadata disable should be ignored"
|
|
|
|
# Test 4: Admin-configured disable via user_api_key_metadata IS respected
|
|
data_with_admin_disable = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}},
|
|
}
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data_with_admin_disable, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
assert result is False, "Admin-configured disable should be respected"
|
|
|
|
# Test 5: Admin config in metadata isn't shadowed by user-supplied litellm_metadata
|
|
data_cross_key = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}},
|
|
"litellm_metadata": {"request_tags": ["user-supplied"]},
|
|
}
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data_cross_key, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
assert (
|
|
result is False
|
|
), "Admin config in metadata must not be shadowed by user-supplied litellm_metadata"
|
|
|
|
# Test 6: After the pre-call strip runs, user-injected
|
|
# user_api_key_metadata in the non-authoritative metadata key is gone.
|
|
# _get_admin_metadata must then surface admin config unchanged.
|
|
data_post_strip = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}},
|
|
"litellm_metadata": {}, # post-strip: attacker payload removed
|
|
}
|
|
result = custom_guardrail.should_run_guardrail(
|
|
data=data_post_strip, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
assert (
|
|
result is False
|
|
), "Admin config in metadata must be respected when other metadata key is empty"
|
|
|
|
def test_should_run_guardrail_with_opted_out_global_guardrails(self):
|
|
"""Test that per-guardrail opt-out only works from admin metadata"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
custom_guardrail = CustomGuardrail(
|
|
guardrail_name="global_guardrail",
|
|
default_on=True,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# Test 1: User-injected opt-out at root level is IGNORED
|
|
data_root = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"opted_out_global_guardrails": ["global_guardrail"],
|
|
}
|
|
assert (
|
|
custom_guardrail.should_run_guardrail(
|
|
data=data_root, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
is True
|
|
)
|
|
|
|
# Test 2: User-injected opt-out in metadata is IGNORED
|
|
data_metadata = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"opted_out_global_guardrails": ["global_guardrail"]},
|
|
}
|
|
assert (
|
|
custom_guardrail.should_run_guardrail(
|
|
data=data_metadata, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
is True
|
|
)
|
|
|
|
# Test 4: a different guardrail in the opt-out list → still runs
|
|
data_other = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"opted_out_global_guardrails": ["some_other_guardrail"]},
|
|
}
|
|
assert (
|
|
custom_guardrail.should_run_guardrail(
|
|
data=data_other, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
is True
|
|
)
|
|
|
|
# Test 5: empty opt-out list → still runs
|
|
data_empty = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"opted_out_global_guardrails": []},
|
|
}
|
|
assert (
|
|
custom_guardrail.should_run_guardrail(
|
|
data=data_empty, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
is True
|
|
)
|
|
|
|
# Test 6: malformed value (bool instead of list) → safely ignored, guardrail runs
|
|
data_malformed = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {"opted_out_global_guardrails": True},
|
|
}
|
|
assert (
|
|
custom_guardrail.should_run_guardrail(
|
|
data=data_malformed, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
is True
|
|
)
|
|
|
|
def test_should_run_guardrail_opt_out_does_not_affect_non_global(self):
|
|
"""Opt-out list only matters for default_on=True guardrails"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
non_global = CustomGuardrail(
|
|
guardrail_name="opt_in_guardrail",
|
|
default_on=False,
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
# An opt-in guardrail named in opted_out_global_guardrails is still controlled
|
|
# by the explicit `guardrails` request list, not by the global opt-out list.
|
|
data = {
|
|
"model": "gpt-3.5-turbo",
|
|
"messages": [{"role": "user", "content": "test"}],
|
|
"metadata": {
|
|
"opted_out_global_guardrails": ["opt_in_guardrail"],
|
|
"guardrails": ["opt_in_guardrail"],
|
|
},
|
|
}
|
|
assert (
|
|
non_global.should_run_guardrail(
|
|
data=data, event_type=GuardrailEventHooks.pre_call
|
|
)
|
|
is True
|
|
)
|
|
|
|
|
|
class TestApplyGuardrailCheck:
|
|
def test_apply_guardrail_check_only_on_direct_implementation(self):
|
|
"""
|
|
Test that "apply_guardrail" in type(callback).__dict__ only returns True
|
|
when the object's own class implements the method, not when it's inherited
|
|
from a parent class.
|
|
|
|
This is critical for properly routing guardrail handling to the unified
|
|
guardrail handler vs the guardrail's own implementation.
|
|
"""
|
|
|
|
# Parent class with apply_guardrail (CustomGuardrail already has it)
|
|
class ParentGuardrail(CustomGuardrail):
|
|
"""Parent that inherits apply_guardrail from CustomGuardrail"""
|
|
|
|
pass
|
|
|
|
# Child class that only inherits apply_guardrail (doesn't override)
|
|
class ChildGuardrailWithoutOverride(ParentGuardrail):
|
|
"""Child that only inherits apply_guardrail"""
|
|
|
|
pass
|
|
|
|
# Child class that overrides apply_guardrail
|
|
class ChildGuardrailWithOverride(ParentGuardrail):
|
|
"""Child that overrides apply_guardrail"""
|
|
|
|
async def apply_guardrail(self, text, language=None, entities=None):
|
|
return f"modified: {text}"
|
|
|
|
# Instantiate the classes
|
|
parent_instance = ParentGuardrail()
|
|
child_without_override = ChildGuardrailWithoutOverride()
|
|
child_with_override = ChildGuardrailWithOverride()
|
|
|
|
# Test: CustomGuardrail itself has apply_guardrail in its __dict__
|
|
assert (
|
|
"apply_guardrail" in type(CustomGuardrail()).__dict__
|
|
), "CustomGuardrail should have apply_guardrail in its own __dict__"
|
|
|
|
# Test: ParentGuardrail inherits but doesn't override, so it should NOT be in __dict__
|
|
assert (
|
|
"apply_guardrail" not in type(parent_instance).__dict__
|
|
), "ParentGuardrail should NOT have apply_guardrail in its own __dict__ (only inherited)"
|
|
|
|
# Test: ChildGuardrailWithoutOverride only inherits, should NOT be in __dict__
|
|
assert (
|
|
"apply_guardrail" not in type(child_without_override).__dict__
|
|
), "ChildGuardrailWithoutOverride should NOT have apply_guardrail in its own __dict__ (only inherited)"
|
|
|
|
# Test: ChildGuardrailWithOverride overrides the method, SHOULD be in __dict__
|
|
assert (
|
|
"apply_guardrail" in type(child_with_override).__dict__
|
|
), "ChildGuardrailWithOverride SHOULD have apply_guardrail in its own __dict__ (overridden)"
|
|
|
|
# Verify that all instances still have the method via inheritance (hasattr)
|
|
assert hasattr(
|
|
parent_instance, "apply_guardrail"
|
|
), "All instances should have apply_guardrail via inheritance"
|
|
assert hasattr(
|
|
child_without_override, "apply_guardrail"
|
|
), "All instances should have apply_guardrail via inheritance"
|
|
assert hasattr(
|
|
child_with_override, "apply_guardrail"
|
|
), "All instances should have apply_guardrail via inheritance"
|
|
|
|
|
|
class TestGuardrailLoggingAggregation:
|
|
def _make_guardrail(self):
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
return CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
def _invoke_add_log(self, request_data: dict) -> None:
|
|
guardrail = self._make_guardrail()
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response={"result": "ok"},
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
start_time=1.0,
|
|
end_time=2.0,
|
|
duration=1.0,
|
|
masked_entity_count={"EMAIL": 1},
|
|
guardrail_provider="presidio",
|
|
)
|
|
|
|
def test_appends_to_existing_metadata_list(self):
|
|
request_data = {
|
|
"metadata": {
|
|
"standard_logging_guardrail_information": [
|
|
{"guardrail_name": "existing_guardrail"}
|
|
]
|
|
}
|
|
}
|
|
|
|
self._invoke_add_log(request_data)
|
|
|
|
info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert isinstance(info, list)
|
|
assert len(info) == 2
|
|
assert info[0]["guardrail_name"] == "existing_guardrail"
|
|
assert info[1]["guardrail_name"] == "test_guardrail"
|
|
|
|
def test_converts_existing_metadata_dict_to_list(self):
|
|
request_data = {
|
|
"metadata": {
|
|
"standard_logging_guardrail_information": {"guardrail_name": "legacy"}
|
|
}
|
|
}
|
|
|
|
self._invoke_add_log(request_data)
|
|
|
|
info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert isinstance(info, list)
|
|
assert len(info) == 2
|
|
assert info[0]["guardrail_name"] == "legacy"
|
|
assert info[1]["guardrail_name"] == "test_guardrail"
|
|
|
|
def test_appends_to_litellm_metadata(self):
|
|
request_data = {
|
|
"litellm_metadata": {
|
|
"standard_logging_guardrail_information": [
|
|
{"guardrail_name": "litellm_existing"}
|
|
]
|
|
}
|
|
}
|
|
|
|
self._invoke_add_log(request_data)
|
|
|
|
info = request_data["litellm_metadata"][
|
|
"standard_logging_guardrail_information"
|
|
]
|
|
assert isinstance(info, list)
|
|
assert len(info) == 2
|
|
assert info[1]["guardrail_name"] == "test_guardrail"
|
|
|
|
|
|
class TestGuardrailOtelSpanEmission:
|
|
"""Recording a guardrail emits its otel span inline, so every guardrail
|
|
execution produces a span — including the pass-through allow path that never
|
|
reaches a post-call hook."""
|
|
|
|
def _make_guardrail(self):
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
return CustomGuardrail(
|
|
guardrail_name="emit_guard",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
def _record(self, guardrail, request_data):
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response={"result": "ok"},
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
start_time=1.0,
|
|
end_time=2.0,
|
|
duration=1.0,
|
|
)
|
|
|
|
def test_emits_span_for_recorded_entry(self, monkeypatch):
|
|
captured = []
|
|
monkeypatch.setattr(
|
|
"litellm.integrations.otel.logger.emit_guardrail_span",
|
|
captured.append,
|
|
)
|
|
|
|
request_data = {"metadata": {}}
|
|
self._record(self._make_guardrail(), request_data)
|
|
|
|
assert len(captured) == 1
|
|
emitted = captured[0]
|
|
recorded = request_data["metadata"]["standard_logging_guardrail_information"][
|
|
-1
|
|
]
|
|
assert emitted is recorded
|
|
assert emitted["guardrail_name"] == "emit_guard"
|
|
assert emitted["start_time"] == 1.0
|
|
assert emitted["end_time"] == 2.0
|
|
|
|
def test_span_emission_failure_does_not_break_recording(self, monkeypatch):
|
|
def _boom(_entry):
|
|
raise RuntimeError("otel exporter down")
|
|
|
|
monkeypatch.setattr(
|
|
"litellm.integrations.otel.logger.emit_guardrail_span", _boom
|
|
)
|
|
|
|
request_data = {"metadata": {}}
|
|
self._record(self._make_guardrail(), request_data)
|
|
|
|
info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(info) == 1
|
|
assert info[0]["guardrail_name"] == "emit_guard"
|
|
|
|
|
|
class TestGuardrailSensitiveFieldStripping:
|
|
"""Tests that secret_fields is stripped from guardrail responses before logging.
|
|
|
|
Matches the pattern used by Langfuse and Arize integrations which also
|
|
pop("secret_fields") to prevent raw Authorization headers from being persisted.
|
|
"""
|
|
|
|
def _make_guardrail(self):
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
return CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
def test_secret_fields_stripped_from_guardrail_response(self):
|
|
"""Ensure secret_fields (containing raw Authorization headers) is not persisted."""
|
|
guardrail = self._make_guardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
guardrail_response_with_secrets = {
|
|
"model": "gpt-4",
|
|
"messages": [{"role": "user", "content": "hello"}],
|
|
"secret_fields": {
|
|
"raw_headers": {
|
|
"authorization": "Bearer sk-live-secret-key-12345",
|
|
"content-type": "application/json",
|
|
}
|
|
},
|
|
"proxy_server_request": {"url": "http://localhost:4000/chat/completions"},
|
|
}
|
|
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response=guardrail_response_with_secrets,
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
duration=1.0,
|
|
)
|
|
|
|
info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(info) == 1
|
|
logged_response = info[0]["guardrail_response"]
|
|
|
|
# secret_fields must be stripped
|
|
assert "secret_fields" not in logged_response
|
|
|
|
# Other fields should be preserved
|
|
assert "model" in logged_response
|
|
assert "messages" in logged_response
|
|
assert "proxy_server_request" in logged_response
|
|
|
|
def test_string_guardrail_response_not_affected(self):
|
|
"""String responses (e.g. 'allow', 'deny') should pass through unchanged."""
|
|
guardrail = self._make_guardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response="allow",
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
duration=0.5,
|
|
)
|
|
|
|
info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert info[0]["guardrail_response"] == "allow"
|
|
|
|
def test_no_authorization_header_in_logged_response(self):
|
|
"""Verify no plaintext Authorization header ends up in the logged guardrail response."""
|
|
import json
|
|
|
|
guardrail = self._make_guardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response={
|
|
"model": "gpt-4",
|
|
"secret_fields": {
|
|
"raw_headers": {
|
|
"authorization": "Bearer sk-live-SHOULD-NOT-APPEAR",
|
|
}
|
|
},
|
|
},
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
duration=1.0,
|
|
)
|
|
|
|
logged_response = request_data["metadata"][
|
|
"standard_logging_guardrail_information"
|
|
][0]["guardrail_response"]
|
|
assert "secret_fields" not in logged_response
|
|
assert "sk-live-SHOULD-NOT-APPEAR" not in json.dumps(logged_response)
|
|
|
|
def test_secret_fields_stripped_from_list_dict_response(self):
|
|
"""Ensure secret_fields is stripped from List[dict] guardrail responses too."""
|
|
guardrail = self._make_guardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response=[
|
|
{
|
|
"result": "ok",
|
|
"secret_fields": {
|
|
"raw_headers": {"authorization": "Bearer sk-secret"}
|
|
},
|
|
},
|
|
{"result": "also_ok"},
|
|
],
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
duration=1.0,
|
|
)
|
|
|
|
import json
|
|
|
|
serialized = json.dumps(request_data)
|
|
assert "secret_fields" not in serialized
|
|
assert "sk-secret" not in serialized
|
|
|
|
|
|
class TestCustomGuardrailPassthroughSupport:
|
|
"""Tests for passthrough endpoint guardrail support - Issue fixes."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_post_call_success_deployment_hook_with_httpx_response(self):
|
|
"""
|
|
Test that async_post_call_success_deployment_hook handles raw httpx.Response objects
|
|
from passthrough endpoints without crashing with TypeError.
|
|
|
|
This tests Fix #3: TypeError: TypedDict does not support instance and class checks
|
|
"""
|
|
import httpx
|
|
|
|
custom_guardrail = CustomGuardrail()
|
|
|
|
# Mock the async_post_call_success_hook to return None (guardrail didn't modify response)
|
|
custom_guardrail.async_post_call_success_hook = AsyncMock(return_value=None)
|
|
|
|
# Create a mock httpx.Response object (typical passthrough response)
|
|
mock_response = AsyncMock(spec=httpx.Response)
|
|
mock_response.status_code = 200
|
|
mock_response.text = "Mock response"
|
|
|
|
request_data = {
|
|
"guardrails": ["test_guardrail"],
|
|
"user_api_key_user_id": "test_user",
|
|
"user_api_key_team_id": "test_team",
|
|
"user_api_key_end_user_id": "test_end_user",
|
|
"user_api_key_hash": "test_hash",
|
|
"user_api_key_request_route": "passthrough_route",
|
|
}
|
|
|
|
# This should not raise TypeError: TypedDict does not support instance and class checks
|
|
result = await custom_guardrail.async_post_call_success_deployment_hook(
|
|
request_data=request_data,
|
|
response=mock_response,
|
|
call_type=CallTypes.allm_passthrough_route,
|
|
)
|
|
|
|
# When result is None, should return the original response
|
|
assert result == mock_response
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_post_call_success_deployment_hook_with_none_call_type(self):
|
|
"""
|
|
Test that async_post_call_success_deployment_hook handles None call_type gracefully.
|
|
|
|
This ensures that even if call_type is None (before fix #1), the guardrail doesn't crash.
|
|
"""
|
|
custom_guardrail = CustomGuardrail()
|
|
|
|
# Mock the async_post_call_success_hook to return None
|
|
custom_guardrail.async_post_call_success_hook = AsyncMock(return_value=None)
|
|
|
|
mock_response = AsyncMock()
|
|
|
|
request_data = {
|
|
"guardrails": ["test_guardrail"],
|
|
"user_api_key_user_id": "test_user",
|
|
}
|
|
|
|
# Call with None call_type - should not crash
|
|
result = await custom_guardrail.async_post_call_success_deployment_hook(
|
|
request_data=request_data,
|
|
response=mock_response,
|
|
call_type=None,
|
|
)
|
|
|
|
# Should return the original response when result is None
|
|
assert result == mock_response
|
|
|
|
def test_is_valid_response_type_with_none(self):
|
|
"""
|
|
Test _is_valid_response_type helper method correctly identifies None as invalid.
|
|
|
|
This is part of Fix #3: Safely handling TypedDict types that don't support isinstance checks.
|
|
"""
|
|
custom_guardrail = CustomGuardrail()
|
|
|
|
# None should be invalid
|
|
assert custom_guardrail._is_valid_response_type(None) is False
|
|
|
|
def test_is_valid_response_type_with_typeddict_error(self):
|
|
"""
|
|
Test _is_valid_response_type gracefully handles TypeError from TypedDict.
|
|
|
|
This tests Fix #3: When isinstance() is called with TypedDict types, it raises TypeError.
|
|
The method should catch this and allow the response through.
|
|
"""
|
|
from litellm.types.utils import ModelResponse
|
|
|
|
custom_guardrail = CustomGuardrail()
|
|
|
|
# Create a valid LiteLLM response object
|
|
response = ModelResponse(
|
|
id="test-id",
|
|
choices=[],
|
|
created=0,
|
|
model="test-model",
|
|
object="chat.completion",
|
|
)
|
|
|
|
# This should return True (it's a valid response type or TypeError is caught)
|
|
result = custom_guardrail._is_valid_response_type(response)
|
|
assert result is True
|
|
|
|
|
|
class TestEventTypeLogging:
|
|
"""Tests for event_type logging in guardrail information."""
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_guardrail_information_infers_event_type_from_async_pre_call_hook(
|
|
self,
|
|
):
|
|
"""
|
|
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.pre_call
|
|
from async_pre_call_hook function name.
|
|
"""
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="test_event_type_guardrail",
|
|
event_hook=[
|
|
GuardrailEventHooks.pre_call,
|
|
GuardrailEventHooks.post_call,
|
|
],
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def async_pre_call_hook(self, data: dict, **kwargs):
|
|
return {"result": "pre_call_executed"}
|
|
|
|
guardrail = TestGuardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
await guardrail.async_pre_call_hook(data=request_data)
|
|
|
|
# Check that the guardrail_mode was set to pre_call (not the full list)
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_guardrail_information_infers_event_type_from_async_post_call_success_hook(
|
|
self,
|
|
):
|
|
"""
|
|
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call
|
|
from async_post_call_success_hook function name.
|
|
"""
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="test_event_type_guardrail",
|
|
event_hook=[
|
|
GuardrailEventHooks.pre_call,
|
|
GuardrailEventHooks.post_call,
|
|
],
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def async_post_call_success_hook(self, data: dict, **kwargs):
|
|
return {"result": "post_call_executed"}
|
|
|
|
guardrail = TestGuardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
await guardrail.async_post_call_success_hook(data=request_data)
|
|
|
|
# Check that the guardrail_mode was set to post_call (not the full list)
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_guardrail_information_infers_event_type_from_async_moderation_hook(
|
|
self,
|
|
):
|
|
"""
|
|
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.during_call
|
|
from async_moderation_hook function name.
|
|
"""
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="test_event_type_guardrail",
|
|
event_hook=[
|
|
GuardrailEventHooks.during_call,
|
|
GuardrailEventHooks.post_call,
|
|
],
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def async_moderation_hook(self, data: dict, **kwargs):
|
|
return {"result": "moderation_executed"}
|
|
|
|
guardrail = TestGuardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
await guardrail.async_moderation_hook(data=request_data)
|
|
|
|
# Check that the guardrail_mode was set to during_call (not the full list)
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.during_call
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_guardrail_information_infers_event_type_from_async_post_call_streaming_hook(
|
|
self,
|
|
):
|
|
"""
|
|
Test that log_guardrail_information decorator correctly infers GuardrailEventHooks.post_call
|
|
from async_post_call_streaming_hook function name.
|
|
"""
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="test_event_type_guardrail",
|
|
event_hook=[
|
|
GuardrailEventHooks.pre_call,
|
|
GuardrailEventHooks.post_call,
|
|
],
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def async_post_call_streaming_hook(self, data: dict, **kwargs):
|
|
return {"result": "streaming_executed"}
|
|
|
|
guardrail = TestGuardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
await guardrail.async_post_call_streaming_hook(data=request_data)
|
|
|
|
# Check that the guardrail_mode was set to post_call (not the full list)
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_guardrail_information_returns_none_for_unknown_function_name(
|
|
self,
|
|
):
|
|
"""
|
|
Test that log_guardrail_information decorator returns None for event_type
|
|
when function name doesn't match known patterns, and falls back to self.event_hook.
|
|
"""
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="test_event_type_guardrail",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def some_other_hook(self, data: dict, **kwargs):
|
|
return {"result": "other_hook_executed"}
|
|
|
|
guardrail = TestGuardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
await guardrail.some_other_hook(data=request_data)
|
|
|
|
# Check that the guardrail_mode falls back to self.event_hook
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
|
|
|
|
def test_add_standard_logging_uses_event_type_over_event_hook(self):
|
|
"""
|
|
Test that add_standard_logging_guardrail_information_to_request_data
|
|
prioritizes event_type parameter over self.event_hook.
|
|
"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
guardrail = CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
event_hook=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call],
|
|
)
|
|
|
|
request_data = {"metadata": {}}
|
|
|
|
# Call with explicit event_type
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response={"result": "ok"},
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
event_type=GuardrailEventHooks.post_call,
|
|
)
|
|
|
|
# Should use the provided event_type (post_call), not the full event_hook list
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.post_call
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_guardrail_information_skips_auto_record_if_function_already_recorded(
|
|
self,
|
|
):
|
|
"""When a wrapped guardrail function records its own entry directly
|
|
(e.g. block_code_execution.apply_guardrail records a rich
|
|
``[detections...]`` payload), the decorator must NOT also append its
|
|
own ``"allow"``/raw-response entry — otherwise every backend
|
|
(OTEL spans, Datadog, Langfuse, spend logs) double-records one
|
|
logical guardrail invocation."""
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="block-code",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def apply_guardrail(self, inputs, request_data, **kwargs):
|
|
self.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response=[{"action_taken": "block"}],
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
event_type=GuardrailEventHooks.pre_call,
|
|
)
|
|
return inputs
|
|
|
|
guardrail = TestGuardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
await guardrail.apply_guardrail(
|
|
inputs={"texts": ["x"]}, request_data=request_data
|
|
)
|
|
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1, (
|
|
f"Decorator must not double-record when the wrapped function "
|
|
f"already appended its own entry; got {len(logged_info)} entries"
|
|
)
|
|
assert logged_info[0]["guardrail_response"] == [{"action_taken": "block"}]
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_log_guardrail_information_skips_auto_record_on_exception_if_function_already_recorded(
|
|
self,
|
|
):
|
|
"""Same as above on the failure path: if the wrapped function
|
|
appended an entry in its ``finally`` block before re-raising, the
|
|
decorator must just re-raise without auto-recording on top."""
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class TestGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="block-code",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def apply_guardrail(self, inputs, request_data, **kwargs):
|
|
try:
|
|
raise ValueError("blocked")
|
|
finally:
|
|
self.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response=[{"action_taken": "block"}],
|
|
request_data=request_data,
|
|
guardrail_status="guardrail_intervened",
|
|
event_type=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
guardrail = TestGuardrail()
|
|
request_data = {"metadata": {}}
|
|
|
|
with pytest.raises(ValueError, match="blocked"):
|
|
await guardrail.apply_guardrail(
|
|
inputs={"texts": ["x"]}, request_data=request_data
|
|
)
|
|
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_status"] == "guardrail_intervened"
|
|
|
|
def test_add_standard_logging_falls_back_to_event_hook_when_event_type_is_none(
|
|
self,
|
|
):
|
|
"""
|
|
Test that add_standard_logging_guardrail_information_to_request_data
|
|
falls back to self.event_hook when event_type is None.
|
|
"""
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
guardrail = CustomGuardrail(
|
|
guardrail_name="test_guardrail",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
request_data = {"metadata": {}}
|
|
|
|
# Call with event_type=None
|
|
guardrail.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response={"result": "ok"},
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
event_type=None,
|
|
)
|
|
|
|
# Should fall back to self.event_hook
|
|
logged_info = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(logged_info) == 1
|
|
assert logged_info[0]["guardrail_mode"] == GuardrailEventHooks.pre_call
|
|
|
|
|
|
class TestTracingFieldsPopulation:
|
|
"""Verify add_standard_logging_guardrail_information_to_request_data passes tracing_detail fields."""
|
|
|
|
def test_new_fields_set_on_slg(self):
|
|
cg = CustomGuardrail(guardrail_name="test-rail")
|
|
request_data = {"metadata": {}}
|
|
cg.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response={"result": "ok"},
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
tracing_detail=GuardrailTracingDetail(
|
|
guardrail_id="rail-123",
|
|
policy_template="EU AI Act Article 5",
|
|
detection_method="regex",
|
|
confidence_score=0.95,
|
|
match_details=[{"type": "pattern", "action_taken": "BLOCK"}],
|
|
patterns_checked=12,
|
|
alert_recipients=["admin@example.com"],
|
|
),
|
|
)
|
|
slg_list = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(slg_list) == 1
|
|
slg = slg_list[0]
|
|
assert slg["guardrail_id"] == "rail-123"
|
|
assert slg["policy_template"] == "EU AI Act Article 5"
|
|
assert slg["detection_method"] == "regex"
|
|
assert slg["confidence_score"] == 0.95
|
|
assert slg["patterns_checked"] == 12
|
|
assert slg["alert_recipients"] == ["admin@example.com"]
|
|
assert len(slg["match_details"]) == 1
|
|
|
|
def test_new_fields_default_to_absent(self):
|
|
"""When tracing_detail is not passed, new fields are absent from the SLG dict."""
|
|
cg = CustomGuardrail(guardrail_name="test-rail")
|
|
request_data = {"metadata": {}}
|
|
cg.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response="ok",
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
)
|
|
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
|
assert slg.get("guardrail_id") is None
|
|
assert slg.get("policy_template") is None
|
|
assert slg.get("confidence_score") is None
|
|
|
|
def test_multiple_guardrails_with_different_policies(self):
|
|
"""One request, multiple guardrails each with own policy_template."""
|
|
cg1 = CustomGuardrail(guardrail_name="rail-1")
|
|
cg2 = CustomGuardrail(guardrail_name="rail-2")
|
|
request_data = {"metadata": {}}
|
|
|
|
cg1.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response="ok",
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
tracing_detail=GuardrailTracingDetail(policy_template="GDPR"),
|
|
)
|
|
cg2.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response="blocked",
|
|
request_data=request_data,
|
|
guardrail_status="guardrail_intervened",
|
|
tracing_detail=GuardrailTracingDetail(
|
|
policy_template="EU AI Act Article 5"
|
|
),
|
|
)
|
|
|
|
slg_list = request_data["metadata"]["standard_logging_guardrail_information"]
|
|
assert len(slg_list) == 2
|
|
assert slg_list[0]["policy_template"] == "GDPR"
|
|
assert slg_list[1]["policy_template"] == "EU AI Act Article 5"
|
|
|
|
def test_classification_field_passed_through(self):
|
|
"""Classification dict for LLM-judge guardrails is passed through."""
|
|
cg = CustomGuardrail(guardrail_name="judge-rail")
|
|
request_data = {"metadata": {}}
|
|
classification = {
|
|
"flagged": True,
|
|
"category": "workplace_emotion_recognition",
|
|
"article_reference": "Article 5(1)(f)",
|
|
"confidence": 0.94,
|
|
"reason": "Request asks to analyze employee sentiment",
|
|
}
|
|
cg.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response="blocked",
|
|
request_data=request_data,
|
|
guardrail_status="guardrail_intervened",
|
|
tracing_detail=GuardrailTracingDetail(
|
|
classification=classification,
|
|
detection_method="llm-judge",
|
|
confidence_score=0.94,
|
|
),
|
|
)
|
|
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
|
assert slg["classification"] == classification
|
|
assert slg["detection_method"] == "llm-judge"
|
|
assert slg["confidence_score"] == 0.94
|
|
|
|
|
|
class TestCustomGuardrailSpendLogMatchRedaction:
|
|
"""Guardrail JSON persisted via standard_logging must not contain raw match spans."""
|
|
|
|
def test_add_standard_logging_redacts_nested_match(self):
|
|
cg = CustomGuardrail(guardrail_name="test-rail")
|
|
raw = {
|
|
"assessments": [
|
|
{
|
|
"sensitiveInformationPolicy": {
|
|
"piiEntities": [
|
|
{"type": "NAME", "match": "GG", "action": "BLOCKED"}
|
|
]
|
|
}
|
|
}
|
|
]
|
|
}
|
|
request_data: dict = {"metadata": {}}
|
|
cg.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response=raw,
|
|
request_data=request_data,
|
|
guardrail_status="guardrail_intervened",
|
|
)
|
|
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
|
assert (
|
|
slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"][
|
|
"piiEntities"
|
|
][0]["match"]
|
|
== "[REDACTED]"
|
|
)
|
|
assert (
|
|
raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
|
|
"match"
|
|
]
|
|
== "GG"
|
|
)
|
|
|
|
def test_add_standard_logging_redacts_regex_field(self):
|
|
cg = CustomGuardrail(guardrail_name="test-rail")
|
|
raw = {"filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}]}
|
|
request_data: dict = {"metadata": {}}
|
|
cg.add_standard_logging_guardrail_information_to_request_data(
|
|
guardrail_json_response=raw,
|
|
request_data=request_data,
|
|
guardrail_status="success",
|
|
)
|
|
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
|
assert slg["guardrail_response"]["filters"][0]["regex"] == "[REDACTED]"
|
|
assert raw["filters"][0]["regex"] == r"\d{3}-\d{2}-\d{4}"
|
|
|
|
|
|
class TestGuardrailInterventionClassification:
|
|
"""A routing decision is a deliberate guardrail intervention, not a failure."""
|
|
|
|
def test_sensitive_data_route_exception_is_intervention(self):
|
|
from litellm.exceptions import SensitiveDataRouteException
|
|
|
|
exc = SensitiveDataRouteException(
|
|
route_to_model="on-prem-model",
|
|
session_id="sess-1",
|
|
guardrail_name="pii-rail",
|
|
)
|
|
assert CustomGuardrail._is_guardrail_intervention(exc) is True
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_routing_logged_as_intervened_not_failed(self):
|
|
from litellm.exceptions import SensitiveDataRouteException
|
|
from litellm.integrations.custom_guardrail import log_guardrail_information
|
|
from litellm.types.guardrails import GuardrailEventHooks
|
|
|
|
class RoutingGuardrail(CustomGuardrail):
|
|
def __init__(self):
|
|
super().__init__(
|
|
guardrail_name="pii-rail",
|
|
event_hook=GuardrailEventHooks.pre_call,
|
|
)
|
|
|
|
@log_guardrail_information
|
|
async def async_pre_call_hook(self, data, **kwargs):
|
|
raise SensitiveDataRouteException(
|
|
route_to_model="on-prem-model",
|
|
session_id="sess-1",
|
|
guardrail_name=self.guardrail_name,
|
|
)
|
|
|
|
guardrail = RoutingGuardrail()
|
|
request_data: dict = {"metadata": {}}
|
|
|
|
with pytest.raises(SensitiveDataRouteException):
|
|
await guardrail.async_pre_call_hook(data=request_data)
|
|
|
|
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
|
assert slg["guardrail_status"] == "guardrail_intervened"
|