From 381cba33eb44064c250f1b44b4c048eacc71646f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 07:20:57 -0700 Subject: [PATCH] fix(guardrails): scan decision and extraction request text with native-hook guardrails Guardrails that keep their own lifecycle hooks (hide-secrets, Model Armor, Azure Prompt Shield, Azure Text Moderation, LLM Shield Proxy) read messages, input, or prompt, so they found no text on /v1/extraction and /v1/systemone and missed the question instructions on /v1/decisions. On those call types they now run through apply_guardrail, a new decisions translation scans the input or state plus each question's instructions, and a requested guardrail is no longer listed as applied unless it actually scanned. Hide-secrets keeps its per-key opt-out on the unified path --- .../enterprise_callbacks/secret_detection.py | 37 ++- litellm/integrations/custom_guardrail.py | 19 ++ .../guardrail_translation/base_translation.py | 2 + .../guardrail_translation/__init__.py | 13 + .../guardrail_translation/handler.py | 168 ++++++++++++ .../unified_guardrail/unified_guardrail.py | 6 +- .../proxy/policy_engine/pipeline_executor.py | 4 +- litellm/proxy/utils.py | 69 +++-- .../test_secret_detection.py | 41 ++- .../integrations/test_custom_guardrail.py | 38 +++ .../guardrail_translation/__init__.py | 0 .../test_decisions_guardrail_handler.py | 176 +++++++++++++ .../policy_engine/test_pipeline_executor.py | 20 ++ .../test_proxy_logging_hook_detection.py | 248 ++++++++++++++++++ .../proxy_logging/test_guardrail_pipeline.py | 8 +- .../utils/proxy_logging/test_pre_call_hook.py | 2 +- 16 files changed, 809 insertions(+), 42 deletions(-) create mode 100644 litellm/llms/openai/decisions/guardrail_translation/__init__.py create mode 100644 litellm/llms/openai/decisions/guardrail_translation/handler.py create mode 100644 tests/unit/llms/openai/decisions/guardrail_translation/__init__.py create mode 100644 tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py index 65e4d50587f..35c0c7dd119 100644 --- a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py @@ -19,7 +19,7 @@ import json import re import tempfile from collections import OrderedDict -from collections.abc import Generator, Iterator, Sequence +from collections.abc import Generator, Iterator, Mapping, Sequence from contextvars import ContextVar from functools import partial from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional @@ -600,6 +600,24 @@ def _next_run_state(budget: int, state: tuple[int, int, int], size: int) -> tupl return run_id, count + 1, total + size +_KEY_PERMISSIONS: Final[TypeAdapter[Mapping[object, object] | None]] = TypeAdapter(Mapping[object, object] | None) +_METADATA_BUCKET: Final = TypeAdapter(Mapping[str, object]) + + +def _key_opted_out(key_auth: object) -> bool: + if not isinstance(key_auth, UserAPIKeyAuth): + return False + permissions: Final = _KEY_PERMISSIONS.validate_python(key_auth.permissions) + return permissions is not None and permissions.get(GUARDRAIL_NAME) is False + + +def _request_key_opted_out(data: Mapping[str, object]) -> bool: + return any( + isinstance(bucket, dict) and _key_opted_out(_METADATA_BUCKET.validate_python(bucket).get("user_api_key_auth")) + for bucket in (data.get("metadata"), data.get("litellm_metadata")) + ) + + def _value_in_texts(value: str, texts: Sequence[str]) -> str | None: """Return the detected value as it appears in ``texts``, decoding it when it came from a JSON-escaped key line.""" candidates: Final = (value, _json_unescaped(value)) @@ -608,7 +626,7 @@ def _value_in_texts(value: str, texts: Sequence[str]) -> str | None: class _ENTERPRISE_SecretDetection(CustomGuardrail): # Keeps proxied traffic on async_pre_call_hook (the unified apply_guardrail - # path skips should_run_check and never sees data["prompt"]). + # path never sees data["prompt"]). use_native_lifecycle_hooks: ClassVar[bool] = True def __init__(self, detect_secrets_config: Optional[dict] = None, **kwargs): @@ -693,12 +711,12 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): return map_tool_json_text(payload, lambda text: pattern.sub("[REDACTED]", text)) async def should_run_check(self, user_api_key_dict: UserAPIKeyAuth) -> bool: - if user_api_key_dict.permissions is not None: - if GUARDRAIL_NAME in user_api_key_dict.permissions: - if user_api_key_dict.permissions[GUARDRAIL_NAME] is False: - return False + return not _key_opted_out(user_api_key_dict) - return True + def should_run_guardrail(self, data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool: + return super().should_run_guardrail( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.should_run_guardrail leaves data untyped + data, event_type + ) and not _request_key_opted_out(data) @log_guardrail_information async def apply_guardrail( @@ -709,8 +727,9 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail): logging_obj: Optional["LiteLLMLoggingObj"] = None, ) -> GenericGuardrailAPIInputs: """Unified-interface entrypoint, used by /guardrails/apply_guardrail - (the UI test playground). Proxied traffic keeps using - ``async_pre_call_hook``, see ``use_native_lifecycle_hooks``.""" + (the UI test playground) and by the decision and extraction routes, + whose request text the native hook never reads. Other proxied traffic + keeps using ``async_pre_call_hook``, see ``use_native_lifecycle_hooks``.""" texts = inputs.get("texts") if not texts or not any(texts): return inputs diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 9e2dc902d55..4418ac20895 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -72,6 +72,25 @@ _PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16) _GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422}) +APPLY_GUARDRAIL_ONLY_CALL_TYPES: Final[frozenset[str]] = frozenset( + call_type.value + for call_type in ( + CallTypes.extraction, + CallTypes.aextraction, + CallTypes.decisions, + CallTypes.adecisions, + CallTypes.systemone, + CallTypes.asystemone, + ) +) + + +def runs_through_apply_guardrail(guardrail: "CustomGuardrail", call_type: str | None) -> bool: + return "apply_guardrail" in type(guardrail).__dict__ and ( + not guardrail.use_native_lifecycle_hooks or call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES + ) + + DEFAULT_ADVISORY_MESSAGE: Final = ( "The user's latest message was flagged for {reason} by a content safety " "guardrail. This may be a false positive. Use your judgment: respond " diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 6399ded9555..104c82812e9 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -72,6 +72,8 @@ class BaseTranslation(ABC): has the legacy post-call hook against that response, so on a translation without it such a guardrail keeps running on its own.""" + scans_responses: ClassVar[bool] = True + def post_call_hook_response(self, response: object) -> object: """The ``response`` this endpoint's non-streaming post-call hooks receive, derived from the object the translation stores under ``request_data["response"]`` while scanning an diff --git a/litellm/llms/openai/decisions/guardrail_translation/__init__.py b/litellm/llms/openai/decisions/guardrail_translation/__init__.py new file mode 100644 index 00000000000..3afe325608d --- /dev/null +++ b/litellm/llms/openai/decisions/guardrail_translation/__init__.py @@ -0,0 +1,13 @@ +from typing import Final + +from litellm.llms.openai.decisions.guardrail_translation.handler import DecisionsHandler, SystemOneHandler +from litellm.types.utils import CallTypes + +guardrail_translation_mappings: Final = { + CallTypes.decisions: DecisionsHandler, + CallTypes.adecisions: DecisionsHandler, + CallTypes.systemone: SystemOneHandler, + CallTypes.asystemone: SystemOneHandler, +} + +__all__ = ["DecisionsHandler", "SystemOneHandler", "guardrail_translation_mappings"] diff --git a/litellm/llms/openai/decisions/guardrail_translation/handler.py b/litellm/llms/openai/decisions/guardrail_translation/handler.py new file mode 100644 index 00000000000..0244b44f688 --- /dev/null +++ b/litellm/llms/openai/decisions/guardrail_translation/handler.py @@ -0,0 +1,168 @@ +from collections.abc import Mapping +from itertools import chain +from types import MappingProxyType +from typing import TYPE_CHECKING, ClassVar, Final + +from pydantic import JsonValue, TypeAdapter + +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.proxy.guardrails._content_utils import iter_tool_json_text, map_tool_json_text +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth + +_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_QUESTION_LIST_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, JsonValue], ...]]] = TypeAdapter( + tuple[Mapping[str, JsonValue], ...] +) +_QUESTION_MAP_ADAPTER: Final[TypeAdapter[Mapping[str, Mapping[str, JsonValue]]]] = TypeAdapter( + Mapping[str, Mapping[str, JsonValue]] +) + + +def _input_text(part: JsonValue) -> str | None: + if not isinstance(part, dict) or part.get("type") != "input_text": + return None + text: Final = part.get("text") + return text if isinstance(text, str) else None + + +def _content_texts(content: JsonValue) -> tuple[str, ...]: + if isinstance(content, list): + return tuple(text for part in content if (text := _input_text(part))) + return (content,) if isinstance(content, str) and content else () + + +def _message_texts(message: JsonValue) -> tuple[str, ...]: + return _content_texts(message.get("content")) if isinstance(message, dict) else () + + +def _input_texts(decision_input: JsonValue) -> tuple[str, ...]: + if not isinstance(decision_input, list): + return _content_texts(decision_input) + return tuple(chain.from_iterable(_message_texts(message) for message in decision_input)) + + +def _guardrailed_part(part: JsonValue, guardrailed: Mapping[str, str]) -> JsonValue: + text: Final = _input_text(part) + if text is None or not isinstance(part, dict): + return part + return {**part, "text": guardrailed.get(text, text)} + + +def _guardrailed_content(content: JsonValue, guardrailed: Mapping[str, str]) -> JsonValue: + if isinstance(content, list): + return [_guardrailed_part(part, guardrailed) for part in content] + return guardrailed.get(content, content) if isinstance(content, str) else content + + +def _guardrailed_message(message: JsonValue, guardrailed: Mapping[str, str]) -> JsonValue: + return ( + {**message, "content": _guardrailed_content(message["content"], guardrailed)} + if isinstance(message, dict) and "content" in message + else message + ) + + +def _guardrailed_input(decision_input: JsonValue, guardrailed: Mapping[str, str]) -> JsonValue: + if not isinstance(decision_input, list): + return _guardrailed_content(decision_input, guardrailed) + return [_guardrailed_message(message, guardrailed) for message in decision_input] + + +def _guardrailed_json(value: JsonValue, guardrailed: Mapping[str, str]) -> JsonValue: + return map_tool_json_text(value, lambda text: guardrailed.get(text, text)) + + +def _instruction_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]: + return tuple(iter_tool_json_text(question.get("instructions"))) + + +def _with_guardrailed_instructions( + question: Mapping[str, JsonValue], guardrailed: Mapping[str, str] +) -> Mapping[str, JsonValue]: + if "instructions" not in question: + return question + return {**question, "instructions": _guardrailed_json(question["instructions"], guardrailed)} + + +async def _guardrailed_texts( + texts: tuple[str, ...], + guardrail_to_apply: "CustomGuardrail", + request_data: dict[str, object], # mutable-ok: CustomGuardrail.apply_guardrail takes the request as a plain dict + litellm_logging_obj: "LiteLLMLoggingObj | None", +) -> Mapping[str, str]: + unique_texts: Final = tuple(dict.fromkeys(texts)) + model: Final = request_data.get("model") + result: Final = await guardrail_to_apply.apply_guardrail( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.apply_guardrail types request_data as a bare dict + inputs=GenericGuardrailAPIInputs(texts=list(unique_texts), model=model if isinstance(model, str) else None), + request_data=request_data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + returned: Final = result.get("texts", []) + return MappingProxyType( + {text: returned[index] if index < len(returned) else text for index, text in enumerate(unique_texts)} + ) + + +class _DecisionAnswersTranslation(BaseTranslation): + scans_responses: ClassVar[bool] = False + + async def process_output_response( + self, + response: object, + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, + request_data: dict[str, object] | None = None, # mutable-ok: BaseTranslation hands over the request dict + ) -> object: + return response + + +class DecisionsHandler(_DecisionAnswersTranslation): + async def process_input_messages( + self, + data: dict[str, object], # mutable-ok: BaseTranslation hands over the proxy's request dict + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> Mapping[str, object]: + decision_input: Final = _JSON_ADAPTER.validate_python(data.get("input")) + questions: Final = _QUESTION_LIST_ADAPTER.validate_python(data.get("questions") or ()) + texts: Final = (*_input_texts(decision_input), *chain.from_iterable(map(_instruction_texts, questions))) + if not texts: + return data + guardrailed: Final = await _guardrailed_texts(texts, guardrail_to_apply, data, litellm_logging_obj) + return { + **data, + "input": _guardrailed_input(decision_input, guardrailed), + "questions": [_with_guardrailed_instructions(question, guardrailed) for question in questions], + } + + +class SystemOneHandler(_DecisionAnswersTranslation): + async def process_input_messages( + self, + data: dict[str, object], # mutable-ok: BaseTranslation hands over the proxy's request dict + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> Mapping[str, object]: + state: Final = _JSON_ADAPTER.validate_python(data.get("state")) + questions: Final = _QUESTION_MAP_ADAPTER.validate_python(data.get("questions") or {}) + texts: Final = ( + *iter_tool_json_text(state), + *chain.from_iterable(map(_instruction_texts, questions.values())), + ) + if not texts: + return data + guardrailed: Final = await _guardrailed_texts(texts, guardrail_to_apply, data, litellm_logging_obj) + return { + **data, + "state": _guardrailed_json(state, guardrailed), + "questions": { + name: _with_guardrailed_instructions(question, guardrailed) for name, question in questions.items() + }, + } 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..a147580a1f9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -415,7 +415,11 @@ class UnifiedLLMGuardrails(CustomLogger): ) return response - endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) + translation: Final = mappings[CallTypes(call_type)]() + if not translation.scans_responses: + return response + + endpoint_translation: Final = _as_endpoint_translation(translation) try: response = await endpoint_translation.process_output_response( diff --git a/litellm/proxy/policy_engine/pipeline_executor.py b/litellm/proxy/policy_engine/pipeline_executor.py index 4ad29bd3478..09ea1e9d836 100644 --- a/litellm/proxy/policy_engine/pipeline_executor.py +++ b/litellm/proxy/policy_engine/pipeline_executor.py @@ -19,6 +19,7 @@ from litellm.constants import LOGS_GUARDRAIL_INFORMATION_MARKER from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, + runs_through_apply_guardrail, ) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( @@ -573,9 +574,8 @@ class PipelineExecutor: hook_input, scans_raw_request = _prepare_hook_input(step, callback, data, raw_request_snapshot) snapshot_entries_before: Final = len(recorded_guardrail_information(hook_input)) - # Use unified_guardrail path if callback implements apply_guardrail target: CustomLogger = callback - use_unified: Final = PipelineExecutor.supports_unified_execution(callback) + use_unified: Final = runs_through_apply_guardrail(callback, call_type) if use_unified and streaming_chunks is None: hook_input["guardrail_to_apply"] = callback target = UnifiedLLMGuardrails() diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 0ddfc5ce62f..ede7dcd9243 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -144,8 +144,10 @@ from litellm.exceptions import ( SensitiveDataRouteException, ) from litellm.integrations.custom_guardrail import ( + APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail, ModifyResponseException, + runs_through_apply_guardrail, ) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger @@ -342,6 +344,7 @@ _VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0 _VIEW_SETUP_GATE_TABLE: Final = '"LiteLLM_SpendLogs"' _VIEW_SETUP_GATE_PROBE_ROWS: Final = TypeAdapter(tuple[Mapping[str, bool], ...]) _DELEGATED_ATTRIBUTE: Final = TypeAdapter(object) +_REQUESTED_GUARDRAILS: Final = TypeAdapter(tuple[object, ...]) _ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"] _ViewSetupAttempt: TypeAlias = Literal["ready", "table_missing"] | Exception @@ -726,6 +729,24 @@ def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple return (guardrails, others) +def _warn_on_guardrails_that_cannot_scan(requested: tuple[object, ...], call_type: str) -> None: + requested_names: Final = frozenset(name for name in requested if isinstance(name, str)) + unscannable: Final = tuple( + callback.guardrail_name + for callback in ProxyLogging.callback_capabilities().resolved_callbacks + if isinstance(callback, CustomGuardrail) + and callback.guardrail_name is not None + and (callback.default_on or callback.guardrail_name in requested_names) + and "apply_guardrail" not in type(callback).__dict__ + ) + if unscannable: + verbose_proxy_logger.warning( + "Guardrails %s do not implement apply_guardrail, so they cannot scan the text of %s requests", + ", ".join(unscannable), + call_type, + ) + + def _merge_pipeline_metadata_bucket(data: dict, bucket_key: str, modified_bucket_value: object) -> None: if not isinstance(modified_bucket_value, dict): return @@ -2369,11 +2390,7 @@ class ProxyLogging: Returns: Result from the guardrail execution """ - # Use unified_guardrail if callback has apply_guardrail method - has_apply_guardrail: Final = "apply_guardrail" in type(callback).__dict__ and not getattr( - callback, "use_native_lifecycle_hooks", False - ) - use_unified: Final = has_apply_guardrail and not ( + use_unified: Final = runs_through_apply_guardrail(callback, call_type) and not ( hook_type == "during_call" and getattr(callback, "use_native_during_call_hook", False) ) if use_unified: @@ -2698,7 +2715,7 @@ class ProxyLogging: data.pop("prompt_version", None) data.pop("prompt_environment", None) - def _process_guardrail_metadata(self, data: dict) -> None: + def _process_guardrail_metadata(self, data: dict, call_type: CallTypesLiteral) -> None: """Process guardrails from metadata and add to applied_guardrails.""" from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, @@ -2713,6 +2730,15 @@ class ProxyLogging: elif isinstance(metadata_litellm, dict) and "guardrails" in metadata_litellm: guardrails_in_metadata = metadata_litellm.get("guardrails", []) + if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + _warn_on_guardrails_that_cannot_scan( + _REQUESTED_GUARDRAILS.validate_python(guardrails_in_metadata) + if isinstance(guardrails_in_metadata, list) + else (), + call_type, + ) + return + if guardrails_in_metadata and isinstance(guardrails_in_metadata, list): applied_guardrails = [] if isinstance(metadata_standard, dict) and "applied_guardrails" in metadata_standard: @@ -3014,7 +3040,7 @@ class ProxyLogging: and (guardrails_only or not caps.has_pre_call_override) ): if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, call_type) return data parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = ( @@ -3118,17 +3144,17 @@ class ProxyLogging: ) if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, call_type) return data except SensitiveDataRouteException as e: data = await self._handle_sensitive_data_route_exception(e, data, user_api_key_dict) if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, call_type) return data except Exception: if data is not None: - self._process_guardrail_metadata(data) + self._process_guardrail_metadata(data, call_type) raise async def _run_parallel_pre_call_guardrails( @@ -3445,14 +3471,14 @@ class ProxyLogging: _callback_capabilities = callback_capabilities @staticmethod - def _stream_requires_guardrail_translation(user_api_key_dict: UserAPIKeyAuth) -> bool: + def _route_call_type(user_api_key_dict: UserAPIKeyAuth) -> CallTypes | None: route: Final = user_api_key_dict.request_route - if not route: - return False - call_types: Final = get_call_types_for_route(route) - if not call_types: - return False - return call_types[0] in NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES + call_types: Final = get_call_types_for_route(route) if route else None + return call_types[0] if call_types else None + + @staticmethod + def _stream_requires_guardrail_translation(user_api_key_dict: UserAPIKeyAuth) -> bool: + return ProxyLogging._route_call_type(user_api_key_dict) in NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES @staticmethod def has_post_call_response_headers_callbacks() -> bool: @@ -3568,8 +3594,7 @@ class ProxyLogging: call_type: CallTypesLiteral, ) -> None: if ( - "apply_guardrail" in type(callback).__dict__ - and not callback.use_native_lifecycle_hooks + runs_through_apply_guardrail(callback, call_type) and user_api_key_dict is not None and not callback.use_native_during_call_hook ): @@ -4137,6 +4162,7 @@ class ProxyLogging: pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call") guardrail_callbacks, other_callbacks = _partition_post_call_callbacks() + call_type: Final = ProxyLogging._route_call_type(user_api_key_dict) try: # Merge model-level guardrails before checking which guardrails to run guardrail_data: Final = check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) @@ -4168,7 +4194,7 @@ class ProxyLogging: guardrail_response: Any | None = None - if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks: + if runs_through_apply_guardrail(callback, call_type): data["guardrail_to_apply"] = callback guardrail_response = await self._run_guardrail_with_metrics( callback, @@ -4239,11 +4265,12 @@ class ProxyLogging: immediately before awaiting, and the unified hook pops it before its first suspension point, so concurrent guardrails never race on that key. """ + call_type: Final = ProxyLogging._route_call_type(user_api_key_dict) async def _run_one(callback: CustomGuardrail) -> None: if callback.should_run_guardrail(data=guardrail_data, event_type=GuardrailEventHooks.post_call) is not True: return - if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks: + if runs_through_apply_guardrail(callback, call_type): data["guardrail_to_apply"] = callback await self._run_guardrail_with_metrics( callback, diff --git a/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py index 9816eb91102..1b591b54304 100644 --- a/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py +++ b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py @@ -6,8 +6,9 @@ Covers the three defects from the ticket: Spend Logs / the guardrails monitor show activity, with hits ("mask" + masked_entity_count) distinguishable from clean requests ("allow"). - Defining ``apply_guardrail`` must NOT reroute proxied traffic off the - native ``async_pre_call_hook`` (per-key opt-out and ``data["prompt"]`` - handling live only on the native path). + native ``async_pre_call_hook`` (``data["prompt"]`` handling lives only on + the native path), except on the decision and extraction routes, whose + request text the native hook never reads. """ import cProfile @@ -1134,13 +1135,45 @@ async def test_pre_call_hook_still_redacts_text_completion_prompt(): def test_proxied_traffic_stays_on_native_hooks(): """Implementing apply_guardrail must not reroute proxied requests onto the - unified path: that path skips ``should_run_check`` (per-key opt-out) and - never sees ``data["prompt"]``.""" + unified path: that path never sees ``data["prompt"]``.""" guardrail = _guardrail() assert guardrail.uses_apply_guardrail_interface() is True assert guardrail._deployment_hook_target() is guardrail +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("permissions", "expected_state", "listed"), + ( + ({}, "my key is [REDACTED]", True), + ({"hide_secrets": False}, f"my key is {AWS_KEY}", False), + ), +) +async def test_proxy_pre_call_hook_redacts_extraction_state_unless_the_key_opted_out( + monkeypatch, permissions, expected_state, listed +): + import litellm + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "callbacks", [_guardrail()]) + key = UserAPIKeyAuth(api_key="sk-9876", permissions=permissions) + data = { + "model": "invoice-reader", + "state": f"my key is {AWS_KEY}", + "fields": {"total": {"type": "number"}}, + "metadata": {"user_api_key_auth": key}, + } + + result = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=key, data=data, call_type="aextraction" + ) + + assert result["state"] == expected_state + buckets = (result.get("metadata") or {}, result.get("litellm_metadata") or {}) + applied = [*buckets[0].get("applied_guardrails", []), *buckets[1].get("applied_guardrails", [])] + assert ("hide-secrets" in applied) is listed + + @pytest.mark.asyncio async def test_apply_guardrail_without_texts_records_nothing(): """No inputs means nothing was inspected, so no "allow" row is recorded. diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 614ff97553a..16068cd3735 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -15,6 +15,7 @@ from litellm.integrations.custom_guardrail import ( _request_is_streaming, guardrail_request_data_with_streaming, log_guardrail_information, + runs_through_apply_guardrail, without_server_streaming_classification, ) from litellm.litellm_core_utils.litellm_logging import Logging @@ -3786,3 +3787,40 @@ def test_default_on_guardrail() -> None: ) is True ) + + +class _NativeHooksWithApplyGuardrail(_UndecoratedGuardrail): + use_native_lifecycle_hooks: ClassVar[bool] = True + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + return inputs + + +@pytest.mark.parametrize( + ("guardrail_type", "call_type", "expected"), + ( + (_UndecoratedGuardrail, "acompletion", True), + (_UndecoratedGuardrail, "aextraction", True), + (_NativeHooksWithApplyGuardrail, "acompletion", False), + (_NativeHooksWithApplyGuardrail, "aresponses", False), + (_NativeHooksWithApplyGuardrail, None, False), + (_NativeHooksWithApplyGuardrail, "extraction", True), + (_NativeHooksWithApplyGuardrail, "aextraction", True), + (_NativeHooksWithApplyGuardrail, "decisions", True), + (_NativeHooksWithApplyGuardrail, "adecisions", True), + (_NativeHooksWithApplyGuardrail, "systemone", True), + (_NativeHooksWithApplyGuardrail, "asystemone", True), + (_NativeApplyGuardrail, "aextraction", False), + (_InheritedApplyGuardrail, "aextraction", False), + ), +) +def test_native_hook_guardrails_take_apply_guardrail_only_where_their_hooks_read_no_text( + guardrail_type: type[CustomGuardrail], call_type: str | None, expected: bool +) -> None: + assert runs_through_apply_guardrail(guardrail_type(guardrail_name="g"), call_type) is expected diff --git a/tests/unit/llms/openai/decisions/guardrail_translation/__init__.py b/tests/unit/llms/openai/decisions/guardrail_translation/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py b/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py new file mode 100644 index 00000000000..a2993ac1cea --- /dev/null +++ b/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py @@ -0,0 +1,176 @@ +from typing import Final, Literal + +import pytest + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms import get_guardrail_translation_mapping +from litellm.llms.openai.decisions.guardrail_translation.handler import DecisionsHandler, SystemOneHandler +from litellm.types.decisions import OpenAIDecisionResponse +from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs + + +class _TaggingGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="tagging") + self.seen: list[tuple[str, tuple[str, ...]]] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + texts: Final = tuple(inputs.get("texts", [])) + self.seen.append((input_type, texts)) + return GenericGuardrailAPIInputs(texts=[f"{text} [GUARDRAILED]" for text in texts]) + + +class _ReturnsNoTexts(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + return GenericGuardrailAPIInputs() + + +@pytest.mark.parametrize( + ("call_type", "handler"), + ( + (CallTypes.decisions, DecisionsHandler), + (CallTypes.adecisions, DecisionsHandler), + (CallTypes.systemone, SystemOneHandler), + (CallTypes.asystemone, SystemOneHandler), + ), +) +def test_decision_call_types_resolve_to_their_handler(call_type: CallTypes, handler: type[object]) -> None: + assert get_guardrail_translation_mapping(call_type) is handler + + +@pytest.mark.asyncio +async def test_decision_message_text_parts_and_question_instructions_are_guardrailed() -> None: + guardrail: Final = _TaggingGuardrail() + image: Final = {"type": "input_image", "image_url": "data:image/png;base64,AA=="} + data: Final = { + "model": "judge", + "input": [ + {"role": "user", "content": [{"type": "input_text", "text": "ticket one"}, image]}, + {"role": "user", "content": "ticket two"}, + ], + "questions": [ + {"type": "predicate", "name": "urgent", "instructions": "Is it urgent?"}, + {"type": "box", "name": "logo", "instructions": "ticket one"}, + ], + } + + result: Final = await DecisionsHandler().process_input_messages(data, guardrail) + + assert guardrail.seen == [("request", ("ticket one", "ticket two", "Is it urgent?"))] + assert {key: value for key, value in result.items() if key != "metadata"} == { + "model": "judge", + "input": [ + {"role": "user", "content": [{"type": "input_text", "text": "ticket one [GUARDRAILED]"}, image]}, + {"role": "user", "content": "ticket two [GUARDRAILED]"}, + ], + "questions": [ + {"type": "predicate", "name": "urgent", "instructions": "Is it urgent? [GUARDRAILED]"}, + {"type": "box", "name": "logo", "instructions": "ticket one [GUARDRAILED]"}, + ], + } + assert data["input"][0]["content"][0]["text"] == "ticket one" + + +@pytest.mark.asyncio +async def test_systemone_state_leaves_and_question_instructions_are_guardrailed() -> None: + guardrail: Final = _TaggingGuardrail() + criteria_question: Final = {"type": "noul", "criteria": {"yes": "refund"}} + data: Final = { + "model": "judge", + "state": {"subject": "refund", "lines": ["order 7", "", 3, True], "meta": {"id": None}}, + "questions": { + "refund": {"type": "noul", "instructions": {"ask": "Does the customer want a refund?"}}, + "routed": criteria_question, + }, + } + + result: Final = await SystemOneHandler().process_input_messages(data, guardrail) + + assert guardrail.seen == [("request", ("refund", "order 7", "Does the customer want a refund?"))] + assert {key: value for key, value in result.items() if key != "metadata"} == { + "model": "judge", + "state": { + "subject": "refund [GUARDRAILED]", + "lines": ["order 7 [GUARDRAILED]", "", 3, True], + "meta": {"id": None}, + }, + "questions": { + "refund": {"type": "noul", "instructions": {"ask": "Does the customer want a refund? [GUARDRAILED]"}}, + "routed": criteria_question, + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("handler", "data"), + ( + ( + SystemOneHandler, + {"model": "judge", "state": {"count": 3}, "questions": {"q": {"type": "noul", "criteria": {}}}}, + ), + ( + DecisionsHandler, + { + "model": "judge", + "input": [{"role": "user", "content": ""}], + "questions": [{"type": "predicate", "instructions": ""}], + }, + ), + ), +) +async def test_request_without_text_is_returned_unscanned( + handler: type[DecisionsHandler] | type[SystemOneHandler], data: dict[str, object] +) -> None: + guardrail: Final = _TaggingGuardrail() + + result: Final = await handler().process_input_messages(data, guardrail) + + assert result is data + assert guardrail.seen == [] + + +@pytest.mark.asyncio +async def test_guardrail_that_returns_no_texts_leaves_the_request_text_as_sent() -> None: + data: Final = { + "model": "judge", + "input": "ticket", + "questions": [{"type": "predicate", "instructions": "Is it urgent?"}], + } + + result: Final = await DecisionsHandler().process_input_messages(data, _ReturnsNoTexts(guardrail_name="noop")) + + assert result == data + + +@pytest.mark.asyncio +@pytest.mark.parametrize("handler", (DecisionsHandler, SystemOneHandler)) +async def test_decision_answers_are_returned_as_the_provider_sent_them( + handler: type[DecisionsHandler] | type[SystemOneHandler], +) -> None: + guardrail: Final = _TaggingGuardrail() + response: Final = OpenAIDecisionResponse.model_validate( + { + "model": "judge", + "answers": [{"type": "refusal", "name": "q"}], + "usage": {"input_tokens": 1, "output_tokens": 0, "total_tokens": 1}, + } + ) + + result: Final = await handler().process_output_response(response, guardrail) + + assert result is response + assert guardrail.seen == [] diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index c42878bd4f7..8fbe997b3e7 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -1137,8 +1137,10 @@ class _PolicyOptOutGuardrail(CustomGuardrail): def __init__(self): super().__init__(guardrail_name="policy-opt-out", default_on=True) self.native_pre_call_ran = False + self.scanned = [] async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + self.scanned.extend(inputs.get("texts", [])) return inputs async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): @@ -1164,6 +1166,24 @@ async def test_pipeline_step_keeps_native_hook_when_opted_out(monkeypatch): assert "guardrail_to_apply" not in data +@pytest.mark.asyncio +async def test_pipeline_step_scans_extraction_state_through_apply_guardrail_when_opted_out(monkeypatch): + guardrail = _PolicyOptOutGuardrail() + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + outcome, _, _, _ = await PipelineExecutor._run_step( + step=PipelineStep(guardrail="policy-opt-out", on_fail="block", on_pass="allow"), + mode="pre_call", + data={"model": "m", "state": "invoice from vendor", "fields": {"total": {"type": "number"}}}, + user_api_key_dict=UserAPIKeyAuth(), + call_type="aextraction", + ) + + assert outcome == "pass" + assert guardrail.native_pre_call_ran is False + assert guardrail.scanned == ["invoice from vendor"] + + class _TextReturningGuardrail(CustomGuardrail): def __init__(self, returned_texts): super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True) diff --git a/tests/unit/proxy/test_proxy_logging_hook_detection.py b/tests/unit/proxy/test_proxy_logging_hook_detection.py index 182ef833069..a6671aba7a6 100644 --- a/tests/unit/proxy/test_proxy_logging_hook_detection.py +++ b/tests/unit/proxy/test_proxy_logging_hook_detection.py @@ -3,6 +3,7 @@ from fastapi import HTTPException import litellm from litellm.caching import DualCache +from litellm.exceptions import SensitiveDataRouteException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth @@ -1060,3 +1061,250 @@ async def test_parallel_post_call_guardrails_keep_native_hook_when_opted_out(mon assert opted_out.native_hooks_ran == ["post_call"] assert routed.native_hooks_ran == [] + + +_SECRET = "sk-live-1234" + + +def _applied_guardrails(data): + buckets = (data.get("metadata") or {}, data.get("litellm_metadata") or {}) + return [*buckets[0].get("applied_guardrails", []), *buckets[1].get("applied_guardrails", [])] + + +class _MasksSecretWithNativeHooks(_KeepsNativeHooks): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.scanned = [] + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + texts = inputs.get("texts", []) + self.scanned.append((input_type, tuple(texts))) + return {**inputs, "texts": [text.replace(_SECRET, "[MASKED]") for text in texts]} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "data", "expected"), + ( + ( + "aextraction", + {"model": "m", "state": f"paid with {_SECRET}", "fields": {"total": {"type": "number"}}}, + {"state": "paid with [MASKED]"}, + ), + ( + "asystemone", + { + "model": "m", + "state": {"ticket": f"my key is {_SECRET}"}, + "questions": {"leak": {"type": "noul", "instructions": f"Is {_SECRET} a key?"}}, + }, + { + "state": {"ticket": "my key is [MASKED]"}, + "questions": {"leak": {"type": "noul", "instructions": "Is [MASKED] a key?"}}, + }, + ), + ( + "adecisions", + { + "model": "m", + "input": f"my key is {_SECRET}", + "questions": [{"type": "predicate", "instructions": f"Is {_SECRET} a key?"}], + }, + { + "input": "my key is [MASKED]", + "questions": [{"type": "predicate", "instructions": "Is [MASKED] a key?"}], + }, + ), + ), +) +async def test_pre_call_hook_scans_decision_and_extraction_text_with_a_native_hook_guardrail( + monkeypatch, call_type, data, expected +): + guardrail = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + result = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data=data, + call_type=call_type, + ) + + assert {key: result[key] for key in expected} == expected + assert guardrail.native_hooks_ran == [] + assert _applied_guardrails(result) == ["keeps_native"] + + +@pytest.mark.asyncio +async def test_during_call_hook_scans_extraction_state_with_a_native_hook_guardrail(monkeypatch): + guardrail = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.during_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + await ProxyLogging(user_api_key_cache=DualCache()).during_call_hook( + data={"model": "m", "state": f"paid with {_SECRET}", "fields": {"total": {"type": "number"}}}, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + call_type="aextraction", + ) + + assert guardrail.native_hooks_ran == [] + assert guardrail.scanned == [("request", (f"paid with {_SECRET}",))] + + +@pytest.mark.asyncio +async def test_pre_call_hook_keeps_native_hook_guardrail_on_its_own_hook_for_chat(monkeypatch): + guardrail = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data={"model": "m", "messages": [{"role": "user", "content": f"my key is {_SECRET}"}]}, + call_type="acompletion", + ) + + assert guardrail.native_hooks_ran == ["pre_call"] + assert guardrail.scanned == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("run_in_parallel", (False, True)) +async def test_post_call_success_hook_scans_extraction_answers_with_a_native_hook_guardrail( + monkeypatch, run_in_parallel +): + from litellm.types.extraction import ExtractionResponse + + guardrail = _MasksSecretWithNativeHooks( + event_hook=GuardrailEventHooks.post_call, default_on=True, run_in_parallel=run_in_parallel + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + response = ExtractionResponse.model_validate( + { + "model": "m", + "answers": {"key": {"type": "text", "value": _SECRET}}, + "truncated": False, + "usage": {"input_tokens": 5, "output_tokens": 0}, + } + ) + + await ProxyLogging(user_api_key_cache=DualCache()).post_call_success_hook( + data={"model": "m", "state": "s", "fields": {"key": {"type": "text"}}}, + response=response, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/v1/extraction"), + ) + + assert guardrail.native_hooks_ran == [] + assert guardrail.scanned == [("response", (_SECRET,))] + + +class _ReadsNoDecisionText(CustomGuardrail): + def __init__(self, guardrail_name, **kwargs): + super().__init__(guardrail_name=guardrail_name, event_hook=GuardrailEventHooks.pre_call, **kwargs) + self.native_hooks_ran = [] + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + self.native_hooks_ran.append(call_type) + return data + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "listed"), (("aextraction", False), ("acompletion", True))) +async def test_requested_guardrail_without_apply_guardrail_is_listed_only_where_it_reads_the_text( + monkeypatch, caplog, call_type, listed +): + requested = _ReadsNoDecisionText("requested-native") + default_on = _ReadsNoDecisionText("default-native", default_on=True) + scanner = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [requested, default_on, scanner]) + + with caplog.at_level("WARNING"): + result = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data={ + "model": "m", + "state": "s", + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"guardrails": ["requested-native"]}, + }, + call_type=call_type, + ) + + assert requested.native_hooks_ran == [call_type] + assert ("requested-native" in _applied_guardrails(result)) is listed + warnings = [record.getMessage() for record in caplog.records if "cannot scan" in record.getMessage()] + assert warnings == ( + [] + if listed + else [ + "Guardrails requested-native, default-native do not implement apply_guardrail, so they cannot scan the text of aextraction requests" + ] + ) + + +class _ReroutesToPrivateModel(CustomGuardrail): + def __init__(self): + super().__init__( + guardrail_name="reroute", event_hook=GuardrailEventHooks.pre_call, default_on=True, run_in_parallel=True + ) + + async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): + raise SensitiveDataRouteException( + route_to_model="private-model", session_id="s1", guardrail_name="reroute", sticky_session_routing=False + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "listed"), (("aextraction", False), ("acompletion", True))) +async def test_rerouted_request_lists_a_requested_guardrail_only_where_it_reads_the_text( + monkeypatch, call_type, listed +): + requested = _ReadsNoDecisionText("requested-native") + monkeypatch.setattr(litellm, "callbacks", [requested, _ReroutesToPrivateModel()]) + + result = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data={ + "model": "m", + "state": "s", + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"guardrails": ["requested-native"]}, + }, + call_type=call_type, + ) + + assert result["model"] == "private-model" + assert ("requested-native" in _applied_guardrails(result)) is listed + + +class _RecordsScans(_AppliesGuardrail): + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.scanned = [] + + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + self.scanned.append(input_type) + return inputs + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ("/v1/decisions", "/v1/systemone")) +async def test_post_call_guardrail_leaves_decision_answers_unscanned_and_unlisted(monkeypatch, route): + from litellm.types.decisions import OpenAIDecisionResponse + + guardrail = _RecordsScans(event_hook=GuardrailEventHooks.post_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + response = OpenAIDecisionResponse.model_validate( + { + "model": "m", + "answers": [{"type": "refusal", "name": "q"}], + "usage": {"input_tokens": 1, "output_tokens": 0, "total_tokens": 1}, + } + ) + data = {"model": "m", "metadata": {}} + + result = await ProxyLogging(user_api_key_cache=DualCache()).post_call_success_hook( + data=data, + response=response, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route=route), + ) + + assert result is response + assert guardrail.scanned == [] + assert _applied_guardrails(data) == [] diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index acc1463b072..a551c42dd3e 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -275,7 +275,7 @@ def test_process_guardrail_metadata_calls_header_helper(proxy_logging, monkeypat monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) data = {"metadata": {"guardrails": ["g1", "g2"]}} - proxy_logging._process_guardrail_metadata(data) + proxy_logging._process_guardrail_metadata(data, "completion") snapshot = { "call_count": len(calls), "first_name": calls[0]["name"], @@ -300,7 +300,7 @@ def test_process_guardrail_metadata_skips_already_applied(proxy_logging, monkeyp monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) data = {"metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g1"]}} - proxy_logging._process_guardrail_metadata(data) + proxy_logging._process_guardrail_metadata(data, "completion") assert calls == ["g2"] @@ -312,12 +312,12 @@ def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging, monkeypat "add_guardrail_to_applied_guardrails_header", MagicMock(side_effect=AssertionError("should not be called")), ) - proxy_logging._process_guardrail_metadata({}) + proxy_logging._process_guardrail_metadata({}, "completion") def test_process_guardrail_metadata_invalid_data_raises(proxy_logging): with pytest.raises(AttributeError): - proxy_logging._process_guardrail_metadata(None) # type: ignore[arg-type] + proxy_logging._process_guardrail_metadata(None, "completion") # type: ignore[arg-type] # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py index b7fc8c6dbeb..01a962f5c37 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -170,7 +170,7 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox proxy_logging.slack_alerting_instance = MagicMock(alerting=None) invoked = {} - def fake_process(d): + def fake_process(d, call_type): invoked["data"] = d proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment]