diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py b/enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py index c190e680347..5fc251fd330 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 @@ -601,6 +601,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)) @@ -609,7 +627,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): @@ -694,12 +712,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( @@ -710,8 +728,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..7fa2721e5b3 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -14,6 +14,7 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_logger from litellm.caching import DualCache from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, @@ -72,6 +73,30 @@ _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 is_apply_guardrail_only_route(route: str | None) -> bool: + call_types: Final = get_call_types_for_route(route) if route else None + return bool(call_types) and call_types[0] in APPLY_GUARDRAIL_ONLY_CALL_TYPES + + +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 " @@ -226,6 +251,7 @@ class CustomGuardrail(CustomLogger): use_native_lifecycle_hooks: ClassVar[bool] = False records_own_guardrail_information: ClassVar[bool] = False + adds_itself_to_applied_guardrails: ClassVar[bool] = False logging_only_scope: LoggingOnlyScope | None logging_only_continue_on_input_failure: bool diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 6399ded9555..6b044a87521 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -72,6 +72,10 @@ 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 + + lists_guardrail_when_called: ClassVar[bool] = False + 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..5b02a037041 --- /dev/null +++ b/litellm/llms/openai/decisions/guardrail_translation/handler.py @@ -0,0 +1,412 @@ +import json +import re +from collections.abc import Iterator, Mapping +from difflib import SequenceMatcher +from itertools import chain, groupby +from types import MappingProxyType +from typing import TYPE_CHECKING, ClassVar, Final + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.llms.base_llm.guardrail_translation.utils import unappliable_request_rewrite +from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler +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.llms.openai import AllMessageValues + from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth + +_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_KEY_ADAPTER: Final[TypeAdapter[str]] = TypeAdapter(str) +_JSON_KEY_PREFIX: Final = re.compile(r'("(?:[^"\\\x00-\x1f]|\\(?:["\\/bfnrt]|u[0-9a-fA-F]{4}))*"): ') +_JSON_ESCAPE: Final = re.compile(r"(\\(?:[bfnrt]|u[0-9a-fA-F]{4}))|\\.") +_MARKED_JSON_ESCAPE: Final = re.compile(r"(\\(?:[bfnrt]|u[0-9a-fA-F]{4}))~|\\.") +_TOKEN: Final = re.compile(r"\s+|\S+") +_DESCRIBED_ITEM_KEYS: Final = ("choices", "levels") +_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 _marked_escape(escape: re.Match[str]) -> str: + return f"{escape.group()}~" if escape.group(1) else escape.group() + + +def _unmarked_escape(escape: re.Match[str]) -> str: + return escape.group(1) or escape.group() + + +def _entry_text(key: str, value: str | float | None) -> str: + return _JSON_ESCAPE.sub(_marked_escape, json.dumps({key: value}, ensure_ascii=False)[1:-1]) + + +def _escaped_texts(text: str, key: str, item: JsonValue) -> tuple[str, ...]: + return tuple(form for form in (key, item) if isinstance(form, str) and form not in text) + + +def _entry_text_forms(key: str, item: JsonValue) -> Iterator[str]: + if isinstance(item, (list, dict)): + yield key + yield from _json_text_forms(item) + else: + text: Final = _entry_text(key, item) + yield text + yield from _escaped_texts(text, key, item) + + +def _json_text_forms(value: JsonValue) -> Iterator[str]: + match value: + case bool() | None: + return + case str(): + yield value + case int() | float(): + yield json.dumps(value) + case list(): + yield from chain.from_iterable(map(_json_text_forms, value)) + case dict(): + for key, item in value.items(): + yield from _entry_text_forms(key, item) + + +def _json_texts(value: JsonValue) -> tuple[str, ...]: + return tuple(text for text in _json_text_forms(value) if text) + + +def _guardrailed_number(value: float, guardrailed: Mapping[str, str]) -> JsonValue: + text: Final = json.dumps(value) + rewritten: Final = guardrailed.get(text, text) + return value if rewritten == text else rewritten + + +def _occurrence_numbers(keys: tuple[str, ...]) -> Iterator[tuple[int, int]]: + for _, positions in groupby(sorted(range(len(keys)), key=keys.__getitem__), key=keys.__getitem__): + yield from ((position, number) for number, position in enumerate(positions, start=1)) + + +def _distinct_keys(keys: tuple[str, ...]) -> tuple[str, ...]: + if len(frozenset(keys)) == len(keys): + return keys + numbers: Final = MappingProxyType(dict(_occurrence_numbers(keys))) + return tuple(key if numbers[position] == 1 else f"{key}_{numbers[position]}" for position, key in enumerate(keys)) + + +def _guardrailed_json(value: JsonValue, guardrailed: Mapping[str, str], guardrail_name: str | None) -> JsonValue: + match value: + case bool() | None: + return value + case str(): + return guardrailed.get(value, value) + case int() | float(): + return _guardrailed_number(value, guardrailed) + case list(): + return [_guardrailed_json(item, guardrailed, guardrail_name) for item in value] + case dict(): + entries: Final = tuple( + _guardrailed_entry(key, item, guardrailed, guardrail_name) for key, item in value.items() + ) + keys: Final = _distinct_keys(tuple(key for key, _ in entries)) + if len(frozenset(keys)) != len(keys): + raise unappliable_request_rewrite(guardrail_name) + return dict(zip(keys, (item for _, item in entries), strict=True)) + + +def _rewritten_value(text: str) -> JsonValue: + try: + return _JSON_ADAPTER.validate_json(text) + except ValidationError: + return text + + +def _rewritten_entry(rewritten: str, guardrail_name: str | None) -> tuple[str, JsonValue]: + unmarked: Final = _MARKED_JSON_ESCAPE.sub(_unmarked_escape, rewritten) + key: Final = _JSON_KEY_PREFIX.match(unmarked) + if key is None: + raise unappliable_request_rewrite(guardrail_name) + return _KEY_ADAPTER.validate_json(key.group(1)), _rewritten_value(unmarked[key.end() :]) + + +def _decoded_text(text: str, form: str, guardrailed: Mapping[str, str]) -> str: + return form if form in text else guardrailed.get(form, form) + + +def _tokens(text: str) -> tuple[str, ...]: + return tuple(_TOKEN.findall(text)) + + +def _edits(original: tuple[str, ...], rewritten: tuple[str, ...]) -> frozenset[tuple[int, int, str]]: + opcodes: Final = SequenceMatcher(None, original, rewritten, autojunk=False).get_opcodes() + return frozenset( + (start, end, "".join(rewritten[low:high])) for tag, start, end, low, high in opcodes if tag != "equal" + ) + + +def _covered(edit: tuple[int, int, str], edits: frozenset[tuple[int, int, str]]) -> bool: + start, end, _ = edit + return any(other[0] <= start and end <= other[1] and other[1] - other[0] > end - start for other in edits) + + +def _merged_text(original: str, rewritten: str, decoded: str, guardrail_name: str | None) -> str: + if decoded in (original, rewritten): + return rewritten + if rewritten == original: + return decoded + tokens: Final = _tokens(original) + edits: Final = _edits(tokens, _tokens(rewritten)) | _edits(tokens, _tokens(decoded)) + kept: Final = tuple(sorted(edit for edit in edits if not _covered(edit, edits))) + if any(later[0] < earlier[1] for earlier, later in zip(kept, kept[1:])): + raise unappliable_request_rewrite(guardrail_name) + kept_starts: Final = (0, *(end for _, end, _ in kept)) + kept_ends: Final = (*(start for start, _, _ in kept), len(tokens)) + replacements: Final = (*(text for _, _, text in kept), "") + return "".join( + "".join(tokens[start:end]) + text for start, end, text in zip(kept_starts, kept_ends, replacements, strict=True) + ) + + +def _merged_value(original: str, rewritten: JsonValue, decoded: str, guardrail_name: str | None) -> JsonValue: + if decoded == original: + return rewritten + if isinstance(rewritten, str): + return _merged_text(original, rewritten, decoded, guardrail_name) + raise unappliable_request_rewrite(guardrail_name) + + +def _guardrailed_entry( + key: str, item: JsonValue, guardrailed: Mapping[str, str], guardrail_name: str | None +) -> tuple[str, JsonValue]: + if isinstance(item, (list, dict)): + return guardrailed.get(key, key), _guardrailed_json(item, guardrailed, guardrail_name) + text: Final = _entry_text(key, item) + rewritten: Final = guardrailed.get(text, text) + entry: Final = (key, item) if rewritten == text else _rewritten_entry(rewritten, guardrail_name) + merged_key: Final = _merged_text(key, entry[0], _decoded_text(text, key, guardrailed), guardrail_name) + if not isinstance(item, str): + return merged_key, entry[1] + return merged_key, _merged_value(item, entry[1], _decoded_text(text, item, guardrailed), guardrail_name) + + +def _instruction_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]: + return _json_texts(question.get("instructions")) + + +def _described_items(question: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]: + return tuple(chain.from_iterable(_list_items(question.get(key)) for key in _DESCRIBED_ITEM_KEYS)) + + +def _list_items(value: JsonValue) -> tuple[JsonValue, ...]: + return tuple(value) if isinstance(value, list) else () + + +def _item_description(item: JsonValue) -> tuple[str, ...]: + description: Final = item.get("description") if isinstance(item, dict) else None + return (description,) if isinstance(description, str) and description else () + + +def _decision_question_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]: + descriptions: Final = chain.from_iterable(map(_item_description, _described_items(question))) + return (*_instruction_texts(question), *descriptions) + + +def _guardrailed_item(item: JsonValue, guardrailed: Mapping[str, str]) -> JsonValue: + if not isinstance(item, dict) or not isinstance(description := item.get("description"), str): + return item + return {**item, "description": guardrailed.get(description, description)} + + +def _with_guardrailed_decision_question( + question: Mapping[str, JsonValue], guardrailed: Mapping[str, str], guardrail_name: str | None +) -> Mapping[str, JsonValue]: + described: Final = { + key: [_guardrailed_item(item, guardrailed) for item in _list_items(question[key])] + for key in _DESCRIBED_ITEM_KEYS + if key in question + } + return {**_with_guardrailed_instructions(question, guardrailed, guardrail_name), **described} + + +def _criteria_texts(criteria: JsonValue) -> tuple[str, ...]: + descriptions: Final = tuple(criteria.values()) if isinstance(criteria, dict) else (criteria,) + return tuple(chain.from_iterable(map(_json_texts, descriptions))) + + +def _systemone_question_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]: + return (*_instruction_texts(question), *_criteria_texts(question.get("criteria"))) + + +def _guardrailed_criteria(criteria: JsonValue, guardrailed: Mapping[str, str], guardrail_name: str | None) -> JsonValue: + if not isinstance(criteria, dict): + return _guardrailed_json(criteria, guardrailed, guardrail_name) + return { + label: _guardrailed_json(description, guardrailed, guardrail_name) for label, description in criteria.items() + } + + +def _with_guardrailed_systemone_question( + question: Mapping[str, JsonValue], guardrailed: Mapping[str, str], guardrail_name: str | None +) -> Mapping[str, JsonValue]: + with_instructions: Final = _with_guardrailed_instructions(question, guardrailed, guardrail_name) + if "criteria" not in question: + return with_instructions + return {**with_instructions, "criteria": _guardrailed_criteria(question["criteria"], guardrailed, guardrail_name)} + + +def _with_guardrailed_instructions( + question: Mapping[str, JsonValue], guardrailed: Mapping[str, str], guardrail_name: str | None +) -> Mapping[str, JsonValue]: + if "instructions" not in question: + return question + return {**question, "instructions": _guardrailed_json(question["instructions"], guardrailed, guardrail_name)} + + +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, + ) + if not ( + guardrail_to_apply.records_own_guardrail_information or guardrail_to_apply.adds_itself_to_applied_guardrails + ): + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict + ) + + add_guardrail_to_applied_guardrails_header( + request_data=request_data, guardrail_name=guardrail_to_apply.guardrail_name + ) + 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 + lists_guardrail_when_called: ClassVar[bool] = True + + 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): + def get_structured_messages( + self, + data: dict[str, object], # mutable-ok: BaseTranslation hands over the request dict + ) -> "list[AllMessageValues] | None": # mutable-ok: overrides BaseTranslation.get_structured_messages + return OpenAIResponsesHandler().get_structured_messages({"input": data.get("input")}) + + 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(_decision_question_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_decision_question(question, guardrailed, guardrail_to_apply.guardrail_name) + 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 = (*_json_texts(state), *chain.from_iterable(map(_systemone_question_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, guardrail_to_apply.guardrail_name), + "questions": { + name: _with_guardrailed_systemone_question(question, guardrailed, guardrail_to_apply.guardrail_name) + for name, question in questions.items() + }, + } diff --git a/litellm/proxy/common_utils/callback_utils.py b/litellm/proxy/common_utils/callback_utils.py index d1787f3ff8b..ad73a99f967 100644 --- a/litellm/proxy/common_utils/callback_utils.py +++ b/litellm/proxy/common_utils/callback_utils.py @@ -1,9 +1,9 @@ import copy import json import os -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass -from itertools import accumulate +from itertools import accumulate, chain from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias from pydantic import TypeAdapter @@ -23,6 +23,7 @@ from litellm.constants import ( ) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, get_metadata_variable_name_from_kwargs, get_or_create_metadata_bucket, ) @@ -478,16 +479,33 @@ def _serialize_scan_metadata_header(entries: Iterable[object], *, max_length: in return f"[{','.join(encoded[:kept])}]" +_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object]) +_LIST_ENTRIES: Final = TypeAdapter(tuple[object, ...]) + + +def _list_entries(value: object) -> tuple[object, ...]: + return _LIST_ENTRIES.validate_python(value) if isinstance(value, (list, tuple)) else () + + +def _bucket_list_entries(buckets: Sequence[Mapping[str, object]], key: str) -> tuple[object, ...]: + values: Final = tuple(bucket.get(key) for bucket in buckets) + distinct_values: Final = {id(value): value for value in values} + entries: Final = tuple(chain.from_iterable(_list_entries(value) for value in distinct_values.values())) + return entries if key == GUARDRAIL_SCAN_METADATA_METADATA_KEY else tuple(dict.fromkeys(entries)) + + def get_logging_caching_headers(request_data: dict) -> dict | None: _metadata: Final[dict] = {} - metadata_bucket: Final = request_data.get("metadata") - litellm_metadata_bucket: Final = request_data.get("litellm_metadata") - if isinstance(metadata_bucket, dict): - _metadata.update(metadata_bucket) - if isinstance(litellm_metadata_bucket, dict): + request: Final = _REQUEST_BODY.validate_python(request_data) + found_buckets: Final = (as_str_mapping(request.get(name)) for name in ("metadata", "litellm_metadata")) + buckets: Final = tuple(bucket for bucket in found_buckets if bucket is not None) + for bucket in buckets: # Batch/file routes store proxy tracking in litellm_metadata while # user-facing metadata stays in metadata; merge both for headers. - _metadata.update(litellm_metadata_bucket) + _metadata.update(bucket) + for list_key in ("applied_guardrails", GUARDRAIL_SCAN_IDS_METADATA_KEY, GUARDRAIL_SCAN_METADATA_METADATA_KEY): + if any(list_key in bucket for bucket in buckets): + _metadata[list_key] = _bucket_list_entries(buckets, list_key) headers: Final = {} if "applied_guardrails" in _metadata: headers["x-litellm-applied-guardrails"] = ",".join(_metadata["applied_guardrails"]) diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py index 61117fbc55e..9629a1e5e93 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py @@ -17,12 +17,15 @@ from websockets.asyncio.client import ClientConnection, connect from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) from litellm.proxy._types import ProxyException, UserAPIKeyAuth +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict +) from litellm.proxy.guardrails._content_utils import ( apply_redacted_messages_back, build_inspection_messages, @@ -152,7 +155,9 @@ class AimGuardrail(CustomGuardrail): if is_non_conversational_call_type(call_type) and not self.inspect_embeddings: verbose_proxy_logger.debug("Aim: skipping non-conversational call type %s", call_type) return data - return await self.call_aim_guardrail(data, hook="pre_call", key_alias=user_api_key_dict.key_alias) + return await self.call_aim_guardrail( + data, hook="pre_call", key_alias=user_api_key_dict.key_alias, call_type=call_type + ) async def async_moderation_hook( self, @@ -165,10 +170,19 @@ class AimGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Aim: skipping non-conversational call type %s", call_type) return data - await self.call_aim_guardrail(data, hook="moderation", key_alias=user_api_key_dict.key_alias) + await self.call_aim_guardrail( + data, hook="moderation", key_alias=user_api_key_dict.key_alias, call_type=call_type + ) return data - async def call_aim_guardrail(self, data: dict, hook: str, key_alias: str | None) -> dict: + async def call_aim_guardrail( + self, data: dict, hook: str, key_alias: str | None, call_type: CallTypesLiteral | None = None + ) -> dict: + messages: Final = self._build_aim_inspection_messages(data) # pyright: ignore[reportUnknownMemberType] # hook request data is an untyped dict + if not messages: + return data # pyright: ignore[reportUnknownVariableType] # hook request data is an untyped dict + if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) user_email: Final = data.get("metadata", {}).get("headers", {}).get("x-aim-user-email") call_id: Final = data.get("litellm_call_id") headers: Final = self._build_aim_headers( @@ -180,7 +194,7 @@ class AimGuardrail(CustomGuardrail): response: Final = await self.async_handler.post( f"{self.api_base}/fw/v1/analyze", headers=headers, - json={"messages": self._build_aim_inspection_messages(data)}, + json={"messages": messages}, timeout=self.timeout, ) response.raise_for_status() diff --git a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py index 19b9bccd8aa..a80575c3cc1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -22,7 +22,7 @@ from websockets.exceptions import ConnectionClosed from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail from litellm.litellm_core_utils.core_helpers import iter_items from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -30,6 +30,9 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict +) from litellm.proxy.guardrails._content_utils import ( apply_redacted_messages_back, build_inspection_messages, @@ -173,6 +176,7 @@ class CatoNetworksGuardrail(CustomGuardrail): hook="pre_call", key_alias=user_api_key_dict.key_alias, user_email=self._resolve_cato_user_email(user_api_key_dict), + call_type=call_type, ) async def async_moderation_hook( @@ -190,6 +194,7 @@ class CatoNetworksGuardrail(CustomGuardrail): hook="moderation", key_alias=user_api_key_dict.key_alias, user_email=self._resolve_cato_user_email(user_api_key_dict), + call_type=call_type, ) @classmethod @@ -294,7 +299,13 @@ class CatoNetworksGuardrail(CustomGuardrail): hook: str, key_alias: str | None, user_email: str | None = None, + call_type: CallTypesLiteral | None = None, ) -> dict: + messages: Final = self._inspection_messages(data) # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType] # hook request data is an untyped dict + if not messages: + return data # pyright: ignore[reportUnknownVariableType] # hook request data is an untyped dict + if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) call_id: Final = data.get("litellm_call_id") headers: Final = self._build_cato_headers( hook=hook, @@ -305,7 +316,7 @@ class CatoNetworksGuardrail(CustomGuardrail): response: Final = await self.async_handler.post( f"{self.api_base}/fw/v1/analyze", headers=headers, - json={"messages": self._inspection_messages(data)}, + json={"messages": messages}, timeout=self.timeout, ) response.raise_for_status() diff --git a/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py b/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py index f3bd2b3f889..11db4690cde 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py +++ b/litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py @@ -20,6 +20,9 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, # pyright: ignore[reportUnknownVariableType] # decorator is untyped upstream ) from litellm.litellm_core_utils.llm_judge import attached_auth, caller_team_for, default_router_provider, judge_target +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict +) from litellm.types.decisions import MAX_DECISION_QUESTIONS, OpenAIDecisionResponse, OpenAIPredicateAnswer from litellm.types.guardrails import GuardrailEventHooks, Mode from litellm.types.llms.openai import ChatCompletionAssistantToolCall, ChatCompletionToolCallChunk @@ -347,6 +350,7 @@ class DecisionModelGuardrail(CustomGuardrail): status: Final = "guardrail_intervened" if intervened else ("guardrail_flagged" if flagged else "success") self._log(request_data, verdicts, status, start_time) + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) if intervened: raise HTTPException( status_code=400, diff --git a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py index d507df24faa..0bb6454f35f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py @@ -15,7 +15,11 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + APPLY_GUARDRAIL_ONLY_CALL_TYPES, + CustomGuardrail, + is_apply_guardrail_only_route, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -328,8 +332,8 @@ class DynamoAIGuardrails(CustomGuardrail): error_message: Final = self._create_error_message(processed_result) raise ValueError(error_message) - # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + if _messages or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -373,8 +377,8 @@ class DynamoAIGuardrails(CustomGuardrail): error_message: Final = self._create_error_message(processed_result) raise ValueError(error_message) - # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + if _messages or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -441,8 +445,10 @@ class DynamoAIGuardrails(CustomGuardrail): error_message: Final = self._create_error_message(processed_result) raise ValueError(error_message) - # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + if isinstance(response, litellm.ModelResponse) or not is_apply_guardrail_only_route( + user_api_key_dict.request_route + ): + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) async def async_post_call_streaming_iterator_hook( self, diff --git a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py index 83dd90dafb4..e7ff6ec1039 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py @@ -15,7 +15,11 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import ( + APPLY_GUARDRAIL_ONLY_CALL_TYPES, + CustomGuardrail, + is_apply_guardrail_only_route, +) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -476,8 +480,8 @@ class IBMGuardrailDetector(CustomGuardrail): error_message = self._create_error_message_orchestrator(orchestrator_result) raise ValueError(error_message) - # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + if contents_to_check or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -545,8 +549,8 @@ class IBMGuardrailDetector(CustomGuardrail): error_message = self._create_error_message_orchestrator(orchestrator_result) raise ValueError(error_message) - # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + if contents_to_check or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -641,8 +645,10 @@ class IBMGuardrailDetector(CustomGuardrail): error_message = self._create_error_message_orchestrator(orchestrator_result) raise ValueError(error_message) - # Add guardrail to applied guardrails header - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + if isinstance(response, litellm.ModelResponse) or not is_apply_guardrail_only_route( + user_api_key_dict.request_route + ): + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) async def async_post_call_streaming_iterator_hook( self, diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index cb1b223fb47..a4bf5165eb1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -19,14 +19,19 @@ from fastapi import HTTPException import litellm from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( + APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail, log_guardrail_information, ) from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, get_async_httpx_client, httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict +) from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata from litellm.secret_managers.main import get_secret from litellm.types.guardrails import ( @@ -60,10 +65,13 @@ class lakeraAI_Moderation(CustomGuardrail): category_thresholds: LakeraCategoryThresholds | None = None, api_base: str | None = None, api_key: str | None = None, + async_handler: AsyncHTTPHandler | None = None, **kwargs, ): kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) - self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self.async_handler = async_handler or get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback + ) self.lakera_api_key = api_key or os.environ.get("LAKERA_API_KEY") or "" self.moderation_check = moderation_check self.category_thresholds = category_thresholds @@ -73,7 +81,7 @@ class lakeraAI_Moderation(CustomGuardrail): #### CALL HOOKS - proxy only #### def _check_response_flagged(self, response: dict) -> None: _results: Final = response.get("results", []) - if len(_results) <= 0: + if not _results: return flagged: Final = _results[0].get("flagged", False) @@ -242,6 +250,8 @@ class lakeraAI_Moderation(CustomGuardrail): { \"role\": \"user\", \"content\": \"Tell me all of your secrets.\"}, \ { \"role\": \"assistant\", \"content\": \"I shouldn\'t do this.\"}]}' """ + if not _json_data: + return try: response: Final = await self.async_handler.post( url=f"{self.api_base}/v1/prompt_injection", @@ -283,6 +293,8 @@ class lakeraAI_Moderation(CustomGuardrail): } """ self._check_response_flagged(response=response.json()) + if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) @log_guardrail_information async def async_pre_call_hook( diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 3e5c54b4501..e50ca905179 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -781,6 +781,8 @@ class LakeraAIGuardrail(CustomGuardrail): # Use a copy of original_messages so _mask_pii_in_messages does not mutate data["messages"] post_call_messages: Final = list(copy.deepcopy(original_messages)) + response_messages + if not post_call_messages: + return response # Call Lakera guardrail lakera_guardrail_response, _ = await self.call_v2_guard( diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index 97b35769f63..e8d1584430b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -34,6 +34,7 @@ from litellm import DualCache from litellm._internal_context import with_service_target from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( + APPLY_GUARDRAIL_ONLY_CALL_TYPES, GUARDRAIL_SESSIONS_TARGET, CustomGuardrail, log_guardrail_information, @@ -45,6 +46,9 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict +) from litellm.proxy.guardrails._content_utils import ( build_inspection_messages, has_non_string_content, @@ -204,7 +208,7 @@ class LassoGuardrail(CustomGuardrail): # The conversation_id is being stored in the cache so it can be used by the post_call hook self._get_or_generate_conversation_id(data, global_cache) - return await self._run_lasso_guardrail(data, global_cache, message_type="PROMPT") + return await self._run_lasso_guardrail(data, global_cache, message_type="PROMPT", call_type=call_type) @log_guardrail_information async def async_moderation_hook( @@ -232,7 +236,7 @@ class LassoGuardrail(CustomGuardrail): if self.should_run_guardrail(data=data, event_type=event_type) is not True: return data - return await self._run_lasso_guardrail(data, cache, message_type="PROMPT") + return await self._run_lasso_guardrail(data, cache, message_type="PROMPT", call_type=call_type) @log_guardrail_information async def async_post_call_success_hook( @@ -372,6 +376,7 @@ class LassoGuardrail(CustomGuardrail): data: dict, cache: DualCache, message_type: Literal["PROMPT", "COMPLETION"] = "PROMPT", + call_type: str | None = None, ): """ Run the Lasso guardrail with the specified message type. @@ -392,6 +397,7 @@ class LassoGuardrail(CustomGuardrail): data: The request data containing messages cache: The cache instance for storing conversation_id (optional for post-call) message_type: Either "PROMPT" for input or "COMPLETION" for output + call_type: The request's call type, which decides whether Lasso lists itself as applied Raises: LassoGuardrailAPIError: If the Lasso API call fails @@ -409,6 +415,8 @@ class LassoGuardrail(CustomGuardrail): messages: Final[Sequence[Mapping[str, object]]] = [*expanded_messages, *input_messages] if not messages: return data + if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) # Lasso's classifix endpoint returns masked text that we copy back # into ``data["messages"]``. For multimodal/Responses-API input we diff --git a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py index b9871b1e276..5eb8e3df6fd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -185,6 +185,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): """ use_native_lifecycle_hooks: ClassVar[bool] = True + adds_itself_to_applied_guardrails: ClassVar[bool] = True @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: @@ -291,7 +292,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): return _redact_scanned_content(armor_response) return armor_response - def _raise_if_fail_closed(self, e: ModelArmorAPIError) -> None: + def _raise_if_fail_closed(self, e: ModelArmorAPIError | httpx.HTTPError) -> None: if self.optional_params.get("fail_on_error", True): raise e from None @@ -1325,8 +1326,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): end_time=error_end_time, duration=error_end_time - start_time, ) + self._raise_if_fail_closed(e) return inputs + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) + + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) flagged: Final = self._should_block_content(armor_response, allow_sanitization=False) end_time: Final = time.time() self.add_standard_logging_guardrail_information_to_request_data( diff --git a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py index 2194238cea0..808bb68acf3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py @@ -167,6 +167,8 @@ class PangeaHandler(CustomGuardrail): messages = transformer.get_messages() else: messages = data.get("messages") + if not messages: + return ai_guard_payload: Final = { "debug": False, diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index 3f025cfc8b3..2444eee6abb 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -22,6 +22,7 @@ from litellm import DualCache from litellm._logging import verbose_proxy_logger from litellm._version import version as litellm_version from litellm.integrations.custom_guardrail import ( + APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail, log_guardrail_information, ) @@ -351,8 +352,8 @@ class PillarGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Pillar Guardrail: Pre-call hook") result: Final = await self.run_pillar_guardrail(data, user_api_key_dict) - # Add guardrail name to response headers - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + if data.get("messages") or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return result @@ -394,8 +395,8 @@ class PillarGuardrail(CustomGuardrail): verbose_proxy_logger.debug("Pillar Guardrail: During-call moderation hook") result: Final = await self.run_pillar_guardrail(data, user_api_key_dict) - # Add guardrail name to response headers - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) + if data.get("messages") or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return result 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 5ce2dfb3b87..bbb03af2119 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -20,10 +20,10 @@ from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.cost_calculator import infer_call_type -from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_guardrail import APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route -from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket +from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_or_create_metadata_bucket from litellm.llms import get_guardrail_translation_mapping, load_guardrail_translation_mappings from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks @@ -222,16 +222,25 @@ def _a2a_jsonrpc_error_chunk(exc: HTTPException, request_id: str | None) -> Mapp } -def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> None: - """Populate data['litellm_metadata'] from user_api_key_dict if absent.""" - if "litellm_metadata" not in data: - from litellm.llms.base_llm.guardrail_translation.base_translation import ( - BaseTranslation, - ) +def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: str | None = None) -> None: + """Populate data['litellm_metadata'] from user_api_key_dict if absent. - user_metadata: Final = BaseTranslation.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - data["litellm_metadata"] = user_metadata + Readers pick ``litellm_metadata`` once it exists, so on routes whose bucket is ``metadata`` it starts as a copy. + """ + if "litellm_metadata" in data: + return + from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + ) + + user_metadata: Final = BaseTranslation.transform_user_api_key_dict_to_metadata(user_api_key_dict) + request: Final = cast("Mapping[str, object]", data) # cast-ok: request bodies are str-keyed + route_metadata: Final = ( + as_str_mapping(request.get("metadata")) if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES else None + ) + seeded: Final = {**route_metadata, **user_metadata} if route_metadata is not None else user_metadata + if seeded: + data["litellm_metadata"] = seeded class UnifiedLLMGuardrails(CustomLogger): @@ -291,9 +300,10 @@ class UnifiedLLMGuardrails(CustomLogger): except ValueError: return data # handle unmapped call types - endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) + translation: Final = mappings[CallTypes(call_type)]() + endpoint_translation: Final = _as_endpoint_translation(translation) - _ensure_litellm_metadata(data, user_api_key_dict) + _ensure_litellm_metadata(data, user_api_key_dict, call_type) await _scan_request_attachments(guardrail_to_apply, data, call_type) # pyright: ignore[reportUnknownArgumentType] # hook data is an untyped dict @@ -303,7 +313,11 @@ class UnifiedLLMGuardrails(CustomLogger): litellm_logging_obj=data.get("litellm_logging_obj"), ) - if not guardrail_to_apply.records_own_guardrail_information: + if not ( + guardrail_to_apply.records_own_guardrail_information + or guardrail_to_apply.adds_itself_to_applied_guardrails + or translation.lists_guardrail_when_called + ): add_guardrail_to_applied_guardrails_header( request_data=data, guardrail_name=guardrail_to_apply.guardrail_name ) @@ -341,7 +355,7 @@ class UnifiedLLMGuardrails(CustomLogger): endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]()) - _ensure_litellm_metadata(data, user_api_key_dict) + _ensure_litellm_metadata(data, user_api_key_dict, call_type) return await endpoint_translation.process_input_messages( data=data, @@ -422,7 +436,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( @@ -440,7 +458,9 @@ class UnifiedLLMGuardrails(CustomLogger): if e.original_response is None: e.original_response = response raise - if not guardrail_to_apply.records_own_guardrail_information: + if not ( + guardrail_to_apply.records_own_guardrail_information or guardrail_to_apply.adds_itself_to_applied_guardrails + ): add_guardrail_to_applied_guardrails_header( request_data=data, guardrail_name=guardrail_to_apply.guardrail_name ) 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 9e125b95a68..175a72012c7 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,9 @@ _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, ...]) +_REQUEST_SCAN_HOOKS: Final = (GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call) +_SHIPPED_GUARDRAIL_MODULES: Final = ("litellm.", "litellm_enterprise.") _ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"] _ViewSetupAttempt: TypeAlias = Literal["ready", "table_missing"] | Exception @@ -726,6 +731,45 @@ def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple return (guardrails, others) +def _scans_requests( + guardrail: CustomGuardrail, + data: dict[str, object], # mutable-ok: the proxy's request dict +) -> bool: + return any(guardrail.should_run_guardrail(data=data, event_type=hook) is True for hook in _REQUEST_SCAN_HOOKS) + + +def _warn_on_guardrails_that_cannot_scan( + data: dict[str, object], # mutable-ok: the proxy's request dict + call_type: str, +) -> None: + 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 "apply_guardrail" not in type(callback).__dict__ + and _scans_requests(callback, data) + ) + if unscannable: + verbose_proxy_logger.warning( + "Guardrails %s do not implement apply_guardrail, so on %s requests they scan only the fields their own hook reads", + ", ".join(unscannable), + call_type, + ) + + +def _requested_user_defined_guardrails(requested: frozenset[str], call_type: str) -> tuple[str, ...]: + return 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.guardrail_name in requested + and not runs_through_apply_guardrail(callback, call_type) + and not type(callback).__module__.startswith(_SHIPPED_GUARDRAIL_MODULES) + ) + + def _merge_pipeline_metadata_bucket(data: dict, bucket_key: str, modified_bucket_value: object) -> None: if not isinstance(modified_bucket_value, dict): return @@ -2374,11 +2418,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: @@ -2703,7 +2743,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, @@ -2718,6 +2758,22 @@ 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: + request: Final = cast("dict[str, object]", data) # cast-ok: request bodies are str-keyed + _warn_on_guardrails_that_cannot_scan(request, call_type) + requested: Final = frozenset( + name + for name in ( + _REQUESTED_GUARDRAILS.validate_python(guardrails_in_metadata) + if isinstance(guardrails_in_metadata, list) + else () + ) + if isinstance(name, str) + ) + for guardrail_name in _requested_user_defined_guardrails(requested, call_type): + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_name) + return + if guardrails_in_metadata and isinstance(guardrails_in_metadata, list): applied_guardrails = [] if isinstance(metadata_standard, dict) and "applied_guardrails" in metadata_standard: @@ -3019,7 +3075,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, ...]] = ( @@ -3123,17 +3179,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( @@ -3450,14 +3506,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: @@ -3573,8 +3629,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 ): @@ -4144,6 +4199,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) @@ -4175,7 +4231,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, @@ -4246,11 +4302,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/code_coverage_tests/recursive_detector.py b/tests/code_coverage_tests/recursive_detector.py index c943ff2cbd0..73059bf9d8e 100644 --- a/tests/code_coverage_tests/recursive_detector.py +++ b/tests/code_coverage_tests/recursive_detector.py @@ -3,6 +3,8 @@ import os IGNORE_FUNCTIONS = [ "_json_cost", # bounded at depth 32 and consumed under byte/node limits. + "_json_text_forms", # walks a value TypeAdapter(JsonValue) already accepted, which caps depth near 255. + "_guardrailed_json", # walks a value TypeAdapter(JsonValue) already accepted, which caps depth near 255. "_format_type", "remove_additional_properties", "remove_strict_from_schema", diff --git a/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py b/tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py index 9816eb91102..7d2f3b091bc 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,99 @@ 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: pytest.MonkeyPatch, permissions: dict[str, bool], expected_state: str, listed: bool +) -> None: + import litellm + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "callbacks", [_guardrail()]) + key: Final = UserAPIKeyAuth(api_key="sk-9876", permissions=permissions) + data: Final = { + "model": "invoice-reader", + "state": f"my key is {AWS_KEY}", + "fields": {"total": {"type": "number"}}, + "metadata": {"user_api_key_auth": key}, + } + + result: Final = 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: Final = (result.get("metadata") or {}, result.get("litellm_metadata") or {}) + applied: Final = [*buckets[0].get("applied_guardrails", []), *buckets[1].get("applied_guardrails", [])] + assert ("hide-secrets" in applied) is listed + + +@pytest.mark.asyncio +async def test_proxy_pre_call_hook_redacts_extraction_state_when_the_request_carries_no_key( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "callbacks", [_guardrail()]) + + result: Final = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data={ + "model": "invoice-reader", + "state": f"my key is {AWS_KEY}", + "fields": {"total": {"type": "number"}}, + "metadata": {}, + }, + call_type="aextraction", + ) + + assert result["state"] == "my key is [REDACTED]" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "secret_entry", + ( + {"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"}, + {"db_password": "hunter2-Zq9!xk"}, + ), +) +async def test_proxy_pre_call_hook_redacts_a_systemone_secret_stored_under_a_json_key( + monkeypatch: pytest.MonkeyPatch, secret_entry: dict[str, str] +) -> None: + import litellm + from litellm.proxy.utils import ProxyLogging + + monkeypatch.setattr(litellm, "callbacks", [_guardrail()]) + key: Final = UserAPIKeyAuth(api_key="sk-9876") + + result: Final = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=key, + data={ + "model": "judge", + "state": {**secret_entry, "ticket": "refund order 7"}, + "questions": {"leak": {"type": "noul", "instructions": "Does the state hold a credential?"}}, + "metadata": {"user_api_key_auth": key}, + }, + call_type="asystemone", + ) + + assert result["state"] == {**dict.fromkeys(secret_entry, "[REDACTED]"), "ticket": "refund order 7"} + + @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..38b472b1d02 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[str, object], + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = 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..d971f750d4c --- /dev/null +++ b/tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py @@ -0,0 +1,467 @@ +import re +from functools import reduce +from typing import Final, Literal + +import pytest +from litellm_enterprise.enterprise_callbacks.secret_detection import _ENTERPRISE_SecretDetection +from pydantic import JsonValue + +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content +from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms import get_guardrail_translation_mapping +from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite +from litellm.llms.openai.decisions.guardrail_translation.handler import DecisionsHandler, SystemOneHandler +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail +from litellm.types.guardrails import ContentFilterAction, ContentFilterPattern +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 _MasksSecrets(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + secrets: Final = ( + ("AKIAIOSFODNN7EXAMPLE", ""), + ("AKIAI44QH8DHBEXAMPLE", ""), + ("4111111111111111", ""), + ) + return GenericGuardrailAPIInputs( + texts=[reduce(lambda text, secret: text.replace(*secret), secrets, text) for text in inputs.get("texts", [])] + ) + + +class _DropsQuotes(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(texts=[text.replace('"', "") for text in inputs.get("texts", [])]) + + +class _MasksOverlappingPairs(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( + texts=[text.replace("b\nc", "[BC]").replace("a b", "[AB]") for text in inputs.get("texts", [])] + ) + + +class _NullsAMaskedValue(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( + texts=[re.sub(r'": ".*"$', '": null', text).replace("b", "[B]") for text in inputs.get("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_text_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"}, + { + "type": "choice", + "name": "team", + "instructions": "Which team?", + "choices": [{"value": "billing", "description": "Payment issues"}, {"value": True}], + }, + { + "type": "score", + "name": "mood", + "instructions": "How upset?", + "levels": [{"label": "calm"}, {"label": "angry", "description": "Uses caps"}], + }, + ], + } + + result: Final = await DecisionsHandler().process_input_messages(data, guardrail) + + assert guardrail.seen == [ + ( + "request", + ("ticket one", "ticket two", "Is it urgent?", "Which team?", "Payment issues", "How upset?", "Uses caps"), + ) + ] + 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]"}, + { + "type": "choice", + "name": "team", + "instructions": "Which team? [GUARDRAILED]", + "choices": [{"value": "billing", "description": "Payment issues [GUARDRAILED]"}, {"value": True}], + }, + { + "type": "score", + "name": "mood", + "instructions": "How upset? [GUARDRAILED]", + "levels": [{"label": "calm"}, {"label": "angry", "description": "Uses caps [GUARDRAILED]"}], + }, + ], + } + assert data["input"][0]["content"][0]["text"] == "ticket one" + + +@pytest.mark.asyncio +async def test_systemone_state_keys_scalars_and_question_text_are_guardrailed() -> None: + guardrail: Final = _TaggingGuardrail() + 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": {"type": "noul", "criteria": {"true": "The customer asks for money back", "false": None}}, + "tier": {"type": "score", "criteria": ["calm", {"level": "angry"}]}, + }, + } + + result: Final = await SystemOneHandler().process_input_messages(data, guardrail) + + assert guardrail.seen == [ + ( + "request", + ( + '"subject": "refund"', + "lines", + "order 7", + "3", + "meta", + '"id": null', + '"ask": "Does the customer want a refund?"', + "The customer asks for money back", + "calm", + '"level": "angry"', + ), + ) + ] + assert {key: value for key, value in result.items() if key != "metadata"} == { + "model": "judge", + "state": { + "subject": '"refund" [GUARDRAILED]', + "lines [GUARDRAILED]": ["order 7 [GUARDRAILED]", "", "3 [GUARDRAILED]", True], + "meta [GUARDRAILED]": {"id": "null [GUARDRAILED]"}, + }, + "questions": { + "refund": { + "type": "noul", + "instructions": {"ask": '"Does the customer want a refund?" [GUARDRAILED]'}, + }, + "routed": { + "type": "noul", + "criteria": {"true": "The customer asks for money back [GUARDRAILED]", "false": None}, + }, + "tier": {"type": "score", "criteria": ["calm [GUARDRAILED]", {"level": '"angry" [GUARDRAILED]'}]}, + }, + } + + +@pytest.mark.asyncio +async def test_systemone_secret_key_and_card_number_are_masked_before_the_provider() -> None: + data: Final = { + "model": "judge", + "state": {"AKIAIOSFODNN7EXAMPLE": "", "card": 4111111111111111, "items": 2}, + "questions": {"q": {"type": "noul", "criteria": {}}}, + } + + result: Final = await SystemOneHandler().process_input_messages(data, _MasksSecrets(guardrail_name="mask")) + + assert result["state"] == {"": "", "card": "", "items": 2} + + +@pytest.mark.asyncio +async def test_systemone_masks_a_card_number_that_json_escaping_splits() -> None: + guardrail: Final = ContentFilterGuardrail( + guardrail_name="cards", + patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="visa", action=ContentFilterAction.MASK)], + ) + data: Final = { + "model": "judge", + "state": {"card": "4111\t1111\t1111\t1111", "4111\n1111\n1111\n1111": "on file", "note": "first\nsecond"}, + "questions": {"q": {"type": "noul", "criteria": {}}}, + } + + result: Final = await SystemOneHandler().process_input_messages(data, guardrail) + + assert result["state"] == {"card": "[VISA_REDACTED]", "[VISA_REDACTED]": "on file", "note": "first\nsecond"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("pattern_names", "state", "masked"), + ( + pytest.param( + ("visa",), + {"4111111111111111": "4111\t1111\t1111\t1111"}, + {"[VISA_REDACTED]": "[VISA_REDACTED]"}, + id="key-and-escaped-value", + ), + pytest.param( + ("email",), + {"note": "Please call me back.\ndana.quill@example.com"}, + {"note": "Please call me back.\n[EMAIL_REDACTED]"}, + id="match-after-an-escape", + ), + pytest.param( + ("au_tfn", "email"), + {"TFN": "123456789 contact dana@example.com\n"}, + {"TFN": "[AU_TFN_REDACTED] contact [EMAIL_REDACTED]\n"}, + id="key-context-match-in-an-escaped-value", + ), + pytest.param( + ("au_tfn", "visa"), + {"TFN": "123456789\n4111\t1111\t1111\t1111"}, + {"TFN": "[AU_TFN_REDACTED]\n[VISA_REDACTED]"}, + id="key-context-match-beside-a-split-match", + ), + ), +) +async def test_systemone_masks_findings_on_both_sides_of_a_json_escape( + pattern_names: tuple[str, ...], state: dict[str, JsonValue], masked: dict[str, JsonValue] +) -> None: + guardrail: Final = ContentFilterGuardrail( + guardrail_name="mask", + patterns=[ + ContentFilterPattern(pattern_type="prebuilt", pattern_name=name, action=ContentFilterAction.MASK) + for name in pattern_names + ], + ) + data: Final = {"model": "judge", "state": state, "questions": {"q": {"type": "noul", "criteria": {}}}} + + result: Final = await SystemOneHandler().process_input_messages(data, guardrail) + + assert result["state"] == masked + + +@pytest.mark.asyncio +async def test_systemone_keeps_a_secret_mask_that_covers_an_escaped_value() -> None: + data: Final = { + "model": "judge", + "state": {"password": "hunter2\nAKIAIOSFODNN7EXAMPLE"}, + "questions": {"q": {"type": "noul", "criteria": {}}}, + } + + result: Final = await SystemOneHandler().process_input_messages( + data, _ENTERPRISE_SecretDetection(guardrail_name="hide-secrets") + ) + + assert result["state"] == {"password": "[REDACTED]"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("guardrail", "note"), + ( + pytest.param(_MasksOverlappingPairs(guardrail_name="pairs"), "a b\nc", id="partly-overlapping-masks"), + pytest.param(_NullsAMaskedValue(guardrail_name="nulls"), "a\nb", id="value-turned-into-null"), + ), +) +async def test_systemone_rejects_an_escaped_entry_whose_two_rewrites_cannot_merge( + guardrail: CustomGuardrail, note: str +) -> None: + data: Final = {"model": "judge", "state": {"note": note}, "questions": {"q": {"type": "noul", "criteria": {}}}} + + with pytest.raises(UnappliableRequestRewrite): + await SystemOneHandler().process_input_messages(data, guardrail) + + +@pytest.mark.asyncio +async def test_systemone_keeps_every_state_entry_when_masked_keys_collide() -> None: + data: Final = { + "model": "judge", + "state": {"AKIAIOSFODNN7EXAMPLE": "first", "items": 2, "AKIAI44QH8DHBEXAMPLE": "second"}, + "questions": {"q": {"type": "noul", "criteria": {}}}, + } + + result: Final = await SystemOneHandler().process_input_messages(data, _MasksSecrets(guardrail_name="mask")) + + assert result["state"] == {"": "first", "items": 2, "_2": "second"} + + +@pytest.mark.asyncio +async def test_systemone_rejects_a_rewrite_that_loses_the_state_key() -> None: + data: Final = { + "model": "judge", + "state": {"ticket": "refund order 7"}, + "questions": {"q": {"type": "noul", "criteria": {}}}, + } + + with pytest.raises(UnappliableRequestRewrite): + await SystemOneHandler().process_input_messages(data, _DropsQuotes(guardrail_name="drops")) + + +@pytest.mark.asyncio +async def test_systemone_rejects_masked_keys_that_stay_ambiguous() -> None: + data: Final = { + "model": "judge", + "state": {"AKIAIOSFODNN7EXAMPLE": "first", "AKIAI44QH8DHBEXAMPLE": "second", "_2": "third"}, + "questions": {"q": {"type": "noul", "criteria": {}}}, + } + + with pytest.raises(UnappliableRequestRewrite): + await SystemOneHandler().process_input_messages(data, _MasksSecrets(guardrail_name="mask")) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("handler", "data"), + ( + ( + SystemOneHandler, + {"model": "judge", "state": {}, "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 +@pytest.mark.parametrize( + ("handler", "data"), + ( + ( + DecisionsHandler, + {"model": "judge", "input": "ticket", "questions": [{"type": "predicate", "instructions": "Is it urgent?"}]}, + ), + ( + SystemOneHandler, + {"model": "judge", "state": {"total": 3, "paid": 2.5}, "questions": {"q": {"type": "noul", "criteria": {}}}}, + ), + ), +) +async def test_guardrail_that_returns_no_texts_leaves_the_request_text_as_sent( + handler: type[DecisionsHandler] | type[SystemOneHandler], data: dict[str, object] +) -> None: + result: Final = await handler().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 == [] + + +def test_router_sees_the_image_in_a_decision_input_routed_by_its_proxy_route() -> None: + structured: Final = resolve_structured_messages( + messages=None, + request_kwargs={ + "input": [ + { + "role": "user", + "content": [ + {"type": "input_text", "text": "Is this a cat?"}, + {"type": "input_image", "image_url": "https://example.com/cat.png"}, + ], + } + ], + "litellm_metadata": {"user_api_key_request_route": "/v1/decisions"}, + }, + ) + + assert structured is not None and request_contains_image_content(structured) diff --git a/tests/unit/proxy/common_utils/test_callback_utils.py b/tests/unit/proxy/common_utils/test_callback_utils.py index c8d909ac2eb..4f22a698715 100644 --- a/tests/unit/proxy/common_utils/test_callback_utils.py +++ b/tests/unit/proxy/common_utils/test_callback_utils.py @@ -205,6 +205,28 @@ def test_get_logging_caching_headers_merges_metadata_and_litellm_metadata(): assert headers["x-litellm-policy-sources"] == "global-baseline=team_default" +@pytest.mark.parametrize( + ("metadata_guardrails", "litellm_metadata_guardrails", "expected"), + ( + (["lakera-v1"], ["hide-secrets"], "lakera-v1,hide-secrets"), + (["hide-secrets"], ["hide-secrets", "lakera-v1"], "hide-secrets,lakera-v1"), + ), +) +def test_applied_guardrails_header_lists_guardrails_from_both_metadata_buckets( + metadata_guardrails: list[str], litellm_metadata_guardrails: list[str], expected: str +) -> None: + request_data: Final = { + "metadata": {"applied_guardrails": metadata_guardrails}, + "litellm_metadata": {"applied_guardrails": litellm_metadata_guardrails, "guardrail_scan_ids": ("scan-1",)}, + } + + headers: Final = get_logging_caching_headers(request_data) + + assert headers is not None + assert headers["x-litellm-applied-guardrails"] == expected + assert headers["x-litellm-guardrail-scan-id"] == "scan-1" + + def _record( request_data: dict[str, object], scan_id: str | None, diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py b/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py index 80bde77afdc..1f691540145 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py @@ -7,12 +7,15 @@ import json import os from unittest.mock import AsyncMock, MagicMock, patch +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import Final import pytest import respx import litellm from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import CallTypesLiteral from litellm.proxy.guardrails.guardrail_hooks.dynamoai import DynamoAIGuardrails from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -203,6 +206,96 @@ async def test_post_call_hook_skips_dynamoai_when_the_reply_has_no_text(respx_mo assert route.called is False + +_INPUT_HOOKS = [ + pytest.param( + "pre_call", + lambda guardrail, data, call_type: guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type + ), + id="pre_call", + ), + pytest.param( + "during_call", + lambda guardrail, data, call_type: guardrail.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type + ), + id="during_call", + ), +] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS) +@pytest.mark.parametrize( + ("request_fields", "call_type", "applied"), + [ + pytest.param({"messages": [{"role": "user", "content": "hi"}]}, "completion", ["test-dynamoai"], id="messages"), + pytest.param({"state": "hi"}, "aextraction", None, id="no-messages-on-extraction"), + pytest.param({"state": "hi"}, "completion", ["test-dynamoai"], id="no-messages-on-chat"), + ], +) +async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( + respx_mock: respx.MockRouter, + event_hook: str, + run_hook: Callable[[DynamoAIGuardrails, Mapping[str, object], CallTypesLiteral], Awaitable[object]], + request_fields: Mapping[str, object], + call_type: CallTypesLiteral, + applied: Sequence[str] | None, +) -> None: + respx_mock.post(_ANALYZE_URL).respond(json={"finalAction": "NONE", "appliedPolicies": []}) + guardrail: Final = DynamoAIGuardrails( + guardrail_name="test-dynamoai", + api_key="test-api-key", + api_base="https://api.dynamo.ai", + event_hook=event_hook, + default_on=True, + ) + metadata: Final[dict[str, object]] = {} + + await run_hook(guardrail, {**request_fields, "metadata": metadata}, call_type) + + assert metadata.get("applied_guardrails") == applied + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize( + ("response", "request_route", "applied"), + [ + pytest.param(_assistant_reply("Hello, how are you?"), "/v1/chat/completions", ["test-dynamoai"], id="chat-reply"), + pytest.param( + {"answers": {"aws_key": {"type": "text", "value": None}}}, + "/v1/extraction", + None, + id="no-choices-on-extraction", + ), + pytest.param( + {"answers": {"aws_key": {"type": "text", "value": None}}}, + "/v1/responses", + ["test-dynamoai"], + id="no-choices-on-responses", + ), + ], +) +async def test_post_call_hook_lists_the_guardrail_only_when_it_scanned_text( + respx_mock: respx.MockRouter, + response: litellm.ModelResponse | Mapping[str, object], + request_route: str, + applied: Sequence[str] | None, +) -> None: + respx_mock.post(_ANALYZE_URL).respond(json={"finalAction": "NONE", "appliedPolicies": []}) + metadata: Final[dict[str, object]] = {} + + await _post_call_guardrail().async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "hi"}], "metadata": metadata}, + user_api_key_dict=UserAPIKeyAuth(request_route=request_route), + response=response, + ) + + assert metadata.get("applied_guardrails") == applied + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py b/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py index 4a591c08bf3..1396c0599ac 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py @@ -1,10 +1,15 @@ import json +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import Final import pytest import respx +import litellm + from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import CallTypesLiteral from litellm.proxy.guardrails.guardrail_hooks.ibm_guardrails.ibm_detector import IBMGuardrailDetector pytestmark = pytest.mark.usefixtures("httpx_transport") @@ -122,3 +127,107 @@ async def test_orchestrator_detection_blocks_the_request_and_names_the_detector( ) assert "- HAP (detector: hap-detector, score: 0.970)\n Text: 'hello'" in str(blocked.value) + + +_INPUT_HOOKS = [ + pytest.param( + "pre_call", + lambda guardrail, data, call_type: guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type + ), + id="pre_call", + ), + pytest.param( + "during_call", + lambda guardrail, data, call_type: guardrail.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type + ), + id="during_call", + ), +] +_REQUEST_TEXT_CASES = [ + pytest.param({"messages": [{"role": "user", "content": "hello there"}]}, "completion", ["ibm-guard"], id="messages"), + pytest.param({"state": "hello there"}, "aextraction", None, id="no-messages-on-extraction"), + pytest.param({"state": "hello there"}, "completion", ["ibm-guard"], id="no-messages-on-chat"), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS) +@pytest.mark.parametrize(("request_fields", "call_type", "applied"), _REQUEST_TEXT_CASES) +async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( + respx_mock: respx.MockRouter, + event_hook: str, + run_hook: Callable[[IBMGuardrailDetector, Mapping[str, object], CallTypesLiteral], Awaitable[object]], + request_fields: Mapping[str, object], + call_type: CallTypesLiteral, + applied: Sequence[str] | None, +) -> None: + respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[]]) + guardrail: Final = IBMGuardrailDetector( + guardrail_name="ibm-guard", + auth_token="ibm-token", + base_url=_BASE_URL, + detector_id="hap-detector", + is_detector_server=True, + event_hook=event_hook, + default_on=True, + ) + metadata: Final[dict[str, object]] = {} + + await run_hook(guardrail, {**request_fields, "metadata": metadata}, call_type) + + assert metadata.get("applied_guardrails") == applied + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "request_route", "applied"), + [ + pytest.param( + litellm.ModelResponse( + choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}] + ), + "/v1/chat/completions", + ["ibm-guard"], + id="chat-reply", + ), + pytest.param( + {"answers": {"aws_key": {"type": "text", "value": None}}}, + "/v1/extraction", + None, + id="no-choices-on-extraction", + ), + pytest.param( + {"answers": {"aws_key": {"type": "text", "value": None}}}, + "/v1/responses", + ["ibm-guard"], + id="no-choices-on-responses", + ), + ], +) +async def test_post_call_hook_lists_the_guardrail_only_when_it_scanned_text( + respx_mock: respx.MockRouter, + response: litellm.ModelResponse | Mapping[str, object], + request_route: str, + applied: Sequence[str] | None, +) -> None: + respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[]]) + guardrail: Final = IBMGuardrailDetector( + guardrail_name="ibm-guard", + auth_token="ibm-token", + base_url=_BASE_URL, + detector_id="hap-detector", + is_detector_server=True, + event_hook="post_call", + default_on=True, + ) + metadata: Final[dict[str, object]] = {} + + await guardrail.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "hello"}], "metadata": metadata}, + user_api_key_dict=UserAPIKeyAuth(request_route=request_route), + response=response, + ) + + assert metadata.get("applied_guardrails") == applied diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py index 276fcb86e4d..20b9f0cc726 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py @@ -7,6 +7,9 @@ from copy import deepcopy from unittest.mock import AsyncMock, patch import httpx +from collections.abc import Mapping, Sequence +from typing import Final +from litellm.types.utils import CallTypesLiteral import pytest import respx from httpx import Request, Response @@ -316,7 +319,7 @@ async def test_aim_anonymize_action_blocks_malformed_redacted_messages( ) assert exc_info.value.code == "400" - assert data == request_body + assert {key: value for key, value in data.items() if key != "metadata"} == request_body _OUTPUT_REQUEST = { @@ -815,3 +818,39 @@ async def test_aim_streaming_block_stops_after_verified_chunks(monkeypatch: pyte {"choices": [{"delta": {"content": "A"}}]}, {"done": True}, ] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize( + ("call_type", "request_fields", "aim_calls", "applied"), + [ + pytest.param( + "adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, 1, ["aim"], id="decisions" + ), + pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"), + pytest.param("acompletion", {"messages": [{"role": "user", "content": "Refund?"}]}, 1, None, id="chat"), + ], +) +async def test_aim_calls_aim_and_lists_itself_only_when_the_request_has_text( + respx_mock: respx.MockRouter, + call_type: CallTypesLiteral, + request_fields: Mapping[str, object], + aim_calls: int, + applied: Sequence[str] | None, +) -> None: + analyze: Final = respx_mock.post("http://aim.test/fw/v1/analyze").respond( + json={"analysis_result": {"analysis_time_ms": 1, "policy_drill_down": {}}, "required_action": None} + ) + guardrail: Final = AimGuardrail(api_key="hs-aim-key", api_base="http://aim.test", guardrail_name="aim") + metadata: Final[dict[str, object]] = {} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data={**request_fields, "metadata": metadata}, + call_type=call_type, + ) + + assert analyze.call_count == aim_calls + assert metadata.get("applied_guardrails") == applied diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py index 349030c6c75..8eba3360320 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py @@ -3,7 +3,11 @@ import json import ssl from unittest.mock import AsyncMock, MagicMock, patch +from collections.abc import Mapping, Sequence +from typing import Final +from litellm.types.utils import CallTypesLiteral import pytest +import respx from fastapi.exceptions import HTTPException from httpx import Request, Response from websockets.exceptions import ConnectionClosed @@ -2785,3 +2789,39 @@ async def test_cato_still_inspects_every_conversational_call_type(call_type: str ) mock_post.assert_called_once() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize( + ("call_type", "request_fields", "cato_calls", "applied"), + [ + pytest.param( + "adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, 1, ["cato"], id="decisions" + ), + pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"), + pytest.param("acompletion", {"messages": [{"role": "user", "content": "Refund?"}]}, 1, None, id="chat"), + ], +) +async def test_cato_calls_cato_and_lists_itself_only_when_the_request_has_text( + respx_mock: respx.MockRouter, + call_type: CallTypesLiteral, + request_fields: Mapping[str, object], + cato_calls: int, + applied: Sequence[str] | None, +) -> None: + analyze: Final = respx_mock.post("http://cato.test/fw/v1/analyze").respond( + json={"analysis_result": {}, "required_action": None} + ) + guardrail: Final = CatoNetworksGuardrail(api_key="cato-key", api_base="http://cato.test", guardrail_name="cato") + metadata: Final[dict[str, object]] = {} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data={**request_fields, "metadata": metadata}, + call_type=call_type, + ) + + assert analyze.call_count == cato_calls + assert metadata.get("applied_guardrails") == applied diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py index 06fbdc972df..7ceb871a3ef 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py @@ -3,7 +3,7 @@ from __future__ import annotations import asyncio -from collections.abc import Callable +from collections.abc import Callable, Sequence from typing import Final, Literal, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -1308,3 +1308,26 @@ async def test_tool_call_in_both_tool_calls_and_structured_messages_is_screened_ await guardrail.apply_guardrail(inputs, _request_data(), "request") assert [call.kwargs["input"] for call in router.adecisions.await_args_list] == ['send_email({"body": "hello"})'] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("texts", "unreachable", "applied"), + [ + pytest.param(["hello"], False, ["test_decision_model"], id="scanned"), + pytest.param([], False, None, id="no-text"), + pytest.param(["hello"], True, None, id="unreachable-fail-open"), + ], +) +async def test_decision_model_lists_itself_only_when_it_returned_verdicts( + texts: list[str], unreachable: bool, applied: Sequence[str] | None +) -> None: + router: Final = _decision_router(probability=0.1) + if unreachable: + router.adecisions.side_effect = httpx.ConnectError("unreachable") + guardrail: Final = _make_guardrail(router_provider=lambda: router, unreachable_fallback="fail_open") + request_data: Final = _request_data() + + await guardrail.apply_guardrail({"texts": texts}, request_data, "request") + + assert cast(dict[str, object], request_data["metadata"]).get("applied_guardrails") == applied diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py new file mode 100644 index 00000000000..e1902480cfd --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py @@ -0,0 +1,104 @@ +from collections.abc import Mapping +from itertools import chain +from typing import Final + +import httpx +import pytest +from pydantic import BaseModel, ConfigDict, TypeAdapter + +from litellm.caching import DualCache +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation +from litellm.types.guardrails import GuardrailEventHooks + +_LAKERA_REQUEST_BODY: Final = TypeAdapter(dict[str, object]) + + +class _Metadata(BaseModel): + model_config = ConfigDict(frozen=True) + + applied_guardrails: tuple[str, ...] = () + + +def _applied_guardrails(data: Mapping[str, object]) -> tuple[str, ...]: + return tuple( + chain.from_iterable( + _Metadata.model_validate(data.get(bucket) or {}).applied_guardrails + for bucket in ("metadata", "litellm_metadata") + ) + ) + + +def _lakera(transport: httpx.MockTransport) -> lakeraAI_Moderation: + return lakeraAI_Moderation( + guardrail_name="lakera-v1", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + api_key="lk-test", + api_base="https://lakera.test", + async_handler=AsyncHTTPHandler(transport=transport), + ) + + +def _refuse_any_request(request: httpx.Request) -> httpx.Response: + raise AssertionError(f"unexpected Lakera request with body {request.content!r}") + + +def _clean_scan_of(expected_input: str) -> httpx.MockTransport: + def respond(request: httpx.Request) -> httpx.Response: + assert _LAKERA_REQUEST_BODY.validate_json(request.content)["input"] == expected_input + return httpx.Response(200, json={"results": [{"flagged": False}]}) + + return httpx.MockTransport(respond) + + +@pytest.mark.asyncio +async def test_request_without_text_lakera_reads_skips_the_vendor_and_is_not_listed() -> None: + guardrail: Final = _lakera(httpx.MockTransport(_refuse_any_request)) + data: Final[dict[str, object]] = { + "model": "m", + "state": "paid with sk-live-1234", + "fields": {"total": {"type": "number"}}, + } + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), cache=DualCache(), data=data, call_type="aextraction" + ) + + assert _applied_guardrails(data) == () + + +@pytest.mark.asyncio +async def test_clean_scan_of_string_input_lists_the_guardrail() -> None: + guardrail: Final = _lakera(_clean_scan_of("is this a refund request?")) + data: Final[dict[str, object]] = { + "model": "m", + "input": "is this a refund request?", + "questions": [{"type": "predicate", "instructions": "Is it a refund request?"}], + } + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), cache=DualCache(), data=data, call_type="adecisions" + ) + + assert _applied_guardrails(data) == ("lakera-v1",) + + +@pytest.mark.asyncio +async def test_clean_scan_of_chat_messages_keeps_the_header_as_before() -> None: + scans: Final[list[httpx.Request]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + scans.append(request) + return httpx.Response(200, json={"results": [{"flagged": False}]}) + + guardrail: Final = _lakera(httpx.MockTransport(respond)) + data: Final[dict[str, object]] = {"model": "m", "messages": [{"role": "user", "content": "is this a refund?"}]} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), cache=DualCache(), data=data, call_type="completion" + ) + + assert len(scans) == 1 + assert _applied_guardrails(data) == () diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py index 23288a6136f..a2c339996b8 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -319,6 +319,25 @@ class TestAsyncPostCallSuccessHookSkipFlags: assert "a@b.com" not in result_dict["choices"][0]["message"]["content"] +@pytest.mark.asyncio +async def test_post_call_skips_a_response_with_no_assistant_text() -> None: + guardrail: Final = LakeraAIGuardrail(api_key="test_key") + data: Final = {"model": "clairic", "metadata": {}} + llm_response: Final = MagicMock() + llm_response.model_dump.return_value = {"answer": "Berlin"} + + with patch.object(guardrail, "call_v2_guard", new_callable=AsyncMock) as mock_call: + result: Final = await guardrail.async_post_call_success_hook( + data=data, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + response=llm_response, + ) + + mock_call.assert_not_called() + assert result is llm_response + assert "applied_guardrails" not in data["metadata"] + + PII_ONLY_LAKERA_RESPONSE = { "flagged": True, "breakdown": [{"detector_type": "pii/email", "detected": True, "message_id": 0}], diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py index b8a219ba863..cd62b9fbe89 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py @@ -1549,3 +1549,46 @@ async def test_api_error_handling_legacy_messages( assert "Failed to verify request safety with Lasso API" in str(excinfo.value) assert error_message in str(excinfo.value) assert route.called + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize( + ("call_type", "request_fields", "lasso_calls", "applied"), + [ + pytest.param( + "adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, 1, ["lasso"], id="decisions" + ), + pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"), + pytest.param("acompletion", {"messages": [{"role": "user", "content": "Refund?"}]}, 1, None, id="chat"), + ], +) +@pytest.mark.parametrize("event_hook", ["pre_call", "during_call"]) +async def test_lasso_lists_itself_only_when_it_classified_request_text( + respx_mock: respx.MockRouter, + event_hook: str, + call_type: str, + request_fields: dict[str, object], + lasso_calls: int, + applied: list[str] | None, +) -> None: + classify: Final = respx_mock.post("https://server.lasso.security/gateway/v3/classify").respond( + json={"deputies": {}, "findings": {}, "violations_detected": False} + ) + guardrail: Final = LassoGuardrail( + lasso_api_key="test-api-key", guardrail_name="lasso", event_hook=event_hook, default_on=True + ) + metadata: Final[dict[str, object]] = {} + data: Final = {**request_fields, "metadata": metadata} + + if event_hook == "pre_call": + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type + ) + else: + await guardrail.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type, cache=DualCache() + ) + + assert classify.call_count == lasso_calls + assert metadata.get("applied_guardrails") == applied diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py index 3da77a8707e..34641d5b854 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py @@ -1,21 +1,28 @@ import asyncio import json from types import SimpleNamespace +from collections.abc import Awaitable, Callable, Mapping from typing import Final from unittest.mock import AsyncMock import pytest from httpx import Request, Response +from pydantic import TypeAdapter import litellm from litellm.caching.caching_handler import _PENDING_CACHE_WRITES +from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.exceptions import GuardrailRaisedException +from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.llm_shield_proxy.llm_shield_proxy import ( GUARDRAIL_NAME, LLMShieldProxyGuardrail, ) from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 +from litellm.proxy.utils import ProxyLogging +from litellm.types.extraction import ExtractionAnswer, ExtractionResponse, ExtractionUsage from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ( FunctionCallArgumentsDeltaEvent, @@ -1887,7 +1894,75 @@ class TestCompletionsRestoration: assert "".join(chunk.choices[0].text or "" for chunk in out) == "Mail [EMAIL_1" +_SHIELD_TEXTS: Final = TypeAdapter(tuple[str, ...]) + + +def _pseudonymizing_shield(original: str, placeholder: str) -> Callable[..., Awaitable[Response]]: + async def post(url: str, headers: Mapping[str, str], json: Mapping[str, object], timeout: float) -> Response: + texts: Final = _SHIELD_TEXTS.validate_python(json["texts"]) + swap: Final = (original, placeholder) if url.endswith("/redact") else (placeholder, original) + return _response({"texts": [text.replace(*swap) for text in texts]}) + + return post + + class TestProxyWiring: + @pytest.mark.asyncio + async def test_extraction_request_is_redacted_and_the_answer_restored(self, monkeypatch: pytest.MonkeyPatch): + guardrail: Final = _guardrail(event_hook=["pre_call", "post_call"]) + guardrail.async_handler.post = _pseudonymizing_shield("alice@example.com", "[EMAIL_1]") + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + caller: Final = UserAPIKeyAuth(api_key="sk-9876", request_route="/v1/extraction") + + sent: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=caller, + data={ + "model": "m", + "state": "Write to alice@example.com", + "fields": {"email": {"type": "text", "instructions": "The address after alice@example.com"}}, + }, + call_type="aextraction", + ) + reply: Final = ExtractionResponse( + model="m", + answers={"email": ExtractionAnswer(type="text", value="[EMAIL_1]")}, + truncated=False, + usage=ExtractionUsage(input_tokens=1, output_tokens=1), + ) + restored: Final = await proxy_logging.post_call_success_hook( + data=sent, response=reply, user_api_key_dict=caller + ) + + assert (sent["state"], sent["fields"]["email"]["instructions"]) == ( + "Write to [EMAIL_1]", + "The address after [EMAIL_1]", + ) + assert restored.answers["email"].value == "alice@example.com" + + @pytest.mark.asyncio + async def test_decision_request_is_redacted_and_keeps_its_routing_tags(self, monkeypatch: pytest.MonkeyPatch): + guardrail: Final = _guardrail() + guardrail.async_handler.post = _pseudonymizing_shield("alice@example.com", "[EMAIL_1]") + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + sent: Final = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data={ + "model": "m", + "input": "Refund alice@example.com", + "questions": [{"type": "predicate", "instructions": "Is alice@example.com owed money?"}], + "metadata": {"tags": ["prod"]}, + }, + call_type="adecisions", + ) + + assert (sent["input"], sent["questions"][0]["instructions"]) == ( + "Refund [EMAIL_1]", + "Is [EMAIL_1] owed money?", + ) + assert sent[get_metadata_variable_name_from_kwargs(sent)]["tags"] == ["prod"] + def test_dashboard_config_model_is_exposed(self): """The guardrail garden reads the provider's fields from `get_config_model`.""" model = LLMShieldProxyGuardrail.get_config_model() diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py index 5aa3c7dfc17..507d979ac43 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -3,7 +3,7 @@ import base64 import io import json from collections.abc import Iterator, Sequence -from typing import cast +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, Mock, patch import pytest @@ -5291,6 +5291,109 @@ async def test_logging_only_flagged_prompt_still_scans_response(): assert len(flagged) == 2 +class _RedactsPhoneNumbers(AsyncHTTPHandler): + def __init__(self) -> None: + pass + + async def post( + self, + url: str, + json: dict[str, object] | None = None, + headers: dict[str, str] | None = None, + **kwargs: object, + ) -> httpx.Response: + payload: Final = cast(dict[str, dict[str, str]], json) + text: Final = (payload.get("userPromptData") or payload["modelResponseData"])["text"] + redacted: Final = text.replace("555-123-4567", "[REDACTED]") + result: Final = ( + _clean_armor_response() + if redacted == text + else { + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "sdp": { + "sdpFilterResult": { + "deidentifyResult": {"matchState": "MATCH_FOUND", "data": {"text": redacted}} + } + } + }, + } + } + ) + return httpx.Response(200, json=result, request=httpx.Request("POST", url)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("mask_request_content", (True, False)) +async def test_apply_guardrail_blocks_a_masking_finding(mask_request_content: bool) -> None: + guardrail: Final = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-block", + event_hook=GuardrailEventHooks.pre_call, + mask_request_content=mask_request_content, + async_handler=_RedactsPhoneNumbers(), + access_token_provider=_async_token_provider, + ) + + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": ["call me at 555-123-4567", "refund order 4411"]}, + request_data={"metadata": {}}, + input_type="request", + ) + + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_returns_clean_text_unchanged() -> None: + inputs: Final = {"texts": ["refund order 4411"]} + guardrail: Final = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-clean", + event_hook=GuardrailEventHooks.pre_call, + async_handler=_RedactsPhoneNumbers(), + access_token_provider=_async_token_provider, + ) + + assert await guardrail.apply_guardrail(inputs=inputs, request_data={"metadata": {}}, input_type="request") is inputs + + +class _UnavailableArmor(AsyncHTTPHandler): + def __init__(self) -> None: + pass + + async def post(self, url: str, **kwargs: object) -> httpx.Response: + return httpx.Response(503, text="unavailable", request=httpx.Request("POST", url)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_on_error", (True, False)) +async def test_apply_guardrail_honors_fail_on_error_when_model_armor_is_down(fail_on_error: bool) -> None: + guardrail: Final = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-down", + event_hook=GuardrailEventHooks.pre_call, + fail_on_error=fail_on_error, + async_handler=_UnavailableArmor(), + access_token_provider=_async_token_provider, + ) + inputs: Final = {"texts": ["refund order 4411"]} + + if fail_on_error: + with pytest.raises(ModelArmorAPIError): + await guardrail.apply_guardrail(inputs=inputs, request_data={"metadata": {}}, input_type="request") + else: + assert await guardrail.apply_guardrail(inputs=inputs, request_data={"metadata": {}}, input_type="request") is inputs + + @pytest.mark.asyncio async def test_apply_guardrail_raises_on_flagged_when_not_logging_only(): """The /guardrails/apply_guardrail endpoint calls apply_guardrail directly; a diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py index 4e31a3c2055..2ef1673717e 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py @@ -1,9 +1,14 @@ from unittest.mock import AsyncMock, patch import httpx +from collections.abc import Mapping, Sequence +from typing import Final import pytest +import respx from fastapi import HTTPException +from litellm import DualCache +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.pangea import initialize_guardrail from litellm.proxy.guardrails.guardrail_hooks.pangea.pangea import ( PangeaGuardrailMissingSecrets, @@ -368,3 +373,41 @@ async def test_pangea_ai_guard_response_transformed(pangea_guardrail): ) assert response.choices[0]["message"]["content"] == "Yes, here is an SSN: " + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize( + ("call_type", "request_fields", "pangea_calls", "applied"), + [ + pytest.param( + "acompletion", {"messages": [{"role": "user", "content": "hi"}]}, 1, ["pangea-ai-guard"], id="chat" + ), + pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"), + pytest.param("adecisions", {"input": "Is this a refund?"}, 0, None, id="decisions"), + ], +) +async def test_pangea_calls_pangea_and_lists_itself_only_when_the_request_has_messages( + respx_mock: respx.MockRouter, + call_type: str, + request_fields: Mapping[str, object], + pangea_calls: int, + applied: Sequence[str] | None, +) -> None: + guard: Final = respx_mock.post("https://ai-guard.aws.us.pangea.cloud/v1beta/guard").respond( + json={"result": {"blocked": False, "transformed": False}} + ) + guardrail: Final = PangeaHandler( + mode="pre_call", guardrail_name="pangea-ai-guard", api_key="pts_pangeatokenid", default_on=True + ) + metadata: Final[dict[str, object]] = {} + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data={**request_fields, "metadata": metadata}, + call_type=call_type, + ) + + assert guard.call_count == pangea_calls + assert metadata.get("applied_guardrails") == applied diff --git a/tests/unit/proxy/guardrails/test_pillar_guardrails.py b/tests/unit/proxy/guardrails/test_pillar_guardrails.py index c6be433399a..0a8dadd5d11 100644 --- a/tests/unit/proxy/guardrails/test_pillar_guardrails.py +++ b/tests/unit/proxy/guardrails/test_pillar_guardrails.py @@ -16,7 +16,10 @@ from unittest.mock import Mock, patch import json from urllib.parse import unquote +from collections.abc import Awaitable, Callable, Mapping, Sequence +from typing import Final import pytest +import respx from fastapi.exceptions import HTTPException from httpx import Request, Response @@ -24,6 +27,7 @@ from httpx import Request, Response import litellm from litellm import DualCache from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import CallTypesLiteral from litellm.proxy.common_utils.callback_utils import get_logging_caching_headers from litellm.proxy.guardrails.guardrail_hooks.pillar import ( PillarGuardrail, @@ -1513,3 +1517,58 @@ async def test_mcp_call_masking( if __name__ == "__main__": # Run the tests pytest.main([__file__, "-v"]) + + +_INPUT_HOOKS = [ + pytest.param( + "pre_call", + lambda guardrail, data, call_type: guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type + ), + id="pre_call", + ), + pytest.param( + "during_call", + lambda guardrail, data, call_type: guardrail.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type + ), + id="during_call", + ), +] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS) +@pytest.mark.parametrize( + ("request_fields", "call_type", "applied"), + [ + pytest.param({"messages": [{"role": "user", "content": "hi"}]}, "completion", ["pillar-test"], id="messages"), + pytest.param({"state": "hi"}, "aextraction", None, id="no-messages-on-extraction"), + pytest.param({"state": "hi"}, "completion", ["pillar-test"], id="no-messages-on-chat"), + ], +) +async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( + respx_mock: respx.MockRouter, + event_hook: str, + run_hook: Callable[[PillarGuardrail, Mapping[str, object], CallTypesLiteral], Awaitable[object]], + request_fields: Mapping[str, object], + call_type: CallTypesLiteral, + applied: Sequence[str] | None, +) -> None: + respx_mock.post(url__startswith="https://api.pillar.security/api/v1/protect").respond( + json={"session_id": "s", "flagged": False, "scanners": {}, "evidence": []} + ) + guardrail: Final = PillarGuardrail( + guardrail_name="pillar-test", + api_key="test-pillar-key", + api_base="https://api.pillar.security", + on_flagged_action="block", + event_hook=event_hook, + default_on=True, + ) + metadata: Final[dict[str, object]] = {} + + await run_hook(guardrail, {**request_fields, "metadata": metadata}, call_type) + + assert metadata.get("applied_guardrails") == applied diff --git a/tests/unit/proxy/policy_engine/test_pipeline_executor.py b/tests/unit/proxy/policy_engine/test_pipeline_executor.py index c42878bd4f7..415975a010a 100644 --- a/tests/unit/proxy/policy_engine/test_pipeline_executor.py +++ b/tests/unit/proxy/policy_engine/test_pipeline_executor.py @@ -10,6 +10,8 @@ import pickle from typing import Literal from unittest.mock import MagicMock +from collections.abc import Mapping +from typing import Final import pytest import litellm @@ -30,7 +32,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import ( GuardrailPipeline, PipelineStep, ) -from litellm.types.utils import CallTypesLiteral +from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs try: from fastapi.exceptions import HTTPException @@ -1137,8 +1139,16 @@ class _PolicyOptOutGuardrail(CustomGuardrail): def __init__(self): super().__init__(guardrail_name="policy-opt-out", default_on=True) self.native_pre_call_ran = False + self.scanned: tuple[str, ...] = () - async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: Mapping[str, object], + input_type: Literal["request", "response"], + logging_obj: object = None, + ) -> GenericGuardrailAPIInputs: + self.scanned = (*self.scanned, *inputs.get("texts", [])) return inputs async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type): @@ -1164,6 +1174,26 @@ 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: pytest.MonkeyPatch, +) -> None: + guardrail: Final = _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..70ac3d5c73c 100644 --- a/tests/unit/proxy/test_proxy_logging_hook_detection.py +++ b/tests/unit/proxy/test_proxy_logging_hook_detection.py @@ -1,14 +1,28 @@ +from collections.abc import Mapping +from itertools import chain +from typing import Final, Literal + +import httpx import pytest from fastapi import HTTPException +from pydantic import BaseModel, ConfigDict 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.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ModelArmorGuardrail +from litellm.proxy.guardrails.guardrail_hooks.tool_permission import ToolPermissionGuardrail from litellm.proxy.utils import ProxyLogging +from litellm.types.decisions import OpenAIDecisionResponse +from litellm.types.extraction import ExtractionResponse from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import CallTypesLiteral +from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs def test_has_post_call_response_headers_callbacks_ignores_empty_callbacks( @@ -1060,3 +1074,434 @@ 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: Final = "sk-live-1234" + + +class _Metadata(BaseModel): + model_config = ConfigDict(frozen=True) + + applied_guardrails: tuple[str, ...] = () + + +class _RoutedMetadata(BaseModel): + model_config = ConfigDict(frozen=True) + + tags: tuple[str, ...] = () + requester_ip_address: str | None = None + user_api_key_key_alias: str | None = None + + +def _applied_guardrails(data: Mapping[str, object]) -> tuple[str, ...]: + return tuple( + chain.from_iterable( + _Metadata.model_validate(data.get(bucket) or {}).applied_guardrails + for bucket in ("metadata", "litellm_metadata") + ) + ) + + +class _MasksSecretWithNativeHooks(_KeepsNativeHooks): + def __init__(self, **kwargs: object) -> None: + super().__init__(**kwargs) + self.scanned: 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") or ()) + self.scanned.append((input_type, texts)) + return GenericGuardrailAPIInputs(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: pytest.MonkeyPatch, + call_type: CallTypesLiteral, + data: dict[str, object], + expected: dict[str, object], +) -> None: + guardrail: Final = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + result: Final[dict[str, object]] = 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 +@pytest.mark.parametrize( + ("call_type", "request_fields"), + ( + ("aextraction", {"state": f"paid with {_SECRET}", "fields": {"total": {"type": "number"}}}), + ("asystemone", {"state": {"ticket": _SECRET}, "questions": {"leak": {"type": "noul", "instructions": "Key?"}}}), + ("adecisions", {"input": _SECRET, "questions": [{"type": "predicate", "instructions": "Is it a key?"}]}), + ), +) +async def test_pre_call_hook_keeps_request_tags_where_the_router_reads_them( + monkeypatch: pytest.MonkeyPatch, call_type: CallTypesLiteral, request_fields: Mapping[str, object] +) -> None: + guardrail: Final = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", key_alias="billing"), + data={"model": "m", **request_fields, "metadata": {"tags": ["prod"], "requester_ip_address": "10.0.0.7"}}, + call_type=call_type, + ) + + routed_metadata: Final = _RoutedMetadata.model_validate(result[get_metadata_variable_name_from_kwargs(result)]) + assert routed_metadata == _RoutedMetadata( + tags=("prod",), requester_ip_address="10.0.0.7", user_api_key_key_alias="billing" + ) + assert _applied_guardrails(result) == ("keeps_native",) + + +class _ModelArmorAnswering(AsyncHTTPHandler): + def __init__(self, status_code: int) -> None: + self.status_code = status_code + + async def post(self, url: str, **kwargs: object) -> httpx.Response: + clean: Final = {"sanitizationResult": {"filterMatchState": "NO_MATCH_FOUND", "filterResults": {}}} + return httpx.Response(self.status_code, json=clean, request=httpx.Request("POST", url)) + + +async def _model_armor_token() -> tuple[str, str]: + return ("test-token", "test-project") + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("status_code", "applied"), ((200, ("model-armor",)), (503, ()))) +@pytest.mark.parametrize( + ("call_type", "request_fields"), + ( + ("aextraction", {"state": "refund order 4411", "fields": {"total": {"type": "number"}}}), + ("asystemone", {"state": {"ticket": "refund order 4411"}, "questions": {"q": {"type": "noul"}}}), + ("adecisions", {"input": "refund order 4411", "questions": [{"type": "predicate", "instructions": "Refund?"}]}), + ), +) +async def test_pre_call_hook_lists_model_armor_only_when_it_answered( + monkeypatch: pytest.MonkeyPatch, + call_type: CallTypesLiteral, + request_fields: Mapping[str, object], + status_code: int, + applied: tuple[str, ...], +) -> None: + guardrail: Final = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + guardrail_name="model-armor", + event_hook=GuardrailEventHooks.pre_call, + default_on=True, + fail_on_error=False, + async_handler=_ModelArmorAnswering(status_code), + access_token_provider=_model_armor_token, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data={"model": "m", **request_fields}, + call_type=call_type, + ) + + assert _applied_guardrails(result) == applied + + +@pytest.mark.asyncio +async def test_pre_call_hook_does_not_list_a_guardrail_for_a_systemone_request_without_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrail: Final = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data={"model": "m", "state": {}, "questions": {"routed": {"type": "choice", "criteria": {"refund": None}}}}, + call_type="asystemone", + ) + + assert guardrail.scanned == [] + assert _applied_guardrails(result) == () + + +@pytest.mark.asyncio +async def test_during_call_hook_scans_extraction_state_with_a_native_hook_guardrail( + monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrail: Final = _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: pytest.MonkeyPatch, +) -> None: + guardrail: Final = _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: pytest.MonkeyPatch, run_in_parallel: bool +) -> None: + guardrail: Final = _MasksSecretWithNativeHooks( + event_hook=GuardrailEventHooks.post_call, default_on=True, run_in_parallel=run_in_parallel + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + response: Final = 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: str, + default_on: bool = False, + event_hook: GuardrailEventHooks = GuardrailEventHooks.pre_call, + ) -> None: + super().__init__(guardrail_name=guardrail_name, event_hook=event_hook, default_on=default_on) + self.native_hooks_ran: list[str] = [] + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + self.native_hooks_ran.append(call_type) + return data + + +_REQUEST_WITH_EVERY_TEXT_FIELD: Final = { + "model": "m", + "state": "s", + "input": "hi", + "questions": [{"type": "predicate", "instructions": "Is it a greeting?"}], + "messages": [{"role": "user", "content": "hi"}], +} + + +def _warning(call_type: str) -> str: + return ( + "Guardrails requested-native, default-native do not implement apply_guardrail, " + f"so on {call_type} requests they scan only the fields their own hook reads" + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("call_type", "warning"), + (("aextraction", _warning("aextraction")), ("adecisions", _warning("adecisions")), ("acompletion", None)), +) +async def test_requested_user_guardrail_without_apply_guardrail_stays_listed_and_is_warned_about( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, + call_type: CallTypesLiteral, + warning: str | None, +) -> None: + requested: Final = _ReadsNoDecisionText("requested-native") + default_on: Final = _ReadsNoDecisionText("default-native", default_on=True) + post_call_only: Final = _ReadsNoDecisionText( + "post-call-native", default_on=True, event_hook=GuardrailEventHooks.post_call + ) + scanner: Final = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [requested, default_on, post_call_only, scanner]) + + with caplog.at_level("WARNING"): + result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data={**_REQUEST_WITH_EVERY_TEXT_FIELD, "metadata": {"guardrails": ["requested-native"]}}, + call_type=call_type, + ) + + assert requested.native_hooks_ran == [call_type] + assert "requested-native" in _applied_guardrails(result) + warnings: Final = [ + record.getMessage() for record in caplog.records if "do not implement apply_guardrail" in record.getMessage() + ] + assert warnings == ([] if warning is None else [warning]) + + +def _tool_permission_guardrail() -> ToolPermissionGuardrail: + return ToolPermissionGuardrail( + guardrail_name="tool-permission", + event_hook=GuardrailEventHooks.pre_call, + rules=[], + default_action="allow", + on_disallowed_action="block", + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("call_type", "listed"), (("aextraction", False), ("adecisions", False), ("acompletion", True))) +async def test_requested_built_in_guardrail_is_listed_only_where_its_hook_reads_the_text( + monkeypatch: pytest.MonkeyPatch, call_type: CallTypesLiteral, listed: bool +) -> None: + monkeypatch.setattr(litellm, "callbacks", [_tool_permission_guardrail()]) + + result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data={**_REQUEST_WITH_EVERY_TEXT_FIELD, "metadata": {"guardrails": ["tool-permission"]}}, + call_type=call_type, + ) + + assert ("tool-permission" in _applied_guardrails(result)) is listed + + +class _ReroutesToPrivateModel(CustomGuardrail): + def __init__(self) -> None: + 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: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + 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_built_in_guardrail_only_where_it_reads_the_text( + monkeypatch: pytest.MonkeyPatch, call_type: CallTypesLiteral, listed: bool +) -> None: + monkeypatch.setattr(litellm, "callbacks", [_tool_permission_guardrail(), _ReroutesToPrivateModel()]) + + result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data={**_REQUEST_WITH_EVERY_TEXT_FIELD, "metadata": {"guardrails": ["tool-permission"]}}, + call_type=call_type, + ) + + assert result["model"] == "private-model" + assert ("tool-permission" in _applied_guardrails(result)) is listed + + +class _RecordsScans(_AppliesGuardrail): + def __init__(self, **kwargs: object) -> None: + super().__init__(**kwargs) + self.scanned: list[str] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + 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: pytest.MonkeyPatch, route: str +) -> None: + guardrail: Final = _RecordsScans(event_hook=GuardrailEventHooks.post_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + response: Final = OpenAIDecisionResponse.model_validate( + { + "model": "m", + "answers": [{"type": "refusal", "name": "q"}], + "usage": {"input_tokens": 1, "output_tokens": 0, "total_tokens": 1}, + } + ) + data: Final[dict[str, object]] = {"model": "m", "metadata": {}} + + result: Final = 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 5640d7d94ec..798d7a358e2 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -278,7 +278,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"], @@ -303,7 +303,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"] @@ -315,12 +315,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]