From 6fa4c1757788d8482ac46a7167db0f171bd81ca2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 22:40:28 -0700 Subject: [PATCH] fix(guardrails): skip during-call input guardrails for routes without a call type --- .../unified_guardrail/unified_guardrail.py | 24 ++++--- .../test_unified_guardrail.py | 25 +++++++ tests/unit/test_utils.py | 70 +++++++++++++++++++ tests/unit/types/test_utils.py | 46 +----------- 4 files changed, 111 insertions(+), 54 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index d713f1078a2..be5a2ccaa03 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -102,6 +102,14 @@ def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTran return translation +def _translated_call_type(call_type: object, mappings: Mapping[CallTypes, object]) -> CallTypes | None: + try: + known_call_type: Final = CallTypes(call_type) + except ValueError: + return None + return known_call_type if known_call_type in mappings else None + + def resolve_endpoint_translation( user_api_key_dict: UserAPIKeyAuth, first_response_item: object | None ) -> "tuple[str, BaseTranslation] | None": @@ -277,14 +285,11 @@ class UnifiedLLMGuardrails(CustomLogger): return data mappings: Final = load_guardrail_translation_mappings() + input_call_type: Final = _translated_call_type(call_type, mappings) + if input_call_type is None: + return data - try: - if CallTypes(call_type) not in mappings: - return data - except ValueError: - return data # handle unmapped call types - - endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) + endpoint_translation: Final = _as_endpoint_translation(mappings[input_call_type]()) _ensure_litellm_metadata(data, user_api_key_dict) @@ -329,10 +334,11 @@ class UnifiedLLMGuardrails(CustomLogger): return data mappings: Final = load_guardrail_translation_mappings() - if call_type is not None and CallTypes(call_type) not in mappings: + input_call_type: Final = _translated_call_type(call_type, mappings) + if input_call_type is None: return data - endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) + endpoint_translation: Final = _as_endpoint_translation(mappings[input_call_type]()) _ensure_litellm_metadata(data, user_api_key_dict) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 76fea7c51bb..279c10153d3 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -3122,3 +3122,28 @@ async def test_compact_input_is_guardrailed_like_create_while_stored_response_lo ) assert [call["inputs"]["texts"] for call in guardrail.apply_calls] == scanned_texts + + +@pytest.mark.parametrize("hook", ["async_pre_call_hook", "async_moderation_hook"]) +@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 +) -> 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") + + 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 + ) + ) + + assert guardrail.apply_calls == [] + assert result == {"model": "gpt-4o", "input": "Remember forbiddenword"} diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 146e2df6413..2adf9f45f0b 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -98,6 +98,7 @@ from litellm.utils import ( _snapshot_exception_for_hook, async_post_call_failure_deployment_hook, async_post_call_success_deployment_hook, + async_pre_call_deployment_hook, calculate_max_parallel_requests, client, get_non_default_completion_params, @@ -10882,3 +10883,72 @@ def test_get_optional_params_forwards_a_tool_config_cache_point_to_a_downstream_ ) assert optional_params["extra_body"] == {"kept": True, "cache_control_injection_points": [point]} + + +@pytest.mark.parametrize("call_type", ["aget_responses", "adelete_responses", "acancel_responses", "acompact_responses"]) +@pytest.mark.asyncio +async def test_client_passes_stored_response_call_types_to_both_deployment_hooks( + call_type: str, monkeypatch: pytest.MonkeyPatch +) -> None: + class RecordingLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.hook_calls: tuple[tuple[str, CallTypes | None], ...] = () + + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object] | None: + self.hook_calls = (*self.hook_calls, ("pre", call_type)) + return None + + async def async_post_call_success_deployment_hook( + self, request_data: dict, response: object, call_type: CallTypes | None + ) -> object: + self.hook_calls = (*self.hook_calls, ("post", call_type)) + return None + + callback: Final = RecordingLogger() + monkeypatch.setattr(litellm, "callbacks", [callback]) + + async def stored_response_call(**kwargs: object) -> dict[str, str]: + return {"id": "resp_abc"} + + stored_response_call.__name__ = call_type + + await client(stored_response_call)(model="gpt-4o", response_id="resp_abc") + + assert callback.hook_calls == (("pre", CallTypes(call_type)), ("post", CallTypes(call_type))) + + +@pytest.mark.parametrize( + "call_type", + [ + "get_responses", + "aget_responses", + "delete_responses", + "adelete_responses", + "cancel_responses", + "acancel_responses", + "compact_responses", + "acompact_responses", + ], +) +@pytest.mark.asyncio +async def test_async_pre_call_deployment_hook_passes_responses_call_type( + call_type: str, monkeypatch: pytest.MonkeyPatch +) -> None: + class RecordingLogger(CustomLogger): + received_call_type: CallTypes | None = None + + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: CallTypes | None + ) -> dict[str, object] | None: + self.received_call_type = call_type + return None + + callback: Final = RecordingLogger() + monkeypatch.setattr(litellm, "callbacks", [callback]) + + await async_pre_call_deployment_hook({}, call_type) + + assert callback.received_call_type == call_type diff --git a/tests/unit/types/test_utils.py b/tests/unit/types/test_utils.py index 4f03b7dc712..32479c62372 100644 --- a/tests/unit/types/test_utils.py +++ b/tests/unit/types/test_utils.py @@ -1,17 +1,7 @@ -from typing import Final - import pytest -import litellm -from litellm import utils -from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.initialize_dynamic_callback_params import validate_excluded_services_value -from litellm.types.utils import ( - CAPTURE_MESSAGE_CONTENT_VALUES, - CallTypes, - CaptureMessageContent, - captures_span_content, -) +from litellm.types.utils import CAPTURE_MESSAGE_CONTENT_VALUES, CaptureMessageContent, captures_span_content @pytest.mark.parametrize("value", ["", " ", "auth", "none,redis", "redis,", ",postgres"]) @@ -55,37 +45,3 @@ def test_capture_message_content_constructor_accepts_any_string() -> None: ) def test_captures_span_content(mode: str | None, captures: bool) -> None: assert captures_span_content(mode) is captures - - -@pytest.mark.parametrize( - "call_type", - [ - "get_responses", - "aget_responses", - "delete_responses", - "adelete_responses", - "cancel_responses", - "acancel_responses", - "compact_responses", - "acompact_responses", - ], -) -@pytest.mark.asyncio -async def test_async_pre_call_deployment_hook_passes_responses_call_type( - call_type: str, monkeypatch: pytest.MonkeyPatch -) -> None: - class RecordingLogger(CustomLogger): - received_call_type: CallTypes | None = None - - async def async_pre_call_deployment_hook( - self, kwargs: dict[str, object], call_type: CallTypes | None - ) -> dict[str, object] | None: - self.received_call_type = call_type - return None - - callback: Final = RecordingLogger() - monkeypatch.setattr(litellm, "callbacks", [callback]) - - await utils.async_pre_call_deployment_hook({}, call_type) - - assert callback.received_call_type == call_type