mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(guardrails): list guardrails that scanned decision and extraction text, and scan every part the model reads
Decision Model, Aim, Cato Networks, and Lasso add themselves to x-litellm-applied-guardrails when they scan a request, so the decisions routes keep listing them. Aim, Cato Networks, and Pangea no longer call their API when a request has no text for them. The systemone translation scans keys and numbers in state plus criteria descriptions, and decisions scans choice and level descriptions. Model Armor honors its mask flags and fail_on_error in apply_guardrail, and Pillar, DynamoAI, and IBM Guardrails list themselves only when they had text to scan.
This commit is contained in:
mateo-berri
parent
a5e41d91ac
commit
b45f9124bc
23 changed files with 832 additions and 65 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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": "<AWS_KEY>", "AKIAI44QH8DHBEXAMPLE": "<AWS_KEY>"}
|
||||
return GenericGuardrailAPIInputs(
|
||||
texts=[
|
||||
"<CARD>" 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": "<X>", "card": 4111111111111111, "items": 2},
|
||||
"questions": {"q": {"type": "noul", "criteria": {}}},
|
||||
}
|
||||
|
||||
result: Final = await SystemOneHandler().process_input_messages(data, _MasksSecrets(guardrail_name="mask"))
|
||||
|
||||
assert result["state"] == {"<AWS_KEY>": "<X>", "card": "<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
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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: <US_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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue