scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

feat(guardrails): centralize and optimize latest role message filtering in proxy dispatch

Browse source
This commit is contained in:
Harshit Jain 2026-02-24 17:12:49 +05:30
parent 7622f26918
commit cb256f5bc8
No known key found for this signature in database
GPG key ID: 36C392CD4415B4CF
7 changed files with 292 additions and 93 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
litellm/integrations/custom_guardrail.py
View file

@ -6,6 +6,7 @@ from typing import (
List,
Literal,
Optional,
Tuple,
Type,
Union,
get_args,
@ -92,6 +93,7 @@ class CustomGuardrail(CustomLogger):
mask_request_content: bool = False,
mask_response_content: bool = False,
violation_message_template: Optional[str] = None,
experimental_use_latest_role_message_only: bool = False,
**kwargs,
):
"""
@ -114,12 +116,41 @@ class CustomGuardrail(CustomLogger):
self.mask_request_content: bool = mask_request_content
self.mask_response_content: bool = mask_response_content
self.violation_message_template: Optional[str] = violation_message_template
self.experimental_use_latest_role_message_only: bool = (
experimental_use_latest_role_message_only
)
if supported_event_hooks:
## validate event_hook is in supported_event_hooks
self._validate_event_hook(event_hook, supported_event_hooks)
super().__init__(**kwargs)
def filter_messages_for_latest_role(
self, messages: List[AllMessageValues], target_role: str = "user"
) -> Tuple[
Optional[List[AllMessageValues]],
Optional[List[AllMessageValues]],
Optional[List[int]],
]:
"""Filter to only the latest message for target_role."""
for index in range(len(messages) - 1, -1, -1):
if messages[index].get("role") == target_role:
return [messages[index]], list(messages), [index]
return None, None, None
def merge_filtered_messages(
self,
original_messages: List[AllMessageValues],
updated_target_messages: List[AllMessageValues],
target_indices: List[int],
) -> List[AllMessageValues]:
"""Merge filtered guardrail results back into the original message list."""
merged = list(original_messages)
for idx, updated_msg in zip(target_indices, updated_target_messages):
if idx < len(merged):
merged[idx] = updated_msg
return merged
def render_violation_message(
self, default: str, context: Optional[Dict[str, Any]] = None
) -> str:

67
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -243,59 +243,37 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
messages: Optional[List[AllMessageValues]],
) -> GuardrailMessageFilterResult:
"""Return payload + merge metadata for the latest user message."""
# NOTE: This logic probably belongs in CustomGuardrail once other guardrails adopt the feature.
if messages is None:
return GuardrailMessageFilterResult(None, None, None)
if self.experimental_use_latest_role_message_only is not True:
return GuardrailMessageFilterResult(messages, None, None)
latest_index = self._find_latest_message_index(messages, target_role="user")
if latest_index is None:
return GuardrailMessageFilterResult(None, None, None)
original_messages = list(messages)
payload_messages = [messages[latest_index]]
(
filtered_messages,
original_messages,
target_indices,
) = self.filter_messages_for_latest_role(messages, target_role="user")
return GuardrailMessageFilterResult(
payload_messages=payload_messages,
payload_messages=filtered_messages,
original_messages=original_messages,
target_indices=[latest_index],
target_indices=target_indices,
)
def _find_latest_message_index(
self, messages: List[AllMessageValues], target_role: str
) -> Optional[int]:
for index in range(len(messages) - 1, -1, -1):
if messages[index].get("role", None) == target_role:
return index
return None
def _merge_filtered_messages(
self,
original_messages: Optional[List[AllMessageValues]],
updated_target_messages: List[AllMessageValues],
target_indices: Optional[List[int]],
) -> List[AllMessageValues]:
if not target_indices:
if not target_indices or original_messages is None:
return updated_target_messages
if not original_messages:
original_messages = []
merged_messages = list(original_messages)
if not merged_messages:
merged_messages = list(updated_target_messages)
for replacement_index, updated_message in zip(
target_indices, updated_target_messages
):
if replacement_index < len(merged_messages):
merged_messages[replacement_index] = updated_message
return merged_messages
# NOTE: Consider moving these helpers to CustomGuardrail when the filtering
# logic becomes shared across providers.
return self.merge_filtered_messages(
original_messages=original_messages,
updated_target_messages=updated_target_messages,
target_indices=target_indices,
)
#### CALL HOOKS - proxy only ####
def _load_credentials(
@ -470,9 +448,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
response = getattr(e, "response", None)
if isinstance(response, httpx.Response):
try:
status_code, detail_message = (
self._parse_bedrock_guardrail_error_response(response)
)
(
status_code,
detail_message,
) = self._parse_bedrock_guardrail_error_response(response)
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_provider=self.guardrail_provider,
guardrail_json_response={"error": detail_message},
@ -795,9 +774,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
#########################################################
########## 1. Make the Bedrock API request ##########
#########################################################
bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = (
None
)
bedrock_guardrail_response: Optional[
Union[BedrockGuardrailResponse, str]
] = None
try:
bedrock_guardrail_response = await self.make_bedrock_api_request(
source="INPUT", messages=filtered_messages, request_data=data
@ -867,9 +846,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
#########################################################
########## 1. Make the Bedrock API request ##########
#########################################################
bedrock_guardrail_response: Optional[Union[BedrockGuardrailResponse, str]] = (
None
)
bedrock_guardrail_response: Optional[
Union[BedrockGuardrailResponse, str]
] = None
try:
bedrock_guardrail_response = await self.make_bedrock_api_request(
source="INPUT", messages=filtered_messages, request_data=data

7
litellm/proxy/guardrails/guardrail_initializers.py
View file

@ -46,6 +46,7 @@ def initialize_lakera(litellm_params: LitellmParams, guardrail: Guardrail):
event_hook=litellm_params.mode,
category_thresholds=litellm_params.category_thresholds,
default_on=litellm_params.default_on,
experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only,
)
litellm.logging_callback_manager.add_litellm_callback(_lakera_callback)
return _lakera_callback
@ -66,6 +67,7 @@ def initialize_lakera_v2(litellm_params: LitellmParams, guardrail: Guardrail):
metadata=litellm_params.metadata,
dev_info=litellm_params.dev_info,
on_flagged=litellm_params.on_flagged,
experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only,
)
litellm.logging_callback_manager.add_litellm_callback(_lakera_v2_callback)
return _lakera_v2_callback
@ -94,6 +96,7 @@ def initialize_presidio(litellm_params: LitellmParams, guardrail: Guardrail):
presidio_anonymizer_api_base=litellm_params.presidio_anonymizer_api_base,
presidio_language=litellm_params.presidio_language,
apply_to_output=False,
experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only,
)
params.update(overrides)
callback = _OPTIONAL_PresidioPIIMasking(**params)
@ -139,6 +142,7 @@ def initialize_hide_secrets(litellm_params: LitellmParams, guardrail: Guardrail)
event_hook=litellm_params.mode,
guardrail_name=guardrail.get("guardrail_name", ""),
default_on=litellm_params.default_on,
experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only,
)
litellm.logging_callback_manager.add_litellm_callback(_secret_detection_object)
return _secret_detection_object
@ -166,6 +170,7 @@ def initialize_tool_permission(litellm_params: LitellmParams, guardrail: Guardra
on_disallowed_action=getattr(litellm_params, "on_disallowed_action", "block"),
default_on=litellm_params.default_on,
violation_message_template=litellm_params.violation_message_template,
experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only,
)
litellm.logging_callback_manager.add_litellm_callback(_tool_permission_callback)
return _tool_permission_callback
@ -186,6 +191,7 @@ def initialize_lasso(
mask=litellm_params.mask,
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only,
)
litellm.logging_callback_manager.add_litellm_callback(_lasso_callback)
@ -218,6 +224,7 @@ def initialize_panw_prisma_airs(litellm_params, guardrail):
fallback_on_error=getattr(litellm_params, "fallback_on_error", "block"),
timeout=float(getattr(litellm_params, "timeout", 10.0)),
violation_message_template=litellm_params.violation_message_template,
experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only,
)
litellm.logging_callback_manager.add_litellm_callback(_panw_callback)

70
litellm/proxy/utils.py
View file

@ -870,27 +870,56 @@ class ProxyLogging:
target = unified_guardrail if use_unified else callback
if hook_type == "pre_call":
return await target.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, # type: ignore
cache=self.call_details["user_api_key_cache"],
data=data,
call_type=call_type,
# Filter messages if the flag is enabled
original_messages = None
target_indices = None
if (
hook_type == "pre_call"
and hasattr(callback, "experimental_use_latest_role_message_only")
and callback.experimental_use_latest_role_message_only
and isinstance(data.get("messages"), list)
):
(
filtered,
original_messages,
target_indices,
) = callback.filter_messages_for_latest_role(data["messages"])
verbose_proxy_logger.debug(
"Filtered messages for latest role: %s", filtered
)
elif hook_type == "during_call":
return await target.async_moderation_hook(
data=data,
user_api_key_dict=user_api_key_dict, # type: ignore
call_type=call_type,
)
elif hook_type == "post_call":
return await target.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict, # type: ignore
data=data,
response=response, # type: ignore
)
else:
raise ValueError(f"Unknown hook_type: {hook_type}")
if filtered is not None:
data["messages"] = filtered
try:
if hook_type == "pre_call":
result = await target.async_pre_call_hook(
user_api_key_dict=user_api_key_dict, # type: ignore
cache=self.call_details["user_api_key_cache"],
data=data,
call_type=call_type,
)
elif hook_type == "during_call":
result = await target.async_moderation_hook(
data=data,
user_api_key_dict=user_api_key_dict, # type: ignore
call_type=call_type,
)
elif hook_type == "post_call":
result = await target.async_post_call_success_hook(
user_api_key_dict=user_api_key_dict, # type: ignore
data=data,
response=response, # type: ignore
)
else:
raise ValueError(f"Unknown hook_type: {hook_type}")
finally:
# Restore original messages with modifications merged back
if original_messages is not None and target_indices is not None:
data["messages"] = callback.merge_filtered_messages(
original_messages, data["messages"], target_indices
)
return result
async def _execute_guardrail_with_load_balancing(
self,
@ -3551,6 +3580,7 @@ class PrismaClient:
Run a reconnect cycle with direct db operations and a single overall timeout
budget to avoid long retries on hot paths (e.g. auth).
"""
async def _do_direct_reconnect() -> None:
try:
await self.db.disconnect()

86
tests/proxy_unit_tests/test_guardrail_hook_dispatch.py Normal file
View file

@ -0,0 +1,86 @@
import os
import sys
from typing import Any, Optional
from unittest.mock import MagicMock
import pytest
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.proxy.utils import ProxyLogging
class MockProviderGuardrail(CustomGuardrail):
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.last_hook_type = None
self.last_data_messages = None
async def async_pre_call_hook(
self,
user_api_key_dict: Any,
cache: Any,
data: dict,
call_type: Any,
) -> Optional[dict]:
self.last_hook_type = "pre_call"
# The hook should receive the filtered messages
import copy
self.last_data_messages = copy.deepcopy(data.get("messages"))
# Simulate the guardrail modifying the message (e.g. masking PII)
if self.last_data_messages and len(self.last_data_messages) > 0:
self.last_data_messages[0]["content"] = self.last_data_messages[0][
"content"
].upper()
data["messages"] = self.last_data_messages
return data
@pytest.mark.asyncio
async def test_guardrail_dispatch_message_filtering():
"""
Test that _execute_guardrail_hook correctly filters messages when
experimental_use_latest_role_message_only=True, and then merges the
modifications back.
"""
proxy_logging = ProxyLogging(user_api_key_cache=MagicMock())
guardrail = MockProviderGuardrail()
guardrail.experimental_use_latest_role_message_only = True
# Multi-turn conversation
data = {
"messages": [
{"role": "user", "content": "First message"},
{"role": "assistant", "content": "Assistant reply"},
{"role": "user", "content": "Second message"},
]
}
# Dispatch the pre_call hook
result = await proxy_logging._execute_guardrail_hook(
callback=guardrail,
hook_type="pre_call",
data=data,
user_api_key_dict=MagicMock(),
call_type="completion",
)
# 1. Verify the hook only saw the latest user message
assert guardrail.last_data_messages is not None
assert len(guardrail.last_data_messages) == 1
# Inside the hook, it was uppercased
assert guardrail.last_data_messages[0]["content"] == "SECOND MESSAGE"
# 2. Verify the original data was correctly restored and merged
assert "messages" in data
assert len(data["messages"]) == 3
assert data["messages"][0]["content"] == "First message"
assert data["messages"][1]["content"] == "Assistant reply"
# The last message should have the modification from the hook
assert data["messages"][2]["content"] == "SECOND MESSAGE"

101
tests/test_litellm/integrations/test_custom_guardrail.py
View file

@ -8,7 +8,6 @@ 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"""
@ -86,7 +85,6 @@ class TestCustomGuardrailDeploymentHook:
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
@ -394,21 +392,21 @@ class TestCustomGuardrailPassthroughSupport:
"""
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",
@ -417,14 +415,14 @@ class TestCustomGuardrailPassthroughSupport:
"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
@ -432,53 +430,53 @@ class TestCustomGuardrailPassthroughSupport:
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",
@ -487,13 +485,12 @@ class TestCustomGuardrailPassthroughSupport:
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."""
@ -787,7 +784,9 @@ class TestTracingFieldsPopulation:
guardrail_json_response="blocked",
request_data=request_data,
guardrail_status="guardrail_intervened",
tracing_detail=GuardrailTracingDetail(policy_template="EU AI Act Article 5"),
tracing_detail=GuardrailTracingDetail(
policy_template="EU AI Act Article 5"
),
)
slg_list = request_data["metadata"]["standard_logging_guardrail_information"]
@ -820,3 +819,61 @@ class TestTracingFieldsPopulation:
assert slg["classification"] == classification
assert slg["detection_method"] == "llm-judge"
assert slg["confidence_score"] == 0.94
class TestCustomGuardrailFiltering:
def test_filter_messages_for_latest_role_found(self):
cg = CustomGuardrail()
messages = [
{"role": "system", "content": "system prompt"},
{"role": "user", "content": "first message"},
{"role": "assistant", "content": "response"},
{"role": "user", "content": "latest message"},
]
filtered, original, indices = cg.filter_messages_for_latest_role(
messages, target_role="user"
)
assert filtered == [{"role": "user", "content": "latest message"}]
assert original == messages
assert indices == [3]
def test_filter_messages_for_latest_role_not_found(self):
cg = CustomGuardrail()
messages = [
{"role": "system", "content": "system prompt"},
{"role": "assistant", "content": "response"},
]
filtered, original, indices = cg.filter_messages_for_latest_role(
messages, target_role="user"
)
assert filtered is None
assert original is None
assert indices is None
def test_filter_messages_for_latest_role_empty(self):
cg = CustomGuardrail()
filtered, original, indices = cg.filter_messages_for_latest_role(
[], target_role="user"
)
assert filtered is None
assert original is None
assert indices is None
def test_merge_filtered_messages(self):
cg = CustomGuardrail()
original = [
{"role": "system", "content": "system prompt"},
{"role": "user", "content": "first message"},
{"role": "assistant", "content": "response"},
{"role": "user", "content": "latest message"},
]
updated = [{"role": "user", "content": "modified latest message"}]
indices = [3]
merged = cg.merge_filtered_messages(original, updated, indices)
assert len(merged) == 4
assert merged[0] == original[0]
assert merged[1] == original[1]
assert merged[2] == original[2]
assert merged[3] == {"role": "user", "content": "modified latest message"}

23
tests/test_litellm/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py
View file

@ -8,8 +8,7 @@ with guardrail transformations, including tool calls.
import json
import os
import sys
from typing import Any, List, Literal, Optional, Tuple
from unittest.mock import AsyncMock, MagicMock
from typing import Any, Literal, Optional
import pytest
@ -134,7 +133,9 @@ class TestOpenAIChatCompletionsHandlerToolsInput:
tool = guardrail.last_inputs["tools"][0]
assert tool["type"] == "function"
assert tool["function"]["name"] == "get_weather"
assert tool["function"]["description"] == "Get the current weather in a location"
assert (
tool["function"]["description"] == "Get the current weather in a location"
)
assert "parameters" in tool["function"]
@pytest.mark.asyncio
@ -189,7 +190,10 @@ class TestOpenAIChatCompletionsHandlerToolsInput:
assert guardrail.last_inputs is not None
# tools should not be in inputs if not provided
assert "tools" not in guardrail.last_inputs or guardrail.last_inputs.get("tools") is None
assert (
"tools" not in guardrail.last_inputs
or guardrail.last_inputs.get("tools") is None
)
@pytest.mark.asyncio
async def test_tools_and_tool_calls_both_passed(self):
@ -220,7 +224,10 @@ class TestOpenAIChatCompletionsHandlerToolsInput:
"type": "function",
"function": {
"name": "get_weather",
"parameters": {"type": "object", "properties": {"location": {"type": "string"}}},
"parameters": {
"type": "object",
"properties": {"location": {"type": "string"}},
},
},
}
],
@ -757,7 +764,7 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
This test verifies the fix for the bug where accessing chunk.choices[0]
would raise IndexError when a streaming chunk has an empty choices list.
"""
from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices
from litellm.types.utils import ModelResponseStream
handler = OpenAIChatCompletionsHandler()
guardrail = MockPassThroughGuardrail(guardrail_name="test")
@ -833,7 +840,9 @@ class TestOpenAIChatCompletionsHandlerStreamingOutput:
assert result == responses_so_far
@pytest.mark.asyncio
async def test_process_output_streaming_response_mixed_empty_and_valid_choices_no_finish(self):
async def test_process_output_streaming_response_mixed_empty_and_valid_choices_no_finish(
self,
):
"""Test streaming response with mix of empty and valid choices chunks (stream not finished)
This tests the has_stream_ended check when iterating through chunks with mixed choices.