mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
fix(guardrails): scan compact input in vendor guardrails and warn on skipped input guardrails
This commit is contained in:
mateo-berri
parent
6fa4c17577
commit
47ac65fd67
14 changed files with 208 additions and 271 deletions
|
|
@ -33,6 +33,7 @@ from litellm.types.guardrails import (
|
|||
from litellm.types.llms.openai import AllMessageValues
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel
|
||||
from litellm.types.utils import (
|
||||
RESPONSES_INPUT_CALL_TYPES,
|
||||
CallTypes,
|
||||
GenericGuardrailAPIInputs,
|
||||
GuardrailStatus,
|
||||
|
|
@ -1632,7 +1633,7 @@ class CustomGuardrail(CustomLogger):
|
|||
# /responses
|
||||
# User/System messages are stored in the "input" key, use litellm transformation to get the messages
|
||||
#########################################################
|
||||
if call_type == CallTypes.responses.value or call_type == CallTypes.aresponses.value:
|
||||
if call_type in RESPONSES_INPUT_CALL_TYPES:
|
||||
from typing import cast
|
||||
|
||||
from litellm.responses.litellm_completion_transformation.transformation import (
|
||||
|
|
|
|||
|
|
@ -43,7 +43,7 @@ from litellm.proxy.litellm_pre_call_utils import get_chain_id_from_headers
|
|||
from litellm.proxy.mcp_server.utils import JSONLeafPath, json_string_leaves
|
||||
from litellm.types.guardrails import GuardrailEventHooks, LitellmParams, Mode
|
||||
from litellm.types.proxy.guardrails.guardrail_hooks.akto import AktoGuardrailConfigModelOptionalParams
|
||||
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
|
||||
from litellm.types.utils import RESPONSES_INPUT_CALL_TYPES, CallTypes, GenericGuardrailAPIInputs
|
||||
|
||||
from .akto_attachments import request_attachments, without_attachment_content
|
||||
|
||||
|
|
@ -81,7 +81,6 @@ DEFAULT_REQUEST_PATH: Final = "/v1/chat/completions"
|
|||
MCP_PATH: Final = "/mcp"
|
||||
MCP_TOOL_PREFIX: Final = "mcp"
|
||||
DEFAULT_BLOCK_REASON: Final = "Blocked by Akto Guardrails"
|
||||
RESPONSES_API_CALL_TYPES: Final = frozenset((CallTypes.responses.value, CallTypes.aresponses.value))
|
||||
MESSAGES_API_CALL_TYPES: Final = frozenset((CallTypes.anthropic_messages.value, CallTypes.aanthropic_messages.value))
|
||||
UNMASKABLE_REASON: Final = "Content masked by Akto guardrail policy could not be applied"
|
||||
MALFORMED_ATTACHMENT_REASON: Final = "Attachment could not be read for the Akto guardrail check"
|
||||
|
|
@ -431,7 +430,7 @@ class AktoGuardrail(CustomGuardrail):
|
|||
api_messages: Final = self.messages_api_messages(request_data)
|
||||
# The Responses API sends "input", so a "messages" key there is a decoy
|
||||
raw_messages: Final = (
|
||||
None if call_type_of(request_data) in RESPONSES_API_CALL_TYPES else request_data.get("messages")
|
||||
None if call_type_of(request_data) in RESPONSES_INPUT_CALL_TYPES else request_data.get("messages")
|
||||
)
|
||||
messages: Final = (
|
||||
api_messages
|
||||
|
|
|
|||
|
|
@ -14,7 +14,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
)
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.llms.openai import AllMessageValues, ResponseInputParam
|
||||
from litellm.types.utils import CallTypes, CallTypesLiteral
|
||||
from litellm.types.utils import RESPONSES_INPUT_CALL_TYPES, CallTypesLiteral
|
||||
|
||||
# Azure Content Safety APIs have a 10,000 character limit per request.
|
||||
AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH: Final = 10000
|
||||
|
|
@ -26,7 +26,7 @@ AZURE_CONTENT_SAFETY_TEXT_RECORD_LENGTH: Final = 1000
|
|||
AZURE_CONTENT_SAFETY_DEFAULT_API_VERSION: Final = "2024-09-01"
|
||||
JAVELIN_API_VERSION_STORED_BY_OLDER_RELEASES: Final = "v1"
|
||||
|
||||
RESPONSES_API_CALL_TYPES: Final = frozenset({CallTypes.responses, CallTypes.aresponses})
|
||||
RESPONSES_API_CALL_TYPES: Final = RESPONSES_INPUT_CALL_TYPES
|
||||
|
||||
_RESPONSES_API_CALL_TYPES: Final = RESPONSES_API_CALL_TYPES
|
||||
|
||||
|
|
|
|||
|
|
@ -92,6 +92,7 @@ if TYPE_CHECKING:
|
|||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
|
||||
from litellm.types.utils import (
|
||||
RESPONSES_INPUT_CALL_TYPES,
|
||||
CallTypes,
|
||||
CallTypesLiteral,
|
||||
Choices,
|
||||
|
|
@ -226,9 +227,6 @@ def _redact_assessment_match_fields(assessments: list[dict]) -> list[dict]:
|
|||
return redacted if isinstance(redacted, list) else assessments
|
||||
|
||||
|
||||
_RESPONSES_API_CALL_TYPES: Final = frozenset({CallTypes.responses, CallTypes.aresponses})
|
||||
|
||||
|
||||
def _without_image_bytes(content: Sequence[BedrockContentItem]) -> tuple[BedrockContentItem, ...]:
|
||||
return tuple(
|
||||
BedrockContentItem(
|
||||
|
|
@ -255,7 +253,7 @@ def _is_responses_api_route(request_route: str | None) -> bool:
|
|||
if request_route is None:
|
||||
return False
|
||||
call_types: Final = get_call_types_for_route(request_route)
|
||||
return call_types is not None and any(call_type in _RESPONSES_API_CALL_TYPES for call_type in call_types)
|
||||
return call_types is not None and any(call_type in RESPONSES_INPUT_CALL_TYPES for call_type in call_types)
|
||||
|
||||
|
||||
class _JsonResponse(Protocol):
|
||||
|
|
|
|||
|
|
@ -24,6 +24,7 @@ from litellm.integrations.custom_guardrail import (
|
|||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.utils import (
|
||||
RESPONSES_INPUT_CALL_TYPES,
|
||||
Choices,
|
||||
GuardrailStatus,
|
||||
ModelResponse,
|
||||
|
|
@ -402,7 +403,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
|
|||
user_id: Final = self._resolve_user_id_for_blocking(data, user_api_key_dict)
|
||||
|
||||
prompt_text: str | None = None
|
||||
if call_type in ("responses", "aresponses"):
|
||||
if call_type in RESPONSES_INPUT_CALL_TYPES:
|
||||
# Route Responses API calls to the responses-specific extractor
|
||||
# before the generic ``messages`` branch. This mirrors
|
||||
# ``async_logging_hook`` and ensures ``instructions`` (system
|
||||
|
|
@ -663,7 +664,7 @@ class MicrosoftPurviewDLPGuardrail(PurviewGuardrailBase, CustomGuardrail):
|
|||
# Log prompt (uploadText)
|
||||
try:
|
||||
prompt_text: str | None = None
|
||||
if call_type in ("responses", "aresponses"):
|
||||
if call_type in RESPONSES_INPUT_CALL_TYPES:
|
||||
# Responses API: route to the responses-specific extractor
|
||||
# before the generic ``messages`` branch. litellm's logging
|
||||
# pipeline stores the raw responses ``input`` (a string or a
|
||||
|
|
|
|||
|
|
@ -110,6 +110,19 @@ def _translated_call_type(call_type: object, mappings: Mapping[CallTypes, object
|
|||
return known_call_type if known_call_type in mappings else None
|
||||
|
||||
|
||||
def _warn_untranslated_call_type(
|
||||
guardrail_name: str | None, request_route: str | None, call_type: object, stage: str
|
||||
) -> None:
|
||||
verbose_proxy_logger.warning(
|
||||
"Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; "
|
||||
"skipping %s scanning.",
|
||||
guardrail_name,
|
||||
request_route,
|
||||
call_type,
|
||||
stage,
|
||||
)
|
||||
|
||||
|
||||
def resolve_endpoint_translation(
|
||||
user_api_key_dict: UserAPIKeyAuth, first_response_item: object | None
|
||||
) -> "tuple[str, BaseTranslation] | None":
|
||||
|
|
@ -269,7 +282,9 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
|
||||
verbose_proxy_logger.debug("Running UnifiedLLMGuardrails pre-call hook")
|
||||
|
||||
guardrail_to_apply: Final[CustomGuardrail] = data.pop("guardrail_to_apply", None)
|
||||
guardrail_to_apply: Final = cast( # cast-ok: the proxy stores the selected CustomGuardrail under this key
|
||||
"CustomGuardrail | None", data.pop("guardrail_to_apply", None)
|
||||
)
|
||||
if guardrail_to_apply is None:
|
||||
return data
|
||||
|
||||
|
|
@ -287,6 +302,9 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
mappings: Final = load_guardrail_translation_mappings()
|
||||
input_call_type: Final = _translated_call_type(call_type, mappings)
|
||||
if input_call_type is None:
|
||||
_warn_untranslated_call_type(
|
||||
guardrail_to_apply.guardrail_name, user_api_key_dict.request_route, call_type, "pre-call"
|
||||
)
|
||||
return data
|
||||
|
||||
endpoint_translation: Final = _as_endpoint_translation(mappings[input_call_type]())
|
||||
|
|
@ -318,7 +336,9 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
"""
|
||||
verbose_proxy_logger.debug("Running UnifiedLLMGuardrails moderation hook")
|
||||
|
||||
guardrail_to_apply: Final[CustomGuardrail] = data.pop("guardrail_to_apply", None)
|
||||
guardrail_to_apply: Final = cast( # cast-ok: the proxy stores the selected CustomGuardrail under this key
|
||||
"CustomGuardrail | None", data.pop("guardrail_to_apply", None)
|
||||
)
|
||||
if guardrail_to_apply is None:
|
||||
return data
|
||||
|
||||
|
|
@ -336,6 +356,9 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
mappings: Final = load_guardrail_translation_mappings()
|
||||
input_call_type: Final = _translated_call_type(call_type, mappings)
|
||||
if input_call_type is None:
|
||||
_warn_untranslated_call_type(
|
||||
guardrail_to_apply.guardrail_name, user_api_key_dict.request_route, call_type, "during-call"
|
||||
)
|
||||
return data
|
||||
|
||||
endpoint_translation: Final = _as_endpoint_translation(mappings[input_call_type]())
|
||||
|
|
@ -369,7 +392,9 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
)
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
guardrail_to_apply: Final[CustomGuardrail] = data.pop("guardrail_to_apply", None)
|
||||
guardrail_to_apply: Final = cast( # cast-ok: the proxy stores the selected CustomGuardrail under this key
|
||||
"CustomGuardrail | None", data.pop("guardrail_to_apply", None)
|
||||
)
|
||||
|
||||
if guardrail_to_apply is None:
|
||||
return
|
||||
|
|
@ -412,12 +437,8 @@ class UnifiedLLMGuardrails(CustomLogger):
|
|||
mappings: Final = load_guardrail_translation_mappings()
|
||||
|
||||
if CallTypes(call_type) not in mappings:
|
||||
verbose_proxy_logger.warning(
|
||||
"Guardrail '%s' selected for route '%s' but call type '%s' has no guardrail translation handler; "
|
||||
"skipping post-call scanning.",
|
||||
guardrail_to_apply.guardrail_name,
|
||||
user_api_key_dict.request_route,
|
||||
call_type,
|
||||
_warn_untranslated_call_type(
|
||||
guardrail_to_apply.guardrail_name, user_api_key_dict.request_route, call_type, "post-call"
|
||||
)
|
||||
return response
|
||||
|
||||
|
|
|
|||
|
|
@ -797,6 +797,15 @@ MCP_GUARDRAIL_CALL_TYPES: Final[frozenset[str]] = frozenset(
|
|||
{CallTypes.call_mcp_tool.value, CallTypes.list_mcp_tools.value}
|
||||
)
|
||||
|
||||
RESPONSES_INPUT_CALL_TYPES: Final[frozenset[str]] = frozenset(
|
||||
{
|
||||
CallTypes.responses.value,
|
||||
CallTypes.aresponses.value,
|
||||
CallTypes.compact_responses.value,
|
||||
CallTypes.acompact_responses.value,
|
||||
}
|
||||
)
|
||||
|
||||
# Mapping of API routes to their corresponding call types
|
||||
API_ROUTE_TO_CALL_TYPES: Final[Mapping[str, Sequence[CallTypes]]] = {
|
||||
# Chat Completions
|
||||
|
|
|
|||
|
|
@ -3786,3 +3786,18 @@ def test_default_on_guardrail() -> None:
|
|||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"call_type",
|
||||
[CallTypes.responses, CallTypes.aresponses, CallTypes.compact_responses, CallTypes.acompact_responses],
|
||||
)
|
||||
def test_guardrail_messages_for_responses_and_compact_come_from_input(call_type: CallTypes) -> None:
|
||||
guardrail: Final = CustomGuardrail(guardrail_name="test-guardrail")
|
||||
|
||||
messages: Final = guardrail.get_guardrails_messages_for_call_type(
|
||||
call_type=call_type,
|
||||
data={"input": "Remember forbiddenword", "messages": [{"role": "user", "content": "decoy"}]},
|
||||
)
|
||||
|
||||
assert messages == [{"role": "user", "content": "Remember forbiddenword"}]
|
||||
|
|
|
|||
|
|
@ -1652,11 +1652,12 @@ async def test_a_client_sent_response_key_cannot_skip_recording_or_tool_call_che
|
|||
assert calls["/v1/chat/completions"].get("ingest_data") == "true", "the reply is still recorded"
|
||||
|
||||
|
||||
def test_a_decoy_messages_key_cannot_replace_the_responses_api_input(akto_post_call):
|
||||
@pytest.mark.parametrize("call_type", ["aresponses", "acompact_responses"])
|
||||
def test_a_decoy_messages_key_cannot_replace_the_responses_api_input(akto_post_call, call_type):
|
||||
request_data = {
|
||||
"input": [{"role": "user", "content": "the real prompt"}],
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"litellm_logging_obj": SimpleNamespace(call_type="aresponses", model_call_details={}),
|
||||
"litellm_logging_obj": SimpleNamespace(call_type=call_type, model_call_details={}),
|
||||
}
|
||||
|
||||
payload = akto_post_call.build_akto_payload(
|
||||
|
|
|
|||
|
|
@ -530,8 +530,11 @@ async def test_prompt_shield_messages_less_embeddings_return_data_and_log_allow(
|
|||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("call_type", ["aresponses", "acompact_responses"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_responses_input_is_scanned_and_billing_is_logged(responses_input: object, expected_prompt: str) -> None:
|
||||
async def test_responses_input_is_scanned_and_billing_is_logged(
|
||||
responses_input: object, expected_prompt: str, call_type: str
|
||||
) -> None:
|
||||
guardrail: Final = _priced_shield_guardrail(cost_tier="paid", price_per_1000_text_records=0.38)
|
||||
data: Final[dict[str, object]] = {"input": responses_input}
|
||||
|
||||
|
|
@ -540,7 +543,7 @@ async def test_responses_input_is_scanned_and_billing_is_logged(responses_input:
|
|||
user_api_key_dict=UserAPIKeyAuth(api_key="k"),
|
||||
cache=None,
|
||||
data=data,
|
||||
call_type="aresponses",
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
mock_post.assert_called_once()
|
||||
|
|
|
|||
|
|
@ -54,8 +54,9 @@ async def test_azure_text_moderation_guardrail_pre_call_hook():
|
|||
assert mock_async_make_request.call_args.kwargs["text"] == "Hello, how are you?"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("call_type", ["aresponses", "acompact_responses"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_azure_text_moderation_scans_responses_input() -> None:
|
||||
async def test_azure_text_moderation_scans_responses_input(call_type: str) -> None:
|
||||
guardrail: Final = AzureContentSafetyTextModerationGuardrail(
|
||||
guardrail_name="azure_text_moderation",
|
||||
api_key="azure_text_moderation_api_key",
|
||||
|
|
@ -78,7 +79,7 @@ async def test_azure_text_moderation_scans_responses_input() -> None:
|
|||
user_api_key_dict=UserAPIKeyAuth(api_key="azure_text_moderation_api_key"),
|
||||
cache=None,
|
||||
data={"input": "Review this response input"},
|
||||
call_type="aresponses",
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview.base import (
|
|||
from litellm.proxy.guardrails.guardrail_hooks.microsoft_purview.purview_dlp import (
|
||||
MicrosoftPurviewDLPGuardrail,
|
||||
)
|
||||
from litellm.types.utils import CallTypesLiteral
|
||||
|
||||
|
||||
def _make_guardrail(**kwargs) -> MicrosoftPurviewDLPGuardrail:
|
||||
|
|
@ -201,9 +202,7 @@ class TestResolveUserId:
|
|||
def test_end_user_id_after_key_user_id(self):
|
||||
"""When both key user_id and end_user_id exist, key user_id is used first."""
|
||||
guardrail = _make_guardrail()
|
||||
auth = UserAPIKeyAuth(
|
||||
api_key="test", user_id="key-owner", end_user_id="end-user-101"
|
||||
)
|
||||
auth = UserAPIKeyAuth(api_key="test", user_id="key-owner", end_user_id="end-user-101")
|
||||
assert guardrail._resolve_user_id({}, auth) == "key-owner"
|
||||
|
||||
def test_none_when_missing(self):
|
||||
|
|
@ -222,9 +221,7 @@ class TestPreCallHook:
|
|||
async def test_pre_call_allow(self):
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
await guardrail.async_pre_call_hook(
|
||||
|
|
@ -247,9 +244,7 @@ class TestPreCallHook:
|
|||
"litellm_call_id": "call-abc",
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
out = await guardrail.async_pre_call_hook(
|
||||
|
|
@ -265,9 +260,7 @@ class TestPreCallHook:
|
|||
async def test_pre_call_block(self):
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.side_effect = HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": "Microsoft Purview DLP: Content blocked by policy"},
|
||||
|
|
@ -275,9 +268,7 @@ class TestPreCallHook:
|
|||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="test", user_id="user-123"
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"),
|
||||
cache=None,
|
||||
data={
|
||||
"messages": [
|
||||
|
|
@ -296,9 +287,7 @@ class TestPreCallHook:
|
|||
async def test_pre_call_no_user_id_raises(self):
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test"),
|
||||
|
|
@ -314,9 +303,7 @@ class TestPreCallHook:
|
|||
async def test_pre_call_no_messages_skips(self):
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"),
|
||||
cache=None,
|
||||
|
|
@ -333,9 +320,7 @@ class TestPreCallFullTranscript:
|
|||
"""DLP text must include system / prior turns, not only the last user block."""
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
await guardrail.async_pre_call_hook(
|
||||
|
|
@ -371,16 +356,10 @@ class TestPostCallHook:
|
|||
|
||||
guardrail = _make_guardrail()
|
||||
response = ModelResponse(
|
||||
choices=[
|
||||
Choices(
|
||||
index=0, message=Message(content="Safe response", role="assistant")
|
||||
)
|
||||
],
|
||||
choices=[Choices(index=0, message=Message(content="Safe response", role="assistant"))],
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
result = await guardrail.async_post_call_success_hook(
|
||||
|
|
@ -410,9 +389,7 @@ class TestPostCallHook:
|
|||
],
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.side_effect = HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": "Microsoft Purview DLP: Content blocked by policy"},
|
||||
|
|
@ -421,9 +398,7 @@ class TestPostCallHook:
|
|||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.async_post_call_success_hook(
|
||||
data={},
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="test", user_id="user-123"
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"),
|
||||
response=response,
|
||||
)
|
||||
|
||||
|
|
@ -435,14 +410,10 @@ class TestPostCallHook:
|
|||
|
||||
guardrail = _make_guardrail()
|
||||
response = ModelResponse(
|
||||
choices=[
|
||||
Choices(index=0, message=Message(content="Response", role="assistant"))
|
||||
],
|
||||
choices=[Choices(index=0, message=Message(content="Response", role="assistant"))],
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.async_post_call_success_hook(
|
||||
data={},
|
||||
|
|
@ -471,9 +442,7 @@ class TestPostCallHook:
|
|||
],
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
await guardrail.async_post_call_success_hook(
|
||||
|
|
@ -493,9 +462,7 @@ class TestTextCompletionHooks:
|
|||
async def test_pre_call_text_completion_uses_prompt(self):
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
await guardrail.async_pre_call_hook(
|
||||
|
|
@ -522,9 +489,7 @@ class TestTextCompletionHooks:
|
|||
],
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
await guardrail.async_post_call_success_hook(
|
||||
|
|
@ -605,9 +570,7 @@ class TestResponsesAPIHooks:
|
|||
"""Pre-call hook must scan plain-string ``input`` on responses call type."""
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
await guardrail.async_pre_call_hook(
|
||||
|
|
@ -621,21 +584,20 @@ class TestResponsesAPIHooks:
|
|||
assert mock_check.call_args.kwargs["activity"] == "uploadText"
|
||||
assert "SSN: 123-45-6789" in mock_check.call_args.kwargs["text"]
|
||||
|
||||
@pytest.mark.parametrize("call_type", ["aresponses", "acompact_responses"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_aresponses_string_input(self):
|
||||
async def test_pre_call_aresponses_string_input(self, call_type: CallTypesLiteral):
|
||||
"""Pre-call hook must scan ``input`` on ``aresponses`` call type too."""
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"),
|
||||
cache=None,
|
||||
data={"input": "sensitive content"},
|
||||
call_type="aresponses",
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
mock_check.assert_called_once()
|
||||
|
|
@ -646,17 +608,13 @@ class TestResponsesAPIHooks:
|
|||
"""Pre-call hook must extract text from structured list ``input``."""
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"),
|
||||
cache=None,
|
||||
data={
|
||||
"input": [{"role": "user", "content": "Secret phrase: alpha bravo"}]
|
||||
},
|
||||
data={"input": [{"role": "user", "content": "Secret phrase: alpha bravo"}]},
|
||||
call_type="responses",
|
||||
)
|
||||
|
||||
|
|
@ -668,9 +626,7 @@ class TestResponsesAPIHooks:
|
|||
"""Pre-call hook must not call _check_content when ``input`` is absent."""
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"),
|
||||
cache=None,
|
||||
|
|
@ -685,9 +641,7 @@ class TestResponsesAPIHooks:
|
|||
"""Benign string ``input`` must still scan ``instructions`` (system message)."""
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
await guardrail.async_pre_call_hook(
|
||||
|
|
@ -710,9 +664,7 @@ class TestResponsesAPIHooks:
|
|||
"""Requests with only ``instructions`` (no ``input``) must still be scanned."""
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
await guardrail.async_pre_call_hook(
|
||||
|
|
@ -723,10 +675,7 @@ class TestResponsesAPIHooks:
|
|||
)
|
||||
|
||||
mock_check.assert_called_once()
|
||||
assert (
|
||||
"policy text in instructions only"
|
||||
in mock_check.call_args.kwargs["text"]
|
||||
)
|
||||
assert "policy text in instructions only" in mock_check.call_args.kwargs["text"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_responses_api_output_text(self):
|
||||
|
|
@ -743,16 +692,12 @@ class TestResponsesAPIHooks:
|
|||
"id": "msg-1",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "output_text", "text": "card 4111-1111-1111-1111"}
|
||||
],
|
||||
"content": [{"type": "output_text", "text": "card 4111-1111-1111-1111"}],
|
||||
}
|
||||
],
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
result = await guardrail.async_post_call_success_hook(
|
||||
|
|
@ -778,9 +723,7 @@ class TestResponsesAPIHooks:
|
|||
output=[],
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
await guardrail.async_post_call_success_hook(
|
||||
data={},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"),
|
||||
|
|
@ -789,8 +732,9 @@ class TestResponsesAPIHooks:
|
|||
|
||||
mock_check.assert_not_called()
|
||||
|
||||
@pytest.mark.parametrize("call_type", ["responses", "compact_responses"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_logging_hook_responses_api_input_and_output(self):
|
||||
async def test_logging_hook_responses_api_input_and_output(self, call_type: str):
|
||||
"""Logging hook must scan both ``input`` and ``ResponsesAPIResponse.output``."""
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
|
||||
|
|
@ -809,20 +753,16 @@ class TestResponsesAPIHooks:
|
|||
],
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
await guardrail.async_logging_hook(
|
||||
kwargs={
|
||||
"input": "prompt body",
|
||||
"litellm_params": {
|
||||
"metadata": {"user_api_key_user_id": "user-123"}
|
||||
},
|
||||
"litellm_params": {"metadata": {"user_api_key_user_id": "user-123"}},
|
||||
},
|
||||
result=result_response,
|
||||
call_type="responses",
|
||||
call_type=call_type,
|
||||
)
|
||||
|
||||
assert mock_check.call_count == 2
|
||||
|
|
@ -860,9 +800,7 @@ class TestResponsesAPIHooks:
|
|||
],
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
await guardrail.async_logging_hook(
|
||||
|
|
@ -872,9 +810,7 @@ class TestResponsesAPIHooks:
|
|||
# Simulate litellm's logging path which mirrors the raw
|
||||
# responses input under "messages".
|
||||
"messages": "prompt body",
|
||||
"litellm_params": {
|
||||
"metadata": {"user_api_key_user_id": "user-123"}
|
||||
},
|
||||
"litellm_params": {"metadata": {"user_api_key_user_id": "user-123"}},
|
||||
},
|
||||
result=result_response,
|
||||
call_type="aresponses",
|
||||
|
|
@ -883,11 +819,7 @@ class TestResponsesAPIHooks:
|
|||
assert mock_check.call_count == 2
|
||||
activities = {c.kwargs["activity"] for c in mock_check.call_args_list}
|
||||
assert activities == {"uploadText", "downloadText"}
|
||||
upload_calls = [
|
||||
c
|
||||
for c in mock_check.call_args_list
|
||||
if c.kwargs["activity"] == "uploadText"
|
||||
]
|
||||
upload_calls = [c for c in mock_check.call_args_list if c.kwargs["activity"] == "uploadText"]
|
||||
assert len(upload_calls) == 1
|
||||
upload_text = upload_calls[0].kwargs["text"]
|
||||
assert "prompt body" in upload_text
|
||||
|
|
@ -910,10 +842,7 @@ class TestLoggingResolveUserId:
|
|||
}
|
||||
}
|
||||
}
|
||||
assert (
|
||||
guardrail._resolve_user_id_from_logging_kwargs(kwargs)
|
||||
== "trusted-from-proxy"
|
||||
)
|
||||
assert guardrail._resolve_user_id_from_logging_kwargs(kwargs) == "trusted-from-proxy"
|
||||
|
||||
def test_logging_ignores_caller_supplied_user_id_field(self):
|
||||
"""Caller-controlled ``metadata[user_id_field]`` must not drive Purview audit attribution."""
|
||||
|
|
@ -928,9 +857,7 @@ class TestLoggingResolveUserId:
|
|||
"user_api_key_user_id": "from-top-level",
|
||||
"litellm_params": {"metadata": {}},
|
||||
}
|
||||
assert (
|
||||
guardrail._resolve_user_id_from_logging_kwargs(kwargs) == "from-top-level"
|
||||
)
|
||||
assert guardrail._resolve_user_id_from_logging_kwargs(kwargs) == "from-top-level"
|
||||
|
||||
def test_logging_returns_none_when_no_trusted_identity(self):
|
||||
guardrail = _make_guardrail()
|
||||
|
|
@ -1064,9 +991,7 @@ class TestTokenCaching:
|
|||
async def test_token_cached(self):
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail.async_handler, "post", return_value=_mock_token_response()
|
||||
) as mock_post:
|
||||
with patch.object(guardrail.async_handler, "post", return_value=_mock_token_response()) as mock_post:
|
||||
token1 = await guardrail._get_access_token()
|
||||
token2 = await guardrail._get_access_token()
|
||||
|
||||
|
|
@ -1079,9 +1004,7 @@ class TestTokenCaching:
|
|||
async def test_token_refreshed_on_expiry(self):
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail.async_handler, "post", return_value=_mock_token_response()
|
||||
) as mock_post:
|
||||
with patch.object(guardrail.async_handler, "post", return_value=_mock_token_response()) as mock_post:
|
||||
await guardrail._get_access_token()
|
||||
|
||||
# Expire the token
|
||||
|
|
@ -1130,9 +1053,7 @@ class TestGraphPostHttpError:
|
|||
guardrail = _make_guardrail()
|
||||
|
||||
error_resp = Mock()
|
||||
error_resp.json.return_value = {
|
||||
"error": {"code": "Forbidden", "message": "no access"}
|
||||
}
|
||||
error_resp.json.return_value = {"error": {"code": "Forbidden", "message": "no access"}}
|
||||
error_resp.headers = {}
|
||||
error_resp.raise_for_status = Mock(
|
||||
side_effect=httpx.HTTPStatusError(
|
||||
|
|
@ -1143,9 +1064,7 @@ class TestGraphPostHttpError:
|
|||
)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
guardrail, "_get_access_token", new_callable=AsyncMock
|
||||
) as mock_token,
|
||||
patch.object(guardrail, "_get_access_token", new_callable=AsyncMock) as mock_token,
|
||||
patch.object(guardrail.async_handler, "post", return_value=error_resp),
|
||||
):
|
||||
mock_token.return_value = "mock-token"
|
||||
|
|
@ -1161,9 +1080,7 @@ class TestGraphPostHttpError:
|
|||
"""A Graph error on protectionScopes/compute must not be cached as success."""
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_graph_post", new_callable=AsyncMock
|
||||
) as mock_post:
|
||||
with patch.object(guardrail, "_graph_post", new_callable=AsyncMock) as mock_post:
|
||||
mock_post.side_effect = httpx.HTTPStatusError(
|
||||
"429 Too Many Requests",
|
||||
request=httpx.Request("POST", "https://graph.microsoft.com/"),
|
||||
|
|
@ -1187,15 +1104,9 @@ class TestScopeCaching:
|
|||
async def test_scope_cached(self):
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_graph_post", new_callable=AsyncMock
|
||||
) as mock_post:
|
||||
with patch.object(guardrail, "_graph_post", new_callable=AsyncMock) as mock_post:
|
||||
mock_post.return_value = (
|
||||
{
|
||||
"value": [
|
||||
{"activities": "uploadText", "executionMode": "evaluateInline"}
|
||||
]
|
||||
},
|
||||
{"value": [{"activities": "uploadText", "executionMode": "evaluateInline"}]},
|
||||
{"ETag": "scope-etag"},
|
||||
)
|
||||
|
||||
|
|
@ -1213,17 +1124,11 @@ class TestScopeCaching:
|
|||
guardrail._scope_cache_maxsize = 3
|
||||
|
||||
scope_payload = (
|
||||
{
|
||||
"value": [
|
||||
{"activities": "uploadText", "executionMode": "evaluateInline"}
|
||||
]
|
||||
},
|
||||
{"value": [{"activities": "uploadText", "executionMode": "evaluateInline"}]},
|
||||
{"ETag": "scope-etag"},
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_graph_post", new_callable=AsyncMock
|
||||
) as mock_post:
|
||||
with patch.object(guardrail, "_graph_post", new_callable=AsyncMock) as mock_post:
|
||||
mock_post.return_value = scope_payload
|
||||
|
||||
await guardrail._compute_protection_scopes("user-a")
|
||||
|
|
@ -1258,9 +1163,7 @@ class TestScopeCaching:
|
|||
{"ETag": "scope-etag"},
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_graph_post", new_callable=AsyncMock
|
||||
) as mock_post:
|
||||
with patch.object(guardrail, "_graph_post", new_callable=AsyncMock) as mock_post:
|
||||
mock_post.return_value = scope_payload
|
||||
|
||||
# Populate cache: user-a (older), user-b (newer)
|
||||
|
|
@ -1281,18 +1184,14 @@ class TestScopeCaching:
|
|||
assert mock_post.call_count == 4
|
||||
|
||||
assert "user-a" in guardrail._scope_cache, "user-a was wrongly evicted"
|
||||
assert (
|
||||
"user-b" not in guardrail._scope_cache
|
||||
), "user-b should have been evicted"
|
||||
assert "user-b" not in guardrail._scope_cache, "user-b should have been evicted"
|
||||
assert "user-c" in guardrail._scope_cache
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_scope_invalidated_on_modified(self):
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_graph_post", new_callable=AsyncMock
|
||||
) as mock_post:
|
||||
with patch.object(guardrail, "_graph_post", new_callable=AsyncMock) as mock_post:
|
||||
# First call: compute scopes
|
||||
mock_post.return_value = (
|
||||
{"value": []},
|
||||
|
|
@ -1320,9 +1219,7 @@ class TestGetPromptTextForDlp:
|
|||
def test_single_message_no_extra_separator(self):
|
||||
"""A single message is returned as-is (no leading/trailing separator)."""
|
||||
guardrail = _make_guardrail()
|
||||
result = guardrail.get_prompt_text_for_dlp(
|
||||
[{"role": "user", "content": "Hello"}]
|
||||
)
|
||||
result = guardrail.get_prompt_text_for_dlp([{"role": "user", "content": "Hello"}])
|
||||
assert result == "Hello"
|
||||
|
||||
def test_messages_separated_by_double_newline(self):
|
||||
|
|
@ -1542,9 +1439,7 @@ class TestCheckContentApiErrorHandling:
|
|||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert isinstance(exc_info.value.detail, dict)
|
||||
assert "upstream policy evaluation failed" in exc_info.value.detail.get(
|
||||
"error", ""
|
||||
)
|
||||
assert "upstream policy evaluation failed" in exc_info.value.detail.get("error", "")
|
||||
assert "network failure" in exc_info.value.detail.get("exception", "")
|
||||
assert isinstance(exc_info.value.__cause__, RuntimeError)
|
||||
|
||||
|
|
@ -1633,9 +1528,7 @@ class TestCheckContentApiErrorHandling:
|
|||
headers={"Retry-After": "30"},
|
||||
request=httpx.Request("POST", "https://graph.microsoft.com/v1.0/x"),
|
||||
)
|
||||
upstream_err = httpx.HTTPStatusError(
|
||||
"rate limited", request=upstream_resp.request, response=upstream_resp
|
||||
)
|
||||
upstream_err = httpx.HTTPStatusError("rate limited", request=upstream_resp.request, response=upstream_resp)
|
||||
|
||||
with patch.object(
|
||||
guardrail,
|
||||
|
|
@ -1666,9 +1559,7 @@ class TestCheckContentApiErrorHandling:
|
|||
status_code=401,
|
||||
request=httpx.Request("POST", "https://graph.microsoft.com/v1.0/x"),
|
||||
)
|
||||
upstream_err = httpx.HTTPStatusError(
|
||||
"unauthorized", request=upstream_resp.request, response=upstream_resp
|
||||
)
|
||||
upstream_err = httpx.HTTPStatusError("unauthorized", request=upstream_resp.request, response=upstream_resp)
|
||||
|
||||
with patch.object(
|
||||
guardrail,
|
||||
|
|
@ -1723,9 +1614,7 @@ class TestAsyncLoggingHookIndependence:
|
|||
await guardrail.async_logging_hook(
|
||||
kwargs={
|
||||
"messages": [{"role": "user", "content": "prompt"}],
|
||||
"litellm_params": {
|
||||
"metadata": {"user_api_key_user_id": "user-123"}
|
||||
},
|
||||
"litellm_params": {"metadata": {"user_api_key_user_id": "user-123"}},
|
||||
},
|
||||
result=response,
|
||||
call_type="completion",
|
||||
|
|
@ -1762,9 +1651,7 @@ class TestAsyncLoggingHookIndependence:
|
|||
await guardrail.async_logging_hook(
|
||||
kwargs={
|
||||
"messages": [{"role": "user", "content": "prompt"}],
|
||||
"litellm_params": {
|
||||
"metadata": {"user_api_key_user_id": "user-123"}
|
||||
},
|
||||
"litellm_params": {"metadata": {"user_api_key_user_id": "user-123"}},
|
||||
},
|
||||
result=response,
|
||||
call_type="completion",
|
||||
|
|
@ -1924,15 +1811,13 @@ class TestGetPromptTextToolCalls:
|
|||
"""async_pre_call_hook must include tool_call arguments in the text sent to Purview."""
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=__import__(
|
||||
"litellm.proxy._types", fromlist=["UserAPIKeyAuth"]
|
||||
).UserAPIKeyAuth(api_key="test", user_id="user-123"),
|
||||
user_api_key_dict=__import__("litellm.proxy._types", fromlist=["UserAPIKeyAuth"]).UserAPIKeyAuth(
|
||||
api_key="test", user_id="user-123"
|
||||
),
|
||||
cache=None,
|
||||
data={
|
||||
"messages": [
|
||||
|
|
@ -2055,16 +1940,14 @@ class TestCompletionResponseTextPartsToolCalls:
|
|||
],
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
await guardrail.async_post_call_success_hook(
|
||||
data={},
|
||||
user_api_key_dict=__import__(
|
||||
"litellm.proxy._types", fromlist=["UserAPIKeyAuth"]
|
||||
).UserAPIKeyAuth(api_key="test", user_id="user-123"),
|
||||
user_api_key_dict=__import__("litellm.proxy._types", fromlist=["UserAPIKeyAuth"]).UserAPIKeyAuth(
|
||||
api_key="test", user_id="user-123"
|
||||
),
|
||||
response=response,
|
||||
)
|
||||
|
||||
|
|
@ -2152,9 +2035,7 @@ class TestGraphUserIdEncoding:
|
|||
|
||||
guardrail.async_handler.post = AsyncMock(side_effect=_capture_post)
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_get_access_token", new_callable=AsyncMock
|
||||
) as mock_token:
|
||||
with patch.object(guardrail, "_get_access_token", new_callable=AsyncMock) as mock_token:
|
||||
mock_token.return_value = "tok"
|
||||
await guardrail._compute_protection_scopes("user/with%special")
|
||||
|
||||
|
|
@ -2194,9 +2075,7 @@ class TestResolveTrustedUserId:
|
|||
|
||||
def test_trusted_prefers_key_user_id_over_end_user_id(self):
|
||||
guardrail = _make_guardrail()
|
||||
auth = UserAPIKeyAuth(
|
||||
api_key="test", user_id="key-owner", end_user_id="end-user"
|
||||
)
|
||||
auth = UserAPIKeyAuth(api_key="test", user_id="key-owner", end_user_id="end-user")
|
||||
assert guardrail._resolve_trusted_user_id({}, auth) == "key-owner"
|
||||
|
||||
|
||||
|
|
@ -2285,9 +2164,7 @@ class TestTokenIdPromptHandling:
|
|||
"""Pure token-id prompts must be rejected in blocking pre_call mode."""
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"),
|
||||
|
|
@ -2307,9 +2184,7 @@ class TestTokenIdPromptHandling:
|
|||
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
with caplog.at_level(logging.WARNING):
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"),
|
||||
|
|
@ -2326,9 +2201,7 @@ class TestTokenIdPromptHandling:
|
|||
"""Normal string prompts must still be sent to Purview."""
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"),
|
||||
|
|
@ -2346,9 +2219,7 @@ class TestTokenIdPromptHandling:
|
|||
"""Empty/whitespace-only string prompts must not be flagged as token-id prompts."""
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
data = {"prompt": empty_prompt}
|
||||
result = await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"),
|
||||
|
|
@ -2374,9 +2245,7 @@ class TestTokenIdPromptHandling:
|
|||
"""Nested/mixed token-id prompts must also be rejected in blocking pre_call mode."""
|
||||
guardrail = _make_guardrail()
|
||||
|
||||
with patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check:
|
||||
with patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check:
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await guardrail.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="u1"),
|
||||
|
|
@ -2435,12 +2304,8 @@ class TestStreamingIteratorHook:
|
|||
|
||||
with (
|
||||
patch("litellm.main.stream_chunk_builder", return_value=assembled_response),
|
||||
patch(
|
||||
"litellm.llms.base_llm.base_model_iterator.MockResponseIterator"
|
||||
) as mock_iterator_cls,
|
||||
patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check,
|
||||
patch("litellm.llms.base_llm.base_model_iterator.MockResponseIterator") as mock_iterator_cls,
|
||||
patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check,
|
||||
):
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
|
|
@ -2491,18 +2356,15 @@ class TestStreamingIteratorHook:
|
|||
new_callable=AsyncMock,
|
||||
side_effect=HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": "Microsoft Purview DLP: Content blocked by policy"
|
||||
},
|
||||
detail={"error": "Microsoft Purview DLP: Content blocked by policy"},
|
||||
),
|
||||
),
|
||||
):
|
||||
chunks = []
|
||||
|
||||
async def _drain():
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="test", user_id="user-123"
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"),
|
||||
response=fake_response_stream(),
|
||||
request_data={"metadata": {"user_id": "user-123"}},
|
||||
):
|
||||
|
|
@ -2533,10 +2395,9 @@ class TestStreamingIteratorHook:
|
|||
async def fake_response_stream():
|
||||
yield assembled_response
|
||||
|
||||
with patch(
|
||||
"litellm.main.stream_chunk_builder", return_value=assembled_response
|
||||
):
|
||||
with patch("litellm.main.stream_chunk_builder", return_value=assembled_response):
|
||||
chunks = []
|
||||
|
||||
async def _drain():
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test"), # no user_id
|
||||
|
|
@ -2568,9 +2429,7 @@ class TestStreamingIteratorHook:
|
|||
|
||||
with (
|
||||
patch("litellm.main.stream_chunk_builder", return_value=assembled_response),
|
||||
patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check,
|
||||
patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check,
|
||||
):
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
|
|
@ -2626,9 +2485,7 @@ class TestStreamingIteratorHook:
|
|||
|
||||
with (
|
||||
patch("litellm.main.stream_chunk_builder") as mock_stream_builder,
|
||||
patch.object(
|
||||
guardrail, "_check_content", new_callable=AsyncMock
|
||||
) as mock_check,
|
||||
patch.object(guardrail, "_check_content", new_callable=AsyncMock) as mock_check,
|
||||
):
|
||||
mock_check.return_value = {"policyActions": []}
|
||||
|
||||
|
|
@ -2688,11 +2545,10 @@ class TestStreamingIteratorHook:
|
|||
),
|
||||
):
|
||||
chunks = []
|
||||
|
||||
async def _drain():
|
||||
async for chunk in guardrail.async_post_call_streaming_iterator_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="test", user_id="user-123"
|
||||
),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test", user_id="user-123"),
|
||||
response=fake_response_stream(),
|
||||
request_data={},
|
||||
):
|
||||
|
|
@ -2719,7 +2575,4 @@ class TestRegistration:
|
|||
|
||||
assert "microsoft_purview" in guardrail_initializer_registry
|
||||
assert "microsoft_purview" in guardrail_class_registry
|
||||
assert (
|
||||
guardrail_class_registry["microsoft_purview"]
|
||||
is MicrosoftPurviewDLPGuardrail
|
||||
)
|
||||
assert guardrail_class_registry["microsoft_purview"] is MicrosoftPurviewDLPGuardrail
|
||||
|
|
|
|||
|
|
@ -3128,22 +3128,28 @@ async def test_compact_input_is_guardrailed_like_create_while_stored_response_lo
|
|||
@pytest.mark.parametrize("route_call_type", ["avector_store_list", "not_a_call_type"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_route_without_a_call_type_passes_input_guardrails_unscanned(
|
||||
monkeypatch: pytest.MonkeyPatch, hook: str, route_call_type: str
|
||||
monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, hook: str, route_call_type: str
|
||||
) -> None:
|
||||
_patch_translation_mappings(monkeypatch, load_guardrail_translation_mappings())
|
||||
guardrail: Final = RecordingGuardrail()
|
||||
data: Final = {"model": "gpt-4o", "input": "Remember forbiddenword", "guardrail_to_apply": guardrail}
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(token="hashed-key")
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(token="hashed-key", request_route="/v1/vector_stores")
|
||||
|
||||
result: Final = (
|
||||
await UnifiedLLMGuardrails().async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type=route_call_type
|
||||
with caplog.at_level(logging.WARNING):
|
||||
result: Final = (
|
||||
await UnifiedLLMGuardrails().async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict, cache=DualCache(), data=data, call_type=route_call_type
|
||||
)
|
||||
if hook == "async_pre_call_hook"
|
||||
else await UnifiedLLMGuardrails().async_moderation_hook(
|
||||
data=data, user_api_key_dict=user_api_key_dict, call_type=route_call_type
|
||||
)
|
||||
)
|
||||
if hook == "async_pre_call_hook"
|
||||
else await UnifiedLLMGuardrails().async_moderation_hook(
|
||||
data=data, user_api_key_dict=user_api_key_dict, call_type=route_call_type
|
||||
)
|
||||
)
|
||||
|
||||
assert guardrail.apply_calls == []
|
||||
assert result == {"model": "gpt-4o", "input": "Remember forbiddenword"}
|
||||
stage: Final = "pre-call" if hook == "async_pre_call_hook" else "during-call"
|
||||
assert (
|
||||
f"route '/v1/vector_stores' but call type '{route_call_type}' has no guardrail translation handler; "
|
||||
f"skipping {stage} scanning" in caplog.text
|
||||
)
|
||||
|
|
|
|||
|
|
@ -55,6 +55,7 @@ from litellm.llms.openai_like.json_loader import JSONProviderRegistry
|
|||
from litellm.proxy.utils import is_valid_api_key
|
||||
from litellm.responses.main import aresponses, responses
|
||||
from litellm.types.caching import CachingSupportedCallTypes
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY, STREAM_OPTIONS_STASH_KEYS
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.router import CredentialLiteLLMParams, GenericLiteLLMParams, LiteLLM_Params
|
||||
|
|
@ -10920,6 +10921,34 @@ async def test_client_passes_stored_response_call_types_to_both_deployment_hooks
|
|||
assert callback.hook_calls == (("pre", CallTypes(call_type)), ("post", CallTypes(call_type)))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("call_type", ["aget_responses", "acompact_responses"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_deployment_post_call_guardrail_blocks_a_stored_response_reply(
|
||||
call_type: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
class NoPineappleGuardrail(CustomGuardrail):
|
||||
async def async_post_call_success_hook(
|
||||
self, data: dict[str, object], user_api_key_dict: object, response: object
|
||||
) -> object:
|
||||
if "pineapple" in str(response):
|
||||
raise ValueError("pineapple is not allowed")
|
||||
return response
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"callbacks",
|
||||
[NoPineappleGuardrail(guardrail_name="no-pineapple", event_hook=GuardrailEventHooks.post_call)],
|
||||
)
|
||||
|
||||
async def stored_response_call(**kwargs: object) -> dict[str, str]:
|
||||
return {"id": "resp_abc", "output_text": "I like pineapple"}
|
||||
|
||||
stored_response_call.__name__ = call_type
|
||||
|
||||
with pytest.raises(ValueError, match="pineapple is not allowed"):
|
||||
await client(stored_response_call)(model="gpt-4o", response_id="resp_abc", guardrails=["no-pineapple"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"call_type",
|
||||
[
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue