diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index 104c82812e9..6b044a87521 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -74,6 +74,8 @@ class BaseTranslation(ABC): 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/handler.py b/litellm/llms/openai/decisions/guardrail_translation/handler.py index 0244b44f688..76032483949 100644 --- a/litellm/llms/openai/decisions/guardrail_translation/handler.py +++ b/litellm/llms/openai/decisions/guardrail_translation/handler.py @@ -1,4 +1,5 @@ -from collections.abc import Mapping +import json +from collections.abc import Iterator, Mapping from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar, Final @@ -6,7 +7,7 @@ from typing import TYPE_CHECKING, ClassVar, Final from pydantic import JsonValue, TypeAdapter from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation -from litellm.proxy.guardrails._content_utils import iter_tool_json_text, map_tool_json_text +from litellm.llms.base_llm.guardrail_translation.utils import unappliable_request_rewrite from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -15,6 +16,7 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth _JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) +_DESCRIBED_ITEM_KEYS: Final = ("choices", "levels") _QUESTION_LIST_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, JsonValue], ...]]] = TypeAdapter( tuple[Mapping[str, JsonValue], ...] ) @@ -73,20 +75,123 @@ def _guardrailed_input(decision_input: JsonValue, guardrailed: Mapping[str, str] return [_guardrailed_message(message, guardrailed) for message in decision_input] -def _guardrailed_json(value: JsonValue, guardrailed: Mapping[str, str]) -> JsonValue: - return map_tool_json_text(value, lambda text: guardrailed.get(text, text)) +def _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 key + yield from _json_text_forms(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 _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(): + keys: Final = tuple(guardrailed.get(key, key) for key in value) + if len(frozenset(keys)) != len(keys): + raise unappliable_request_rewrite(guardrail_name) + return { + key: _guardrailed_json(item, guardrailed, guardrail_name) + for key, item in zip(keys, value.values(), strict=True) + } def _instruction_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]: - return tuple(iter_tool_json_text(question.get("instructions"))) + 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] + 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)} + return {**question, "instructions": _guardrailed_json(question["instructions"], guardrailed, guardrail_name)} async def _guardrailed_texts( @@ -103,6 +208,14 @@ async def _guardrailed_texts( input_type="request", logging_obj=litellm_logging_obj, ) + if not guardrail_to_apply.records_own_guardrail_information: + 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)} @@ -111,6 +224,7 @@ async def _guardrailed_texts( class _DecisionAnswersTranslation(BaseTranslation): scans_responses: ClassVar[bool] = False + lists_guardrail_when_called: ClassVar[bool] = True async def process_output_response( self, @@ -132,14 +246,17 @@ class DecisionsHandler(_DecisionAnswersTranslation): ) -> Mapping[str, object]: decision_input: Final = _JSON_ADAPTER.validate_python(data.get("input")) questions: Final = _QUESTION_LIST_ADAPTER.validate_python(data.get("questions") or ()) - texts: Final = (*_input_texts(decision_input), *chain.from_iterable(map(_instruction_texts, questions))) + 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_instructions(question, guardrailed) for question in questions], + "questions": [ + _with_guardrailed_decision_question(question, guardrailed, guardrail_to_apply.guardrail_name) + for question in questions + ], } @@ -152,17 +269,15 @@ class SystemOneHandler(_DecisionAnswersTranslation): ) -> Mapping[str, object]: state: Final = _JSON_ADAPTER.validate_python(data.get("state")) questions: Final = _QUESTION_MAP_ADAPTER.validate_python(data.get("questions") or {}) - texts: Final = ( - *iter_tool_json_text(state), - *chain.from_iterable(map(_instruction_texts, questions.values())), - ) + 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), + "state": _guardrailed_json(state, guardrailed, guardrail_to_apply.guardrail_name), "questions": { - name: _with_guardrailed_instructions(question, guardrailed) for name, question in questions.items() + name: _with_guardrailed_systemone_question(question, guardrailed, guardrail_to_apply.guardrail_name) + for name, question in questions.items() }, } diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py index 61117fbc55e..b564a1df041 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py @@ -23,6 +23,9 @@ from litellm.llms.custom_httpx.http_handler import ( 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, @@ -169,6 +172,10 @@ class AimGuardrail(CustomGuardrail): return data async def call_aim_guardrail(self, data: dict, hook: str, key_alias: str | 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 + 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 +187,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..be7d967836f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py @@ -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, @@ -295,6 +298,10 @@ class CatoNetworksGuardrail(CustomGuardrail): key_alias: str | None, user_email: str | 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 + 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 +312,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..6974c33b0bf 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py @@ -328,8 +328,7 @@ 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) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -373,8 +372,7 @@ 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) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -441,8 +439,7 @@ 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) + 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..ff4ae15b6f1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py +++ b/litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py @@ -476,8 +476,7 @@ 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) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -545,8 +544,7 @@ 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) + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return data @@ -641,8 +639,7 @@ 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) + 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/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index 97b35769f63..ca06b3c1675 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -45,6 +45,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, @@ -409,6 +412,7 @@ class LassoGuardrail(CustomGuardrail): messages: Final[Sequence[Mapping[str, object]]] = [*expanded_messages, *input_messages] if not messages: return data + 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..04f495506a7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py +++ b/litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py @@ -1,6 +1,8 @@ +import asyncio import time from collections.abc import AsyncGenerator, Awaitable, Callable, Generator, Mapping, Sequence from enum import Enum, auto +from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Protocol, runtime_checkable import httpx @@ -291,7 +293,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 @@ -1302,17 +1304,23 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): input_type: Literal["request", "response"], logging_obj: "LiteLLMLoggingObj | None" = None, ) -> GenericGuardrailAPIInputs: - content: Final = "\n".join(text for text in inputs.get("texts") or () if text) + texts: Final = tuple(inputs.get("texts") or ()) + content: Final = "\n".join(text for text in texts if text) if not content: return inputs source: Final[Literal["user_prompt", "model_response"]] = ( "user_prompt" if input_type == "request" else "model_response" ) + masking: Final = self.mask_request_content if input_type == "request" else self.mask_response_content + scanned: Final = tuple(dict.fromkeys(text for text in texts if text)) if masking else (content,) start_time: Final = time.time() try: - armor_response: Final = await self.make_model_armor_request( - content=content, source=source, request_data=request_data + armor_responses: Final = await asyncio.gather( + *( + self.make_model_armor_request(content=text, source=source, request_data=request_data) + for text in scanned + ) ) except (ModelArmorAPIError, httpx.HTTPError) as e: error_end_time: Final = time.time() @@ -1325,12 +1333,19 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): end_time=error_end_time, duration=error_end_time - start_time, ) + self._raise_if_fail_closed(e) return inputs - flagged: Final = self._should_block_content(armor_response, allow_sanitization=False) + flagged: Final = tuple( + response for response in armor_responses if self._should_block_content(response, allow_sanitization=masking) + ) end_time: Final = time.time() self.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response=self._build_logging_response(armor_response), + guardrail_json_response=( + tuple(self._build_logging_response(response) for response in armor_responses) + if masking + else self._build_logging_response(armor_responses[0]) + ), request_data=request_data, guardrail_status="guardrail_flagged" if flagged else "success", guardrail_provider="model_armor", @@ -1338,15 +1353,22 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase): end_time=end_time, duration=end_time - start_time, ) - if flagged and not self._event_hook_is_event_type(GuardrailEventHooks.logging_only): + if self._event_hook_is_event_type(GuardrailEventHooks.logging_only): + return inputs + if flagged: raise HTTPException( status_code=400, detail=self._build_block_error_detail( "Response blocked by Model Armor" if input_type == "response" else "Content blocked by Model Armor", - armor_response, + flagged[0], ), ) - return inputs + if not masking: + return inputs + sanitized: Final = MappingProxyType( + {text: self._get_sanitized_content(response) or text for text, response in zip(scanned, armor_responses)} + ) + return {**inputs, "texts": [sanitized.get(text, text) for text in texts]} @staticmethod def get_config_model() -> type["GuardrailConfigModel"] | None: 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..1983be794b3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -351,8 +351,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"): + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name) return result @@ -394,8 +394,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"): + 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 a147580a1f9..3aff4779235 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -284,7 +284,8 @@ 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) @@ -296,7 +297,7 @@ 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 and not translation.lists_guardrail_when_called: add_guardrail_to_applied_guardrails_header( request_data=data, guardrail_name=guardrail_to_apply.guardrail_name ) 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 index a2993ac1cea..b84769dc8b9 100644 --- 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 @@ -5,6 +5,7 @@ import pytest from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms import get_guardrail_translation_mapping +from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite from litellm.llms.openai.decisions.guardrail_translation.handler import DecisionsHandler, SystemOneHandler from litellm.types.decisions import OpenAIDecisionResponse from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs @@ -27,6 +28,22 @@ class _TaggingGuardrail(CustomGuardrail): 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": ""} + return GenericGuardrailAPIInputs( + texts=[ + "" if text == "4111111111111111" else secrets.get(text, text) for text in inputs.get("texts", []) + ] + ) + + class _ReturnsNoTexts(CustomGuardrail): async def apply_guardrail( self, @@ -52,7 +69,7 @@ def test_decision_call_types_resolve_to_their_handler(call_type: CallTypes, hand @pytest.mark.asyncio -async def test_decision_message_text_parts_and_question_instructions_are_guardrailed() -> None: +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 = { @@ -64,12 +81,29 @@ async def test_decision_message_text_parts_and_question_instructions_are_guardra "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?"))] + 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": [ @@ -79,48 +113,111 @@ async def test_decision_message_text_parts_and_question_instructions_are_guardra "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_leaves_and_question_instructions_are_guardrailed() -> None: +async def test_systemone_state_keys_scalars_and_question_text_are_guardrailed() -> None: guardrail: Final = _TaggingGuardrail() - criteria_question: Final = {"type": "noul", "criteria": {"yes": "refund"}} data: Final = { "model": "judge", "state": {"subject": "refund", "lines": ["order 7", "", 3, True], "meta": {"id": None}}, "questions": { "refund": {"type": "noul", "instructions": {"ask": "Does the customer want a refund?"}}, - "routed": criteria_question, + "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", ("refund", "order 7", "Does the customer want a refund?"))] + assert guardrail.seen == [ + ( + "request", + ( + "subject", + "refund", + "lines", + "order 7", + "3", + "meta", + "id", + "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": ["order 7 [GUARDRAILED]", "", 3, True], - "meta": {"id": None}, + "subject [GUARDRAILED]": "refund [GUARDRAILED]", + "lines [GUARDRAILED]": ["order 7 [GUARDRAILED]", "", "3 [GUARDRAILED]", True], + "meta [GUARDRAILED]": {"id [GUARDRAILED]": None}, }, "questions": { - "refund": {"type": "noul", "instructions": {"ask": "Does the customer want a refund? [GUARDRAILED]"}}, - "routed": criteria_question, + "refund": { + "type": "noul", + "instructions": {"ask [GUARDRAILED]": "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 [GUARDRAILED]": "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_rejects_masked_keys_that_collide() -> None: + data: Final = { + "model": "judge", + "state": {"AKIAIOSFODNN7EXAMPLE": "first", "AKIAI44QH8DHBEXAMPLE": "second"}, + "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": {"count": 3}, "questions": {"q": {"type": "noul", "criteria": {}}}}, + {"model": "judge", "state": {}, "questions": {"q": {"type": "noul", "criteria": {}}}}, ), ( DecisionsHandler, @@ -144,14 +241,23 @@ async def test_request_without_text_is_returned_unscanned( @pytest.mark.asyncio -async def test_guardrail_that_returns_no_texts_leaves_the_request_text_as_sent() -> None: - data: Final = { - "model": "judge", - "input": "ticket", - "questions": [{"type": "predicate", "instructions": "Is it urgent?"}], - } - - result: Final = await DecisionsHandler().process_input_messages(data, _ReturnsNoTexts(guardrail_name="noop")) +@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 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..7208136f69a 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py @@ -203,6 +203,76 @@ 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: guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type="completion" + ), + id="pre_call", + ), + pytest.param( + "during_call", + lambda guardrail, data: guardrail.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type="completion" + ), + 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", "applied"), + [ + pytest.param({"messages": [{"role": "user", "content": "hi"}]}, ["test-dynamoai"], id="messages"), + pytest.param({"state": "hi"}, None, id="no-messages"), + ], +) +async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( + respx_mock: respx.MockRouter, event_hook, run_hook, request_fields, applied +) -> None: + respx_mock.post(_ANALYZE_URL).respond(json={"finalAction": "NONE", "appliedPolicies": []}) + guardrail = DynamoAIGuardrails( + guardrail_name="test-dynamoai", + api_key="test-api-key", + api_base="https://api.dynamo.ai", + event_hook=event_hook, + default_on=True, + ) + metadata: dict[str, object] = {} + + await run_hook(guardrail, {**request_fields, "metadata": metadata}) + + assert metadata.get("applied_guardrails") == applied + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("httpx_transport") +@pytest.mark.parametrize( + ("response", "applied"), + [ + pytest.param(_assistant_reply("Hello, how are you?"), ["test-dynamoai"], id="chat-reply"), + pytest.param({"answers": {"aws_key": {"type": "text", "value": None}}}, None, id="no-choices"), + ], +) +async def test_post_call_hook_lists_the_guardrail_only_when_it_scanned_text( + respx_mock: respx.MockRouter, response, applied +) -> None: + respx_mock.post(_ANALYZE_URL).respond(json={"finalAction": "NONE", "appliedPolicies": []}) + metadata: 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(), + 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..912021e4bcd 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 @@ -3,6 +3,8 @@ import json import pytest import respx +import litellm + from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.ibm_guardrails.ibm_detector import IBMGuardrailDetector @@ -122,3 +124,86 @@ 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: guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type="completion" + ), + id="pre_call", + ), + pytest.param( + "during_call", + lambda guardrail, data: guardrail.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type="completion" + ), + id="during_call", + ), +] +_REQUEST_TEXT_CASES = [ + pytest.param({"messages": [{"role": "user", "content": "hello there"}]}, ["ibm-guard"], id="messages"), + pytest.param({"state": "hello there"}, None, id="no-messages"), +] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS) +@pytest.mark.parametrize(("request_fields", "applied"), _REQUEST_TEXT_CASES) +async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( + respx_mock: respx.MockRouter, event_hook, run_hook, request_fields, applied +) -> None: + respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[]]) + guardrail = 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: dict[str, object] = {} + + await run_hook(guardrail, {**request_fields, "metadata": metadata}) + + assert metadata.get("applied_guardrails") == applied + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "applied"), + [ + pytest.param( + litellm.ModelResponse( + choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}] + ), + ["ibm-guard"], + id="chat-reply", + ), + pytest.param({"answers": {"aws_key": {"type": "text", "value": None}}}, None, id="no-choices"), + ], +) +async def test_post_call_hook_lists_the_guardrail_only_when_it_scanned_text( + respx_mock: respx.MockRouter, response, applied +) -> None: + respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[]]) + guardrail = 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: dict[str, object] = {} + + await guardrail.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "hello"}], "metadata": metadata}, + user_api_key_dict=UserAPIKeyAuth(), + 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..8c7f0d451e7 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py @@ -316,7 +316,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 +815,34 @@ 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"), + ], +) +async def test_aim_calls_aim_and_lists_itself_only_when_the_request_has_text( + respx_mock: respx.MockRouter, call_type: str, request_fields: dict, aim_calls: int, applied: list[str] | None +): + analyze = respx_mock.post("http://aim.test/fw/v1/analyze").respond( + json={"analysis_result": {"analysis_time_ms": 1, "policy_drill_down": {}}, "required_action": None} + ) + guardrail = AimGuardrail(api_key="hs-aim-key", api_base="http://aim.test", guardrail_name="aim") + metadata: 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..66a8f5a85a8 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py @@ -4,6 +4,7 @@ import ssl from unittest.mock import AsyncMock, MagicMock, patch import pytest +import respx from fastapi.exceptions import HTTPException from httpx import Request, Response from websockets.exceptions import ConnectionClosed @@ -2785,3 +2786,34 @@ 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"), + ], +) +async def test_cato_calls_cato_and_lists_itself_only_when_the_request_has_text( + respx_mock, call_type, request_fields, cato_calls, applied +): + analyze = respx_mock.post("http://cato.test/fw/v1/analyze").respond( + json={"analysis_result": {}, "required_action": None} + ) + guardrail = CatoNetworksGuardrail(api_key="cato-key", api_base="http://cato.test", guardrail_name="cato") + metadata = {} + + 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 660cc42a07d..ae6ee15d64a 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py @@ -1307,3 +1307,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: list[str] | 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_lasso.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py index b8a219ba863..dec3510d6ab 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py @@ -1549,3 +1549,35 @@ 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", "applied"), + [ + pytest.param( + "adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, ["lasso"], id="decisions" + ), + pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, None, id="extraction"), + ], +) +async def test_lasso_lists_itself_only_when_it_classified_request_text( + respx_mock: respx.MockRouter, call_type: str, request_fields: dict[str, object], applied: list[str] | None +) -> None: + 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="pre_call", 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 metadata.get("applied_guardrails") == applied 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..7cc553e2aee 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -5291,6 +5291,134 @@ async def test_logging_only_flagged_prompt_still_scans_response(): assert len(flagged) == 2 +class _RedactsPhoneNumbers(AsyncHTTPHandler): + def __init__(self): + self.scanned: list[str] = [] + + async def post( + self, + url: str, + json: dict[str, object] | None = None, + headers: dict[str, str] | None = None, + **kwargs: object, + ) -> httpx.Response: + payload = cast(dict[str, dict[str, str]], json) + text = (payload.get("userPromptData") or payload["modelResponseData"])["text"] + self.scanned.append(text) + redacted = text.replace("555-123-4567", "[REDACTED]") + result = ( + _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("input_type", ("request", "response")) +async def test_apply_guardrail_masks_each_text_when_masking_is_on(input_type): + handler = _RedactsPhoneNumbers() + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-mask", + event_hook=GuardrailEventHooks.pre_call, + mask_request_content=input_type == "request", + mask_response_content=input_type == "response", + async_handler=handler, + access_token_provider=_async_token_provider, + ) + + result = await guardrail.apply_guardrail( + inputs={"texts": ["call me at 555-123-4567", "refund order 4411"]}, + request_data={"metadata": {}}, + input_type=input_type, + ) + + assert result["texts"] == ["call me at [REDACTED]", "refund order 4411"] + assert sorted(handler.scanned) == ["call me at 555-123-4567", "refund order 4411"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_blocks_a_masking_finding_when_masking_is_off(): + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-block", + event_hook=GuardrailEventHooks.pre_call, + 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_when_masking_is_off(): + inputs = {"texts": ["refund order 4411"]} + guardrail = 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): + 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): + guardrail = 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 = {"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") == 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..1a61dc06131 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py @@ -2,8 +2,11 @@ from unittest.mock import AsyncMock, patch import httpx 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 +371,37 @@ 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, call_type, request_fields, pangea_calls, applied +): + guard = respx_mock.post("https://ai-guard.aws.us.pangea.cloud/v1beta/guard").respond( + json={"result": {"blocked": False, "transformed": False}} + ) + guardrail = PangeaHandler( + mode="pre_call", guardrail_name="pangea-ai-guard", api_key="pts_pangeatokenid", default_on=True + ) + metadata = {} + + 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..49c21fa60f4 100644 --- a/tests/unit/proxy/guardrails/test_pillar_guardrails.py +++ b/tests/unit/proxy/guardrails/test_pillar_guardrails.py @@ -17,6 +17,7 @@ import json from urllib.parse import unquote import pytest +import respx from fastapi.exceptions import HTTPException from httpx import Request, Response @@ -1513,3 +1514,52 @@ 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: guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type="completion" + ), + id="pre_call", + ), + pytest.param( + "during_call", + lambda guardrail, data: guardrail.async_moderation_hook( + data=data, user_api_key_dict=UserAPIKeyAuth(), call_type="completion" + ), + 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", "applied"), + [ + pytest.param({"messages": [{"role": "user", "content": "hi"}]}, ["pillar-test"], id="messages"), + pytest.param({"state": "hi"}, None, id="no-messages"), + ], +) +async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text( + respx_mock: respx.MockRouter, event_hook, run_hook, request_fields, applied +): + respx_mock.post(url__startswith="https://api.pillar.security/api/v1/protect").respond( + json={"session_id": "s", "flagged": False, "scanners": {}, "evidence": []} + ) + guardrail = 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: dict[str, object] = {} + + await run_hook(guardrail, {**request_fields, "metadata": metadata}) + + assert metadata.get("applied_guardrails") == applied diff --git a/tests/unit/proxy/test_proxy_logging_hook_detection.py b/tests/unit/proxy/test_proxy_logging_hook_detection.py index d3d00aa0794..17e481db49a 100644 --- a/tests/unit/proxy/test_proxy_logging_hook_detection.py +++ b/tests/unit/proxy/test_proxy_logging_hook_detection.py @@ -1134,6 +1134,21 @@ async def test_pre_call_hook_scans_decision_and_extraction_text_with_a_native_ho assert _applied_guardrails(result) == ["keeps_native"] +@pytest.mark.asyncio +async def test_pre_call_hook_does_not_list_a_guardrail_for_a_systemone_request_without_text(monkeypatch): + guardrail = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + result = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), + data={"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): guardrail = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.during_call, default_on=True)