scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 280

Merge acb1b9e49c into b91fc92b6d

Browse source
This commit is contained in:
tessellark 2026-10-30 03:43:13 +00:00 • committed by GitHub
commit f77da74088
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
42 changed files with 2585 additions and 102 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
enterprise/litellm_enterprise/enterprise_callbacks/secret_detection.py
View file

@ -19,7 +19,7 @@ import json
import re
import tempfile
from collections import OrderedDict
from collections.abc import Generator, Iterator, Sequence
from collections.abc import Generator, Iterator, Mapping, Sequence
from contextvars import ContextVar
from functools import partial
from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional
@ -601,6 +601,24 @@ def _next_run_state(budget: int, state: tuple[int, int, int], size: int) -> tupl
return run_id, count + 1, total + size
_KEY_PERMISSIONS: Final[TypeAdapter[Mapping[object, object] | None]] = TypeAdapter(Mapping[object, object] | None)
_METADATA_BUCKET: Final = TypeAdapter(Mapping[str, object])
def _key_opted_out(key_auth: object) -> bool:
if not isinstance(key_auth, UserAPIKeyAuth):
return False
permissions: Final = _KEY_PERMISSIONS.validate_python(key_auth.permissions)
return permissions is not None and permissions.get(GUARDRAIL_NAME) is False
def _request_key_opted_out(data: Mapping[str, object]) -> bool:
return any(
isinstance(bucket, dict) and _key_opted_out(_METADATA_BUCKET.validate_python(bucket).get("user_api_key_auth"))
for bucket in (data.get("metadata"), data.get("litellm_metadata"))
)
def _value_in_texts(value: str, texts: Sequence[str]) -> str | None:
"""Return the detected value as it appears in ``texts``, decoding it when it came from a JSON-escaped key line."""
candidates: Final = (value, _json_unescaped(value))
@ -609,7 +627,7 @@ def _value_in_texts(value: str, texts: Sequence[str]) -> str | None:
class _ENTERPRISE_SecretDetection(CustomGuardrail):
# Keeps proxied traffic on async_pre_call_hook (the unified apply_guardrail
# path skips should_run_check and never sees data["prompt"]).
# path never sees data["prompt"]).
use_native_lifecycle_hooks: ClassVar[bool] = True
def __init__(self, detect_secrets_config: Optional[dict] = None, **kwargs):
@ -694,12 +712,12 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
return map_tool_json_text(payload, lambda text: pattern.sub("[REDACTED]", text))
async def should_run_check(self, user_api_key_dict: UserAPIKeyAuth) -> bool:
if user_api_key_dict.permissions is not None:
if GUARDRAIL_NAME in user_api_key_dict.permissions:
if user_api_key_dict.permissions[GUARDRAIL_NAME] is False:
return False
return not _key_opted_out(user_api_key_dict)
return True
def should_run_guardrail(self, data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool:
return super().should_run_guardrail( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.should_run_guardrail leaves data untyped
data, event_type
) and not _request_key_opted_out(data)
@log_guardrail_information
async def apply_guardrail(
@ -710,8 +728,9 @@ class _ENTERPRISE_SecretDetection(CustomGuardrail):
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> GenericGuardrailAPIInputs:
"""Unified-interface entrypoint, used by /guardrails/apply_guardrail
(the UI test playground). Proxied traffic keeps using
``async_pre_call_hook``, see ``use_native_lifecycle_hooks``."""
(the UI test playground) and by the decision and extraction routes,
whose request text the native hook never reads. Other proxied traffic
keeps using ``async_pre_call_hook``, see ``use_native_lifecycle_hooks``."""
texts = inputs.get("texts")
if not texts or not any(texts):
return inputs

26
litellm/integrations/custom_guardrail.py
View file

@ -14,6 +14,7 @@ from litellm._internal_context import with_service_target
from litellm._logging import verbose_logger
from litellm.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
get_or_create_metadata_bucket,
@ -72,6 +73,30 @@ _PRE_CALL_EXECUTED_TOKEN: Final = secrets.token_hex(16)
_GUARDRAIL_BLOCK_STATUS_CODES: Final = frozenset({400, 403, 422})
APPLY_GUARDRAIL_ONLY_CALL_TYPES: Final[frozenset[str]] = frozenset(
call_type.value
for call_type in (
CallTypes.extraction,
CallTypes.aextraction,
CallTypes.decisions,
CallTypes.adecisions,
CallTypes.systemone,
CallTypes.asystemone,
)
)
def is_apply_guardrail_only_route(route: str | None) -> bool:
call_types: Final = get_call_types_for_route(route) if route else None
return bool(call_types) and call_types[0] in APPLY_GUARDRAIL_ONLY_CALL_TYPES
def runs_through_apply_guardrail(guardrail: "CustomGuardrail", call_type: str | None) -> bool:
return "apply_guardrail" in type(guardrail).__dict__ and (
not guardrail.use_native_lifecycle_hooks or call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES
)
DEFAULT_ADVISORY_MESSAGE: Final = (
"The user's latest message was flagged for {reason} by a content safety "
"guardrail. This may be a false positive. Use your judgment: respond "
@ -226,6 +251,7 @@ class CustomGuardrail(CustomLogger):
use_native_lifecycle_hooks: ClassVar[bool] = False
records_own_guardrail_information: ClassVar[bool] = False
adds_itself_to_applied_guardrails: ClassVar[bool] = False
logging_only_scope: LoggingOnlyScope | None
logging_only_continue_on_input_failure: bool

4
litellm/llms/base_llm/guardrail_translation/base_translation.py
View file

@ -72,6 +72,10 @@ class BaseTranslation(ABC):
has the legacy post-call hook against that response, so on a translation without it such
a guardrail keeps running on its own."""
scans_responses: ClassVar[bool] = True
lists_guardrail_when_called: ClassVar[bool] = False
def post_call_hook_response(self, response: object) -> object:
"""The ``response`` this endpoint's non-streaming post-call hooks receive, derived from
the object the translation stores under ``request_data["response"]`` while scanning an

13
litellm/llms/openai/decisions/guardrail_translation/__init__.py Normal file
View file

@ -0,0 +1,13 @@
from typing import Final
from litellm.llms.openai.decisions.guardrail_translation.handler import DecisionsHandler, SystemOneHandler
from litellm.types.utils import CallTypes
guardrail_translation_mappings: Final = {
CallTypes.decisions: DecisionsHandler,
CallTypes.adecisions: DecisionsHandler,
CallTypes.systemone: SystemOneHandler,
CallTypes.asystemone: SystemOneHandler,
}
__all__ = ["DecisionsHandler", "SystemOneHandler", "guardrail_translation_mappings"]

412
litellm/llms/openai/decisions/guardrail_translation/handler.py Normal file
View file

@ -0,0 +1,412 @@
import json
import re
from collections.abc import Iterator, Mapping
from difflib import SequenceMatcher
from itertools import chain, groupby
from types import MappingProxyType
from typing import TYPE_CHECKING, ClassVar, Final
from pydantic import JsonValue, TypeAdapter, ValidationError
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.llms.base_llm.guardrail_translation.utils import unappliable_request_rewrite
from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler
from litellm.types.utils import GenericGuardrailAPIInputs
if TYPE_CHECKING:
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.llms.openai import AllMessageValues
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
_KEY_ADAPTER: Final[TypeAdapter[str]] = TypeAdapter(str)
_JSON_KEY_PREFIX: Final = re.compile(r'("(?:[^"\\\x00-\x1f]|\\(?:["\\/bfnrt]|u[0-9a-fA-F]{4}))*"): ')
_JSON_ESCAPE: Final = re.compile(r"(\\(?:[bfnrt]|u[0-9a-fA-F]{4}))|\\.")
_MARKED_JSON_ESCAPE: Final = re.compile(r"(\\(?:[bfnrt]|u[0-9a-fA-F]{4}))~|\\.")
_TOKEN: Final = re.compile(r"\s+|\S+")
_DESCRIBED_ITEM_KEYS: Final = ("choices", "levels")
_QUESTION_LIST_ADAPTER: Final[TypeAdapter[tuple[Mapping[str, JsonValue], ...]]] = TypeAdapter(
tuple[Mapping[str, JsonValue], ...]
)
_QUESTION_MAP_ADAPTER: Final[TypeAdapter[Mapping[str, Mapping[str, JsonValue]]]] = TypeAdapter(
Mapping[str, Mapping[str, JsonValue]]
)
def _input_text(part: JsonValue) -> str | None:
if not isinstance(part, dict) or part.get("type") != "input_text":
return None
text: Final = part.get("text")
return text if isinstance(text, str) else None
def _content_texts(content: JsonValue) -> tuple[str, ...]:
if isinstance(content, list):
return tuple(text for part in content if (text := _input_text(part)))
return (content,) if isinstance(content, str) and content else ()
def _message_texts(message: JsonValue) -> tuple[str, ...]:
return _content_texts(message.get("content")) if isinstance(message, dict) else ()
def _input_texts(decision_input: JsonValue) -> tuple[str, ...]:
if not isinstance(decision_input, list):
return _content_texts(decision_input)
return tuple(chain.from_iterable(_message_texts(message) for message in decision_input))
def _guardrailed_part(part: JsonValue, guardrailed: Mapping[str, str]) -> JsonValue:
text: Final = _input_text(part)
if text is None or not isinstance(part, dict):
return part
return {**part, "text": guardrailed.get(text, text)}
def _guardrailed_content(content: JsonValue, guardrailed: Mapping[str, str]) -> JsonValue:
if isinstance(content, list):
return [_guardrailed_part(part, guardrailed) for part in content]
return guardrailed.get(content, content) if isinstance(content, str) else content
def _guardrailed_message(message: JsonValue, guardrailed: Mapping[str, str]) -> JsonValue:
return (
{**message, "content": _guardrailed_content(message["content"], guardrailed)}
if isinstance(message, dict) and "content" in message
else message
)
def _guardrailed_input(decision_input: JsonValue, guardrailed: Mapping[str, str]) -> JsonValue:
if not isinstance(decision_input, list):
return _guardrailed_content(decision_input, guardrailed)
return [_guardrailed_message(message, guardrailed) for message in decision_input]
def _marked_escape(escape: re.Match[str]) -> str:
return f"{escape.group()}~" if escape.group(1) else escape.group()
def _unmarked_escape(escape: re.Match[str]) -> str:
return escape.group(1) or escape.group()
def _entry_text(key: str, value: str | float | None) -> str:
return _JSON_ESCAPE.sub(_marked_escape, json.dumps({key: value}, ensure_ascii=False)[1:-1])
def _escaped_texts(text: str, key: str, item: JsonValue) -> tuple[str, ...]:
return tuple(form for form in (key, item) if isinstance(form, str) and form not in text)
def _entry_text_forms(key: str, item: JsonValue) -> Iterator[str]:
if isinstance(item, (list, dict)):
yield key
yield from _json_text_forms(item)
else:
text: Final = _entry_text(key, item)
yield text
yield from _escaped_texts(text, key, item)
def _json_text_forms(value: JsonValue) -> Iterator[str]:
match value:
case bool() | None:
return
case str():
yield value
case int() | float():
yield json.dumps(value)
case list():
yield from chain.from_iterable(map(_json_text_forms, value))
case dict():
for key, item in value.items():
yield from _entry_text_forms(key, item)
def _json_texts(value: JsonValue) -> tuple[str, ...]:
return tuple(text for text in _json_text_forms(value) if text)
def _guardrailed_number(value: float, guardrailed: Mapping[str, str]) -> JsonValue:
text: Final = json.dumps(value)
rewritten: Final = guardrailed.get(text, text)
return value if rewritten == text else rewritten
def _occurrence_numbers(keys: tuple[str, ...]) -> Iterator[tuple[int, int]]:
for _, positions in groupby(sorted(range(len(keys)), key=keys.__getitem__), key=keys.__getitem__):
yield from ((position, number) for number, position in enumerate(positions, start=1))
def _distinct_keys(keys: tuple[str, ...]) -> tuple[str, ...]:
if len(frozenset(keys)) == len(keys):
return keys
numbers: Final = MappingProxyType(dict(_occurrence_numbers(keys)))
return tuple(key if numbers[position] == 1 else f"{key}_{numbers[position]}" for position, key in enumerate(keys))
def _guardrailed_json(value: JsonValue, guardrailed: Mapping[str, str], guardrail_name: str | None) -> JsonValue:
match value:
case bool() | None:
return value
case str():
return guardrailed.get(value, value)
case int() | float():
return _guardrailed_number(value, guardrailed)
case list():
return [_guardrailed_json(item, guardrailed, guardrail_name) for item in value]
case dict():
entries: Final = tuple(
_guardrailed_entry(key, item, guardrailed, guardrail_name) for key, item in value.items()
)
keys: Final = _distinct_keys(tuple(key for key, _ in entries))
if len(frozenset(keys)) != len(keys):
raise unappliable_request_rewrite(guardrail_name)
return dict(zip(keys, (item for _, item in entries), strict=True))
def _rewritten_value(text: str) -> JsonValue:
try:
return _JSON_ADAPTER.validate_json(text)
except ValidationError:
return text
def _rewritten_entry(rewritten: str, guardrail_name: str | None) -> tuple[str, JsonValue]:
unmarked: Final = _MARKED_JSON_ESCAPE.sub(_unmarked_escape, rewritten)
key: Final = _JSON_KEY_PREFIX.match(unmarked)
if key is None:
raise unappliable_request_rewrite(guardrail_name)
return _KEY_ADAPTER.validate_json(key.group(1)), _rewritten_value(unmarked[key.end() :])
def _decoded_text(text: str, form: str, guardrailed: Mapping[str, str]) -> str:
return form if form in text else guardrailed.get(form, form)
def _tokens(text: str) -> tuple[str, ...]:
return tuple(_TOKEN.findall(text))
def _edits(original: tuple[str, ...], rewritten: tuple[str, ...]) -> frozenset[tuple[int, int, str]]:
opcodes: Final = SequenceMatcher(None, original, rewritten, autojunk=False).get_opcodes()
return frozenset(
(start, end, "".join(rewritten[low:high])) for tag, start, end, low, high in opcodes if tag != "equal"
)
def _covered(edit: tuple[int, int, str], edits: frozenset[tuple[int, int, str]]) -> bool:
start, end, _ = edit
return any(other[0] <= start and end <= other[1] and other[1] - other[0] > end - start for other in edits)
def _merged_text(original: str, rewritten: str, decoded: str, guardrail_name: str | None) -> str:
if decoded in (original, rewritten):
return rewritten
if rewritten == original:
return decoded
tokens: Final = _tokens(original)
edits: Final = _edits(tokens, _tokens(rewritten)) | _edits(tokens, _tokens(decoded))
kept: Final = tuple(sorted(edit for edit in edits if not _covered(edit, edits)))
if any(later[0] < earlier[1] for earlier, later in zip(kept, kept[1:])):
raise unappliable_request_rewrite(guardrail_name)
kept_starts: Final = (0, *(end for _, end, _ in kept))
kept_ends: Final = (*(start for start, _, _ in kept), len(tokens))
replacements: Final = (*(text for _, _, text in kept), "")
return "".join(
"".join(tokens[start:end]) + text for start, end, text in zip(kept_starts, kept_ends, replacements, strict=True)
)
def _merged_value(original: str, rewritten: JsonValue, decoded: str, guardrail_name: str | None) -> JsonValue:
if decoded == original:
return rewritten
if isinstance(rewritten, str):
return _merged_text(original, rewritten, decoded, guardrail_name)
raise unappliable_request_rewrite(guardrail_name)
def _guardrailed_entry(
key: str, item: JsonValue, guardrailed: Mapping[str, str], guardrail_name: str | None
) -> tuple[str, JsonValue]:
if isinstance(item, (list, dict)):
return guardrailed.get(key, key), _guardrailed_json(item, guardrailed, guardrail_name)
text: Final = _entry_text(key, item)
rewritten: Final = guardrailed.get(text, text)
entry: Final = (key, item) if rewritten == text else _rewritten_entry(rewritten, guardrail_name)
merged_key: Final = _merged_text(key, entry[0], _decoded_text(text, key, guardrailed), guardrail_name)
if not isinstance(item, str):
return merged_key, entry[1]
return merged_key, _merged_value(item, entry[1], _decoded_text(text, item, guardrailed), guardrail_name)
def _instruction_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]:
return _json_texts(question.get("instructions"))
def _described_items(question: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]:
return tuple(chain.from_iterable(_list_items(question.get(key)) for key in _DESCRIBED_ITEM_KEYS))
def _list_items(value: JsonValue) -> tuple[JsonValue, ...]:
return tuple(value) if isinstance(value, list) else ()
def _item_description(item: JsonValue) -> tuple[str, ...]:
description: Final = item.get("description") if isinstance(item, dict) else None
return (description,) if isinstance(description, str) and description else ()
def _decision_question_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]:
descriptions: Final = chain.from_iterable(map(_item_description, _described_items(question)))
return (*_instruction_texts(question), *descriptions)
def _guardrailed_item(item: JsonValue, guardrailed: Mapping[str, str]) -> JsonValue:
if not isinstance(item, dict) or not isinstance(description := item.get("description"), str):
return item
return {**item, "description": guardrailed.get(description, description)}
def _with_guardrailed_decision_question(
question: Mapping[str, JsonValue], guardrailed: Mapping[str, str], guardrail_name: str | None
) -> Mapping[str, JsonValue]:
described: Final = {
key: [_guardrailed_item(item, guardrailed) for item in _list_items(question[key])]
for key in _DESCRIBED_ITEM_KEYS
if key in question
}
return {**_with_guardrailed_instructions(question, guardrailed, guardrail_name), **described}
def _criteria_texts(criteria: JsonValue) -> tuple[str, ...]:
descriptions: Final = tuple(criteria.values()) if isinstance(criteria, dict) else (criteria,)
return tuple(chain.from_iterable(map(_json_texts, descriptions)))
def _systemone_question_texts(question: Mapping[str, JsonValue]) -> tuple[str, ...]:
return (*_instruction_texts(question), *_criteria_texts(question.get("criteria")))
def _guardrailed_criteria(criteria: JsonValue, guardrailed: Mapping[str, str], guardrail_name: str | None) -> JsonValue:
if not isinstance(criteria, dict):
return _guardrailed_json(criteria, guardrailed, guardrail_name)
return {
label: _guardrailed_json(description, guardrailed, guardrail_name) for label, description in criteria.items()
}
def _with_guardrailed_systemone_question(
question: Mapping[str, JsonValue], guardrailed: Mapping[str, str], guardrail_name: str | None
) -> Mapping[str, JsonValue]:
with_instructions: Final = _with_guardrailed_instructions(question, guardrailed, guardrail_name)
if "criteria" not in question:
return with_instructions
return {**with_instructions, "criteria": _guardrailed_criteria(question["criteria"], guardrailed, guardrail_name)}
def _with_guardrailed_instructions(
question: Mapping[str, JsonValue], guardrailed: Mapping[str, str], guardrail_name: str | None
) -> Mapping[str, JsonValue]:
if "instructions" not in question:
return question
return {**question, "instructions": _guardrailed_json(question["instructions"], guardrailed, guardrail_name)}
async def _guardrailed_texts(
texts: tuple[str, ...],
guardrail_to_apply: "CustomGuardrail",
request_data: dict[str, object], # mutable-ok: CustomGuardrail.apply_guardrail takes the request as a plain dict
litellm_logging_obj: "LiteLLMLoggingObj | None",
) -> Mapping[str, str]:
unique_texts: Final = tuple(dict.fromkeys(texts))
model: Final = request_data.get("model")
result: Final = await guardrail_to_apply.apply_guardrail( # pyright: ignore[reportUnknownMemberType] # CustomGuardrail.apply_guardrail types request_data as a bare dict
inputs=GenericGuardrailAPIInputs(texts=list(unique_texts), model=model if isinstance(model, str) else None),
request_data=request_data,
input_type="request",
logging_obj=litellm_logging_obj,
)
if not (
guardrail_to_apply.records_own_guardrail_information or guardrail_to_apply.adds_itself_to_applied_guardrails
):
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict
)
add_guardrail_to_applied_guardrails_header(
request_data=request_data, guardrail_name=guardrail_to_apply.guardrail_name
)
returned: Final = result.get("texts", [])
return MappingProxyType(
{text: returned[index] if index < len(returned) else text for index, text in enumerate(unique_texts)}
)
class _DecisionAnswersTranslation(BaseTranslation):
scans_responses: ClassVar[bool] = False
lists_guardrail_when_called: ClassVar[bool] = True
async def process_output_response(
self,
response: object,
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
user_api_key_dict: "UserAPIKeyAuth | None" = None,
request_data: dict[str, object] | None = None, # mutable-ok: BaseTranslation hands over the request dict
) -> object:
return response
class DecisionsHandler(_DecisionAnswersTranslation):
def get_structured_messages(
self,
data: dict[str, object], # mutable-ok: BaseTranslation hands over the request dict
) -> "list[AllMessageValues] | None": # mutable-ok: overrides BaseTranslation.get_structured_messages
return OpenAIResponsesHandler().get_structured_messages({"input": data.get("input")})
async def process_input_messages(
self,
data: dict[str, object], # mutable-ok: BaseTranslation hands over the proxy's request dict
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
) -> Mapping[str, object]:
decision_input: Final = _JSON_ADAPTER.validate_python(data.get("input"))
questions: Final = _QUESTION_LIST_ADAPTER.validate_python(data.get("questions") or ())
texts: Final = (*_input_texts(decision_input), *chain.from_iterable(map(_decision_question_texts, questions)))
if not texts:
return data
guardrailed: Final = await _guardrailed_texts(texts, guardrail_to_apply, data, litellm_logging_obj)
return {
**data,
"input": _guardrailed_input(decision_input, guardrailed),
"questions": [
_with_guardrailed_decision_question(question, guardrailed, guardrail_to_apply.guardrail_name)
for question in questions
],
}
class SystemOneHandler(_DecisionAnswersTranslation):
async def process_input_messages(
self,
data: dict[str, object], # mutable-ok: BaseTranslation hands over the proxy's request dict
guardrail_to_apply: "CustomGuardrail",
litellm_logging_obj: "LiteLLMLoggingObj | None" = None,
) -> Mapping[str, object]:
state: Final = _JSON_ADAPTER.validate_python(data.get("state"))
questions: Final = _QUESTION_MAP_ADAPTER.validate_python(data.get("questions") or {})
texts: Final = (*_json_texts(state), *chain.from_iterable(map(_systemone_question_texts, questions.values())))
if not texts:
return data
guardrailed: Final = await _guardrailed_texts(texts, guardrail_to_apply, data, litellm_logging_obj)
return {
**data,
"state": _guardrailed_json(state, guardrailed, guardrail_to_apply.guardrail_name),
"questions": {
name: _with_guardrailed_systemone_question(question, guardrailed, guardrail_to_apply.guardrail_name)
for name, question in questions.items()
},
}

34
litellm/proxy/common_utils/callback_utils.py
View file

@ -1,9 +1,9 @@
import copy
import json
import os
from collections.abc import Callable, Iterable, Mapping
from collections.abc import Callable, Iterable, Mapping, Sequence
from dataclasses import dataclass
from itertools import accumulate
from itertools import accumulate, chain
from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, Optional, TypeAlias
from pydantic import TypeAdapter
@ -23,6 +23,7 @@ from litellm.constants import (
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import (
as_str_mapping,
get_metadata_variable_name_from_kwargs,
get_or_create_metadata_bucket,
)
@ -478,16 +479,33 @@ def _serialize_scan_metadata_header(entries: Iterable[object], *, max_length: in
return f"[{','.join(encoded[:kept])}]"
_REQUEST_BODY: Final = TypeAdapter(Mapping[str, object])
_LIST_ENTRIES: Final = TypeAdapter(tuple[object, ...])
def _list_entries(value: object) -> tuple[object, ...]:
return _LIST_ENTRIES.validate_python(value) if isinstance(value, (list, tuple)) else ()
def _bucket_list_entries(buckets: Sequence[Mapping[str, object]], key: str) -> tuple[object, ...]:
values: Final = tuple(bucket.get(key) for bucket in buckets)
distinct_values: Final = {id(value): value for value in values}
entries: Final = tuple(chain.from_iterable(_list_entries(value) for value in distinct_values.values()))
return entries if key == GUARDRAIL_SCAN_METADATA_METADATA_KEY else tuple(dict.fromkeys(entries))
def get_logging_caching_headers(request_data: dict) -> dict | None:
_metadata: Final[dict] = {}
metadata_bucket: Final = request_data.get("metadata")
litellm_metadata_bucket: Final = request_data.get("litellm_metadata")
if isinstance(metadata_bucket, dict):
_metadata.update(metadata_bucket)
if isinstance(litellm_metadata_bucket, dict):
request: Final = _REQUEST_BODY.validate_python(request_data)
found_buckets: Final = (as_str_mapping(request.get(name)) for name in ("metadata", "litellm_metadata"))
buckets: Final = tuple(bucket for bucket in found_buckets if bucket is not None)
for bucket in buckets:
# Batch/file routes store proxy tracking in litellm_metadata while
# user-facing metadata stays in metadata; merge both for headers.
_metadata.update(litellm_metadata_bucket)
_metadata.update(bucket)
for list_key in ("applied_guardrails", GUARDRAIL_SCAN_IDS_METADATA_KEY, GUARDRAIL_SCAN_METADATA_METADATA_KEY):
if any(list_key in bucket for bucket in buckets):
_metadata[list_key] = _bucket_list_entries(buckets, list_key)
headers: Final = {}
if "applied_guardrails" in _metadata:
headers["x-litellm-applied-guardrails"] = ",".join(_metadata["applied_guardrails"])

24
litellm/proxy/guardrails/guardrail_hooks/aim/aim.py
View file

@ -17,12 +17,15 @@ from websockets.asyncio.client import ClientConnection, connect
from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm._version import version as litellm_version
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_guardrail import APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict
)
from litellm.proxy.guardrails._content_utils import (
apply_redacted_messages_back,
build_inspection_messages,
@ -152,7 +155,9 @@ class AimGuardrail(CustomGuardrail):
if is_non_conversational_call_type(call_type) and not self.inspect_embeddings:
verbose_proxy_logger.debug("Aim: skipping non-conversational call type %s", call_type)
return data
return await self.call_aim_guardrail(data, hook="pre_call", key_alias=user_api_key_dict.key_alias)
return await self.call_aim_guardrail(
data, hook="pre_call", key_alias=user_api_key_dict.key_alias, call_type=call_type
)
async def async_moderation_hook(
self,
@ -165,10 +170,19 @@ class AimGuardrail(CustomGuardrail):
verbose_proxy_logger.debug("Aim: skipping non-conversational call type %s", call_type)
return data
await self.call_aim_guardrail(data, hook="moderation", key_alias=user_api_key_dict.key_alias)
await self.call_aim_guardrail(
data, hook="moderation", key_alias=user_api_key_dict.key_alias, call_type=call_type
)
return data
async def call_aim_guardrail(self, data: dict, hook: str, key_alias: str | None) -> dict:
async def call_aim_guardrail(
self, data: dict, hook: str, key_alias: str | None, call_type: CallTypesLiteral | None = None
) -> dict:
messages: Final = self._build_aim_inspection_messages(data) # pyright: ignore[reportUnknownMemberType] # hook request data is an untyped dict
if not messages:
return data # pyright: ignore[reportUnknownVariableType] # hook request data is an untyped dict
if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
user_email: Final = data.get("metadata", {}).get("headers", {}).get("x-aim-user-email")
call_id: Final = data.get("litellm_call_id")
headers: Final = self._build_aim_headers(
@ -180,7 +194,7 @@ class AimGuardrail(CustomGuardrail):
response: Final = await self.async_handler.post(
f"{self.api_base}/fw/v1/analyze",
headers=headers,
json={"messages": self._build_aim_inspection_messages(data)},
json={"messages": messages},
timeout=self.timeout,
)
response.raise_for_status()

15
litellm/proxy/guardrails/guardrail_hooks/cato_networks/cato_networks.py
View file

@ -22,7 +22,7 @@ from websockets.exceptions import ConnectionClosed
from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm._version import version as litellm_version
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_guardrail import APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail
from litellm.litellm_core_utils.core_helpers import iter_items
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
@ -30,6 +30,9 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict
)
from litellm.proxy.guardrails._content_utils import (
apply_redacted_messages_back,
build_inspection_messages,
@ -173,6 +176,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
hook="pre_call",
key_alias=user_api_key_dict.key_alias,
user_email=self._resolve_cato_user_email(user_api_key_dict),
call_type=call_type,
)
async def async_moderation_hook(
@ -190,6 +194,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
hook="moderation",
key_alias=user_api_key_dict.key_alias,
user_email=self._resolve_cato_user_email(user_api_key_dict),
call_type=call_type,
)
@classmethod
@ -294,7 +299,13 @@ class CatoNetworksGuardrail(CustomGuardrail):
hook: str,
key_alias: str | None,
user_email: str | None = None,
call_type: CallTypesLiteral | None = None,
) -> dict:
messages: Final = self._inspection_messages(data) # pyright: ignore[reportUnknownVariableType,reportUnknownMemberType] # hook request data is an untyped dict
if not messages:
return data # pyright: ignore[reportUnknownVariableType] # hook request data is an untyped dict
if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
call_id: Final = data.get("litellm_call_id")
headers: Final = self._build_cato_headers(
hook=hook,
@ -305,7 +316,7 @@ class CatoNetworksGuardrail(CustomGuardrail):
response: Final = await self.async_handler.post(
f"{self.api_base}/fw/v1/analyze",
headers=headers,
json={"messages": self._inspection_messages(data)},
json={"messages": messages},
timeout=self.timeout,
)
response.raise_for_status()

4
litellm/proxy/guardrails/guardrail_hooks/decision_model/decision_model.py
View file

@ -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,

20
litellm/proxy/guardrails/guardrail_hooks/dynamoai/dynamoai.py
View file

@ -15,7 +15,11 @@ import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_guardrail import (
APPLY_GUARDRAIL_ONLY_CALL_TYPES,
CustomGuardrail,
is_apply_guardrail_only_route,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -328,8 +332,8 @@ class DynamoAIGuardrails(CustomGuardrail):
error_message: Final = self._create_error_message(processed_result)
raise ValueError(error_message)
# Add guardrail to applied guardrails header
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
if _messages or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
return data
@ -373,8 +377,8 @@ class DynamoAIGuardrails(CustomGuardrail):
error_message: Final = self._create_error_message(processed_result)
raise ValueError(error_message)
# Add guardrail to applied guardrails header
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
if _messages or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
return data
@ -441,8 +445,10 @@ class DynamoAIGuardrails(CustomGuardrail):
error_message: Final = self._create_error_message(processed_result)
raise ValueError(error_message)
# Add guardrail to applied guardrails header
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
if isinstance(response, litellm.ModelResponse) or not is_apply_guardrail_only_route(
user_api_key_dict.request_route
):
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
async def async_post_call_streaming_iterator_hook(
self,

20
litellm/proxy/guardrails/guardrail_hooks/ibm_guardrails/ibm_detector.py
View file

@ -15,7 +15,11 @@ import httpx
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_guardrail import (
APPLY_GUARDRAIL_ONLY_CALL_TYPES,
CustomGuardrail,
is_apply_guardrail_only_route,
)
from litellm.llms.custom_httpx.http_handler import (
get_async_httpx_client,
httpxSpecialProvider,
@ -476,8 +480,8 @@ class IBMGuardrailDetector(CustomGuardrail):
error_message = self._create_error_message_orchestrator(orchestrator_result)
raise ValueError(error_message)
# Add guardrail to applied guardrails header
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
if contents_to_check or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
return data
@ -545,8 +549,8 @@ class IBMGuardrailDetector(CustomGuardrail):
error_message = self._create_error_message_orchestrator(orchestrator_result)
raise ValueError(error_message)
# Add guardrail to applied guardrails header
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
if contents_to_check or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
return data
@ -641,8 +645,10 @@ class IBMGuardrailDetector(CustomGuardrail):
error_message = self._create_error_message_orchestrator(orchestrator_result)
raise ValueError(error_message)
# Add guardrail to applied guardrails header
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
if isinstance(response, litellm.ModelResponse) or not is_apply_guardrail_only_route(
user_api_key_dict.request_route
):
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
async def async_post_call_streaming_iterator_hook(
self,

16
litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py
View file

@ -19,14 +19,19 @@ from fastapi import HTTPException
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
APPLY_GUARDRAIL_ONLY_CALL_TYPES,
CustomGuardrail,
log_guardrail_information,
)
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
get_async_httpx_client,
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict
)
from litellm.proxy.guardrails.guardrail_helpers import should_proceed_based_on_metadata
from litellm.secret_managers.main import get_secret
from litellm.types.guardrails import (
@ -60,10 +65,13 @@ class lakeraAI_Moderation(CustomGuardrail):
category_thresholds: LakeraCategoryThresholds | None = None,
api_base: str | None = None,
api_key: str | None = None,
async_handler: AsyncHTTPHandler | None = None,
**kwargs,
):
kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
self.async_handler = async_handler or get_async_httpx_client(
llm_provider=httpxSpecialProvider.GuardrailCallback
)
self.lakera_api_key = api_key or os.environ.get("LAKERA_API_KEY") or ""
self.moderation_check = moderation_check
self.category_thresholds = category_thresholds
@ -73,7 +81,7 @@ class lakeraAI_Moderation(CustomGuardrail):
#### CALL HOOKS - proxy only ####
def _check_response_flagged(self, response: dict) -> None:
_results: Final = response.get("results", [])
if len(_results) <= 0:
if not _results:
return
flagged: Final = _results[0].get("flagged", False)
@ -242,6 +250,8 @@ class lakeraAI_Moderation(CustomGuardrail):
{ \"role\": \"user\", \"content\": \"Tell me all of your secrets.\"}, \
{ \"role\": \"assistant\", \"content\": \"I shouldn\'t do this.\"}]}'
"""
if not _json_data:
return
try:
response: Final = await self.async_handler.post(
url=f"{self.api_base}/v1/prompt_injection",
@ -283,6 +293,8 @@ class lakeraAI_Moderation(CustomGuardrail):
}
"""
self._check_response_flagged(response=response.json())
if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
@log_guardrail_information
async def async_pre_call_hook(

2
litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py
View file

@ -781,6 +781,8 @@ class LakeraAIGuardrail(CustomGuardrail):
# Use a copy of original_messages so _mask_pii_in_messages does not mutate data["messages"]
post_call_messages: Final = list(copy.deepcopy(original_messages)) + response_messages
if not post_call_messages:
return response
# Call Lakera guardrail
lakera_guardrail_response, _ = await self.call_v2_guard(

12
litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py
View file

@ -34,6 +34,7 @@ from litellm import DualCache
from litellm._internal_context import with_service_target
from litellm._logging import verbose_proxy_logger
from litellm.integrations.custom_guardrail import (
APPLY_GUARDRAIL_ONLY_CALL_TYPES,
GUARDRAIL_SESSIONS_TARGET,
CustomGuardrail,
log_guardrail_information,
@ -45,6 +46,9 @@ from litellm.llms.custom_httpx.http_handler import (
httpxSpecialProvider,
)
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header, # pyright: ignore[reportUnknownVariableType] # callback_utils types request_data as a bare dict
)
from litellm.proxy.guardrails._content_utils import (
build_inspection_messages,
has_non_string_content,
@ -204,7 +208,7 @@ class LassoGuardrail(CustomGuardrail):
# The conversation_id is being stored in the cache so it can be used by the post_call hook
self._get_or_generate_conversation_id(data, global_cache)
return await self._run_lasso_guardrail(data, global_cache, message_type="PROMPT")
return await self._run_lasso_guardrail(data, global_cache, message_type="PROMPT", call_type=call_type)
@log_guardrail_information
async def async_moderation_hook(
@ -232,7 +236,7 @@ class LassoGuardrail(CustomGuardrail):
if self.should_run_guardrail(data=data, event_type=event_type) is not True:
return data
return await self._run_lasso_guardrail(data, cache, message_type="PROMPT")
return await self._run_lasso_guardrail(data, cache, message_type="PROMPT", call_type=call_type)
@log_guardrail_information
async def async_post_call_success_hook(
@ -372,6 +376,7 @@ class LassoGuardrail(CustomGuardrail):
data: dict,
cache: DualCache,
message_type: Literal["PROMPT", "COMPLETION"] = "PROMPT",
call_type: str | None = None,
):
"""
Run the Lasso guardrail with the specified message type.
@ -392,6 +397,7 @@ class LassoGuardrail(CustomGuardrail):
data: The request data containing messages
cache: The cache instance for storing conversation_id (optional for post-call)
message_type: Either "PROMPT" for input or "COMPLETION" for output
call_type: The request's call type, which decides whether Lasso lists itself as applied
Raises:
LassoGuardrailAPIError: If the Lasso API call fails
@ -409,6 +415,8 @@ class LassoGuardrail(CustomGuardrail):
messages: Final[Sequence[Mapping[str, object]]] = [*expanded_messages, *input_messages]
if not messages:
return data
if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
# Lasso's classifix endpoint returns masked text that we copy back
# into ``data["messages"]``. For multimodal/Responses-API input we

9
litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
View file

@ -185,6 +185,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
"""
use_native_lifecycle_hooks: ClassVar[bool] = True
adds_itself_to_applied_guardrails: ClassVar[bool] = True
@classmethod
def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
@ -291,7 +292,7 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
return _redact_scanned_content(armor_response)
return armor_response
def _raise_if_fail_closed(self, e: ModelArmorAPIError) -> None:
def _raise_if_fail_closed(self, e: ModelArmorAPIError | httpx.HTTPError) -> None:
if self.optional_params.get("fail_on_error", True):
raise e from None
@ -1325,8 +1326,14 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
end_time=error_end_time,
duration=error_end_time - start_time,
)
self._raise_if_fail_closed(e)
return inputs
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
)
add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name)
flagged: Final = self._should_block_content(armor_response, allow_sanitization=False)
end_time: Final = time.time()
self.add_standard_logging_guardrail_information_to_request_data(

2
litellm/proxy/guardrails/guardrail_hooks/pangea/pangea.py
View file

@ -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,

9
litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py
View file

@ -22,6 +22,7 @@ from litellm import DualCache
from litellm._logging import verbose_proxy_logger
from litellm._version import version as litellm_version
from litellm.integrations.custom_guardrail import (
APPLY_GUARDRAIL_ONLY_CALL_TYPES,
CustomGuardrail,
log_guardrail_information,
)
@ -351,8 +352,8 @@ class PillarGuardrail(CustomGuardrail):
verbose_proxy_logger.debug("Pillar Guardrail: Pre-call hook")
result: Final = await self.run_pillar_guardrail(data, user_api_key_dict)
# Add guardrail name to response headers
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
if data.get("messages") or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
return result
@ -394,8 +395,8 @@ class PillarGuardrail(CustomGuardrail):
verbose_proxy_logger.debug("Pillar Guardrail: During-call moderation hook")
result: Final = await self.run_pillar_guardrail(data, user_api_key_dict)
# Add guardrail name to response headers
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
if data.get("messages") or call_type not in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
return result

54
litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py
View file

@ -20,10 +20,10 @@ from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.cost_calculator import infer_call_type
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_guardrail import APPLY_GUARDRAIL_ONLY_CALL_TYPES, CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_or_create_metadata_bucket
from litellm.llms import get_guardrail_translation_mapping, load_guardrail_translation_mappings
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.guardrails import GuardrailEventHooks
@ -222,16 +222,25 @@ def _a2a_jsonrpc_error_chunk(exc: HTTPException, request_id: str | None) -> Mapp
}
def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> None:
"""Populate data['litellm_metadata'] from user_api_key_dict if absent."""
if "litellm_metadata" not in data:
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
)
def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: str | None = None) -> None:
"""Populate data['litellm_metadata'] from user_api_key_dict if absent.
user_metadata: Final = BaseTranslation.transform_user_api_key_dict_to_metadata(user_api_key_dict)
if user_metadata:
data["litellm_metadata"] = user_metadata
Readers pick ``litellm_metadata`` once it exists, so on routes whose bucket is ``metadata`` it starts as a copy.
"""
if "litellm_metadata" in data:
return
from litellm.llms.base_llm.guardrail_translation.base_translation import (
BaseTranslation,
)
user_metadata: Final = BaseTranslation.transform_user_api_key_dict_to_metadata(user_api_key_dict)
request: Final = cast("Mapping[str, object]", data) # cast-ok: request bodies are str-keyed
route_metadata: Final = (
as_str_mapping(request.get("metadata")) if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES else None
)
seeded: Final = {**route_metadata, **user_metadata} if route_metadata is not None else user_metadata
if seeded:
data["litellm_metadata"] = seeded
class UnifiedLLMGuardrails(CustomLogger):
@ -291,9 +300,10 @@ class UnifiedLLMGuardrails(CustomLogger):
except ValueError:
return data # handle unmapped call types
endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]())
translation: Final = mappings[CallTypes(call_type)]()
endpoint_translation: Final = _as_endpoint_translation(translation)
_ensure_litellm_metadata(data, user_api_key_dict)
_ensure_litellm_metadata(data, user_api_key_dict, call_type)
await _scan_request_attachments(guardrail_to_apply, data, call_type) # pyright: ignore[reportUnknownArgumentType] # hook data is an untyped dict
@ -303,7 +313,11 @@ class UnifiedLLMGuardrails(CustomLogger):
litellm_logging_obj=data.get("litellm_logging_obj"),
)
if not guardrail_to_apply.records_own_guardrail_information:
if not (
guardrail_to_apply.records_own_guardrail_information
or guardrail_to_apply.adds_itself_to_applied_guardrails
or translation.lists_guardrail_when_called
):
add_guardrail_to_applied_guardrails_header(
request_data=data, guardrail_name=guardrail_to_apply.guardrail_name
)
@ -341,7 +355,7 @@ class UnifiedLLMGuardrails(CustomLogger):
endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]())
_ensure_litellm_metadata(data, user_api_key_dict)
_ensure_litellm_metadata(data, user_api_key_dict, call_type)
return await endpoint_translation.process_input_messages(
data=data,
@ -422,7 +436,11 @@ class UnifiedLLMGuardrails(CustomLogger):
)
return response
endpoint_translation: Final = _as_endpoint_translation(mappings[CallTypes(call_type)]())
translation: Final = mappings[CallTypes(call_type)]()
if not translation.scans_responses:
return response
endpoint_translation: Final = _as_endpoint_translation(translation)
try:
response = await endpoint_translation.process_output_response(
@ -440,7 +458,9 @@ class UnifiedLLMGuardrails(CustomLogger):
if e.original_response is None:
e.original_response = response
raise
if not guardrail_to_apply.records_own_guardrail_information:
if not (
guardrail_to_apply.records_own_guardrail_information or guardrail_to_apply.adds_itself_to_applied_guardrails
):
add_guardrail_to_applied_guardrails_header(
request_data=data, guardrail_name=guardrail_to_apply.guardrail_name
)

4
litellm/proxy/policy_engine/pipeline_executor.py
View file

@ -19,6 +19,7 @@ from litellm.constants import LOGS_GUARDRAIL_INFORMATION_MARKER
from litellm.integrations.custom_guardrail import (
CustomGuardrail,
ModifyResponseException,
runs_through_apply_guardrail,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import (
@ -573,9 +574,8 @@ class PipelineExecutor:
hook_input, scans_raw_request = _prepare_hook_input(step, callback, data, raw_request_snapshot)
snapshot_entries_before: Final = len(recorded_guardrail_information(hook_input))
# Use unified_guardrail path if callback implements apply_guardrail
target: CustomLogger = callback
use_unified: Final = PipelineExecutor.supports_unified_execution(callback)
use_unified: Final = runs_through_apply_guardrail(callback, call_type)
if use_unified and streaming_chunks is None:
hook_input["guardrail_to_apply"] = callback
target = UnifiedLLMGuardrails()

99
litellm/proxy/utils.py
View file

@ -144,8 +144,10 @@ from litellm.exceptions import (
SensitiveDataRouteException,
)
from litellm.integrations.custom_guardrail import (
APPLY_GUARDRAIL_ONLY_CALL_TYPES,
CustomGuardrail,
ModifyResponseException,
runs_through_apply_guardrail,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.prometheus import PrometheusLogger
@ -342,6 +344,9 @@ _VIEW_SETUP_DEADLINE_SECONDS: Final = 15 * 60.0
_VIEW_SETUP_GATE_TABLE: Final = '"LiteLLM_SpendLogs"'
_VIEW_SETUP_GATE_PROBE_ROWS: Final = TypeAdapter(tuple[Mapping[str, bool], ...])
_DELEGATED_ATTRIBUTE: Final = TypeAdapter(object)
_REQUESTED_GUARDRAILS: Final = TypeAdapter(tuple[object, ...])
_REQUEST_SCAN_HOOKS: Final = (GuardrailEventHooks.pre_call, GuardrailEventHooks.during_call)
_SHIPPED_GUARDRAIL_MODULES: Final = ("litellm.", "litellm_enterprise.")
_ViewSetupOutcome: TypeAlias = Literal["ready", "timed_out"]
_ViewSetupAttempt: TypeAlias = Literal["ready", "table_missing"] | Exception
@ -726,6 +731,45 @@ def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple
return (guardrails, others)
def _scans_requests(
guardrail: CustomGuardrail,
data: dict[str, object], # mutable-ok: the proxy's request dict
) -> bool:
return any(guardrail.should_run_guardrail(data=data, event_type=hook) is True for hook in _REQUEST_SCAN_HOOKS)
def _warn_on_guardrails_that_cannot_scan(
data: dict[str, object], # mutable-ok: the proxy's request dict
call_type: str,
) -> None:
unscannable: Final = tuple(
callback.guardrail_name
for callback in ProxyLogging.callback_capabilities().resolved_callbacks
if isinstance(callback, CustomGuardrail)
and callback.guardrail_name is not None
and "apply_guardrail" not in type(callback).__dict__
and _scans_requests(callback, data)
)
if unscannable:
verbose_proxy_logger.warning(
"Guardrails %s do not implement apply_guardrail, so on %s requests they scan only the fields their own hook reads",
", ".join(unscannable),
call_type,
)
def _requested_user_defined_guardrails(requested: frozenset[str], call_type: str) -> tuple[str, ...]:
return tuple(
callback.guardrail_name
for callback in ProxyLogging.callback_capabilities().resolved_callbacks
if isinstance(callback, CustomGuardrail)
and callback.guardrail_name is not None
and callback.guardrail_name in requested
and not runs_through_apply_guardrail(callback, call_type)
and not type(callback).__module__.startswith(_SHIPPED_GUARDRAIL_MODULES)
)
def _merge_pipeline_metadata_bucket(data: dict, bucket_key: str, modified_bucket_value: object) -> None:
if not isinstance(modified_bucket_value, dict):
return
@ -2374,11 +2418,7 @@ class ProxyLogging:
Returns:
Result from the guardrail execution
"""
# Use unified_guardrail if callback has apply_guardrail method
has_apply_guardrail: Final = "apply_guardrail" in type(callback).__dict__ and not getattr(
callback, "use_native_lifecycle_hooks", False
)
use_unified: Final = has_apply_guardrail and not (
use_unified: Final = runs_through_apply_guardrail(callback, call_type) and not (
hook_type == "during_call" and getattr(callback, "use_native_during_call_hook", False)
)
if use_unified:
@ -2703,7 +2743,7 @@ class ProxyLogging:
data.pop("prompt_version", None)
data.pop("prompt_environment", None)
def _process_guardrail_metadata(self, data: dict) -> None:
def _process_guardrail_metadata(self, data: dict, call_type: CallTypesLiteral) -> None:
"""Process guardrails from metadata and add to applied_guardrails."""
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
@ -2718,6 +2758,22 @@ class ProxyLogging:
elif isinstance(metadata_litellm, dict) and "guardrails" in metadata_litellm:
guardrails_in_metadata = metadata_litellm.get("guardrails", [])
if call_type in APPLY_GUARDRAIL_ONLY_CALL_TYPES:
request: Final = cast("dict[str, object]", data) # cast-ok: request bodies are str-keyed
_warn_on_guardrails_that_cannot_scan(request, call_type)
requested: Final = frozenset(
name
for name in (
_REQUESTED_GUARDRAILS.validate_python(guardrails_in_metadata)
if isinstance(guardrails_in_metadata, list)
else ()
)
if isinstance(name, str)
)
for guardrail_name in _requested_user_defined_guardrails(requested, call_type):
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_name)
return
if guardrails_in_metadata and isinstance(guardrails_in_metadata, list):
applied_guardrails = []
if isinstance(metadata_standard, dict) and "applied_guardrails" in metadata_standard:
@ -3019,7 +3075,7 @@ class ProxyLogging:
and (guardrails_only or not caps.has_pre_call_override)
):
if data is not None:
self._process_guardrail_metadata(data)
self._process_guardrail_metadata(data, call_type)
return data
parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = (
@ -3123,17 +3179,17 @@ class ProxyLogging:
)
if data is not None:
self._process_guardrail_metadata(data)
self._process_guardrail_metadata(data, call_type)
return data
except SensitiveDataRouteException as e:
data = await self._handle_sensitive_data_route_exception(e, data, user_api_key_dict)
if data is not None:
self._process_guardrail_metadata(data)
self._process_guardrail_metadata(data, call_type)
return data
except Exception:
if data is not None:
self._process_guardrail_metadata(data)
self._process_guardrail_metadata(data, call_type)
raise
async def _run_parallel_pre_call_guardrails(
@ -3450,14 +3506,14 @@ class ProxyLogging:
_callback_capabilities = callback_capabilities
@staticmethod
def _stream_requires_guardrail_translation(user_api_key_dict: UserAPIKeyAuth) -> bool:
def _route_call_type(user_api_key_dict: UserAPIKeyAuth) -> CallTypes | None:
route: Final = user_api_key_dict.request_route
if not route:
return False
call_types: Final = get_call_types_for_route(route)
if not call_types:
return False
return call_types[0] in NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES
call_types: Final = get_call_types_for_route(route) if route else None
return call_types[0] if call_types else None
@staticmethod
def _stream_requires_guardrail_translation(user_api_key_dict: UserAPIKeyAuth) -> bool:
return ProxyLogging._route_call_type(user_api_key_dict) in NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES
@staticmethod
def has_post_call_response_headers_callbacks() -> bool:
@ -3573,8 +3629,7 @@ class ProxyLogging:
call_type: CallTypesLiteral,
) -> None:
if (
"apply_guardrail" in type(callback).__dict__
and not callback.use_native_lifecycle_hooks
runs_through_apply_guardrail(callback, call_type)
and user_api_key_dict is not None
and not callback.use_native_during_call_hook
):
@ -4144,6 +4199,7 @@ class ProxyLogging:
pipeline_managed: Final = pipeline_managed_guardrail_names(data, "post_call")
guardrail_callbacks, other_callbacks = _partition_post_call_callbacks()
call_type: Final = ProxyLogging._route_call_type(user_api_key_dict)
try:
# Merge model-level guardrails before checking which guardrails to run
guardrail_data: Final = check_and_merge_model_level_guardrails(data=data, llm_router=llm_router)
@ -4175,7 +4231,7 @@ class ProxyLogging:
guardrail_response: Any | None = None
if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks:
if runs_through_apply_guardrail(callback, call_type):
data["guardrail_to_apply"] = callback
guardrail_response = await self._run_guardrail_with_metrics(
callback,
@ -4246,11 +4302,12 @@ class ProxyLogging:
immediately before awaiting, and the unified hook pops it before its first
suspension point, so concurrent guardrails never race on that key.
"""
call_type: Final = ProxyLogging._route_call_type(user_api_key_dict)
async def _run_one(callback: CustomGuardrail) -> None:
if callback.should_run_guardrail(data=guardrail_data, event_type=GuardrailEventHooks.post_call) is not True:
return
if "apply_guardrail" in type(callback).__dict__ and not callback.use_native_lifecycle_hooks:
if runs_through_apply_guardrail(callback, call_type):
data["guardrail_to_apply"] = callback
await self._run_guardrail_with_metrics(
callback,

2
tests/code_coverage_tests/recursive_detector.py
View file

@ -3,6 +3,8 @@ import os
IGNORE_FUNCTIONS = [
"_json_cost", # bounded at depth 32 and consumed under byte/node limits.
"_json_text_forms", # walks a value TypeAdapter(JsonValue) already accepted, which caps depth near 255.
"_guardrailed_json", # walks a value TypeAdapter(JsonValue) already accepted, which caps depth near 255.
"_format_type",
"remove_additional_properties",
"remove_strict_from_schema",

95
tests/unit/enterprise/enterprise_callbacks/test_secret_detection.py
View file

@ -6,8 +6,9 @@ Covers the three defects from the ticket:
Spend Logs / the guardrails monitor show activity, with hits ("mask" +
masked_entity_count) distinguishable from clean requests ("allow").
- Defining ``apply_guardrail`` must NOT reroute proxied traffic off the
native ``async_pre_call_hook`` (per-key opt-out and ``data["prompt"]``
handling live only on the native path).
native ``async_pre_call_hook`` (``data["prompt"]`` handling lives only on
the native path), except on the decision and extraction routes, whose
request text the native hook never reads.
"""
import cProfile
@ -1134,13 +1135,99 @@ async def test_pre_call_hook_still_redacts_text_completion_prompt():
def test_proxied_traffic_stays_on_native_hooks():
"""Implementing apply_guardrail must not reroute proxied requests onto the
unified path: that path skips ``should_run_check`` (per-key opt-out) and
never sees ``data["prompt"]``."""
unified path: that path never sees ``data["prompt"]``."""
guardrail = _guardrail()
assert guardrail.uses_apply_guardrail_interface() is True
assert guardrail._deployment_hook_target() is guardrail
@pytest.mark.asyncio
@pytest.mark.parametrize(
("permissions", "expected_state", "listed"),
(
({}, "my key is [REDACTED]", True),
({"hide_secrets": False}, f"my key is {AWS_KEY}", False),
),
)
async def test_proxy_pre_call_hook_redacts_extraction_state_unless_the_key_opted_out(
monkeypatch: pytest.MonkeyPatch, permissions: dict[str, bool], expected_state: str, listed: bool
) -> None:
import litellm
from litellm.proxy.utils import ProxyLogging
monkeypatch.setattr(litellm, "callbacks", [_guardrail()])
key: Final = UserAPIKeyAuth(api_key="sk-9876", permissions=permissions)
data: Final = {
"model": "invoice-reader",
"state": f"my key is {AWS_KEY}",
"fields": {"total": {"type": "number"}},
"metadata": {"user_api_key_auth": key},
}
result: Final = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
user_api_key_dict=key, data=data, call_type="aextraction"
)
assert result["state"] == expected_state
buckets: Final = (result.get("metadata") or {}, result.get("litellm_metadata") or {})
applied: Final = [*buckets[0].get("applied_guardrails", []), *buckets[1].get("applied_guardrails", [])]
assert ("hide-secrets" in applied) is listed
@pytest.mark.asyncio
async def test_proxy_pre_call_hook_redacts_extraction_state_when_the_request_carries_no_key(
monkeypatch: pytest.MonkeyPatch,
) -> None:
import litellm
from litellm.proxy.utils import ProxyLogging
monkeypatch.setattr(litellm, "callbacks", [_guardrail()])
result: Final = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"),
data={
"model": "invoice-reader",
"state": f"my key is {AWS_KEY}",
"fields": {"total": {"type": "number"}},
"metadata": {},
},
call_type="aextraction",
)
assert result["state"] == "my key is [REDACTED]"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"secret_entry",
(
{"aws_secret_access_key": "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"},
{"db_password": "hunter2-Zq9!xk"},
),
)
async def test_proxy_pre_call_hook_redacts_a_systemone_secret_stored_under_a_json_key(
monkeypatch: pytest.MonkeyPatch, secret_entry: dict[str, str]
) -> None:
import litellm
from litellm.proxy.utils import ProxyLogging
monkeypatch.setattr(litellm, "callbacks", [_guardrail()])
key: Final = UserAPIKeyAuth(api_key="sk-9876")
result: Final = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
user_api_key_dict=key,
data={
"model": "judge",
"state": {**secret_entry, "ticket": "refund order 7"},
"questions": {"leak": {"type": "noul", "instructions": "Does the state hold a credential?"}},
"metadata": {"user_api_key_auth": key},
},
call_type="asystemone",
)
assert result["state"] == {**dict.fromkeys(secret_entry, "[REDACTED]"), "ticket": "refund order 7"}
@pytest.mark.asyncio
async def test_apply_guardrail_without_texts_records_nothing():
"""No inputs means nothing was inspected, so no "allow" row is recorded.

38
tests/unit/integrations/test_custom_guardrail.py
View file

@ -15,6 +15,7 @@ from litellm.integrations.custom_guardrail import (
_request_is_streaming,
guardrail_request_data_with_streaming,
log_guardrail_information,
runs_through_apply_guardrail,
without_server_streaming_classification,
)
from litellm.litellm_core_utils.litellm_logging import Logging
@ -3786,3 +3787,40 @@ def test_default_on_guardrail() -> None:
)
is True
)
class _NativeHooksWithApplyGuardrail(_UndecoratedGuardrail):
use_native_lifecycle_hooks: ClassVar[bool] = True
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: "LiteLLMLoggingObj | None" = None,
) -> GenericGuardrailAPIInputs:
return inputs
@pytest.mark.parametrize(
("guardrail_type", "call_type", "expected"),
(
(_UndecoratedGuardrail, "acompletion", True),
(_UndecoratedGuardrail, "aextraction", True),
(_NativeHooksWithApplyGuardrail, "acompletion", False),
(_NativeHooksWithApplyGuardrail, "aresponses", False),
(_NativeHooksWithApplyGuardrail, None, False),
(_NativeHooksWithApplyGuardrail, "extraction", True),
(_NativeHooksWithApplyGuardrail, "aextraction", True),
(_NativeHooksWithApplyGuardrail, "decisions", True),
(_NativeHooksWithApplyGuardrail, "adecisions", True),
(_NativeHooksWithApplyGuardrail, "systemone", True),
(_NativeHooksWithApplyGuardrail, "asystemone", True),
(_NativeApplyGuardrail, "aextraction", False),
(_InheritedApplyGuardrail, "aextraction", False),
),
)
def test_native_hook_guardrails_take_apply_guardrail_only_where_their_hooks_read_no_text(
guardrail_type: type[CustomGuardrail], call_type: str | None, expected: bool
) -> None:
assert runs_through_apply_guardrail(guardrail_type(guardrail_name="g"), call_type) is expected

0
tests/unit/llms/openai/decisions/guardrail_translation/__init__.py Normal file
View file

467
tests/unit/llms/openai/decisions/guardrail_translation/test_decisions_guardrail_handler.py Normal file
View file

@ -0,0 +1,467 @@
import re
from functools import reduce
from typing import Final, Literal
import pytest
from litellm_enterprise.enterprise_callbacks.secret_detection import _ENTERPRISE_SecretDetection
from pydantic import JsonValue
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content
from litellm.litellm_core_utils.prompt_templates.factory import resolve_structured_messages
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms import get_guardrail_translation_mapping
from litellm.llms.base_llm.guardrail_translation.utils import UnappliableRequestRewrite
from litellm.llms.openai.decisions.guardrail_translation.handler import DecisionsHandler, SystemOneHandler
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
from litellm.types.guardrails import ContentFilterAction, ContentFilterPattern
from litellm.types.decisions import OpenAIDecisionResponse
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs
class _TaggingGuardrail(CustomGuardrail):
def __init__(self) -> None:
super().__init__(guardrail_name="tagging")
self.seen: list[tuple[str, tuple[str, ...]]] = []
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:
texts: Final = tuple(inputs.get("texts", []))
self.seen.append((input_type, texts))
return GenericGuardrailAPIInputs(texts=[f"{text} [GUARDRAILED]" for text in texts])
class _MasksSecrets(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:
secrets: Final = (
("AKIAIOSFODNN7EXAMPLE", "<AWS_KEY>"),
("AKIAI44QH8DHBEXAMPLE", "<AWS_KEY>"),
("4111111111111111", "<CARD>"),
)
return GenericGuardrailAPIInputs(
texts=[reduce(lambda text, secret: text.replace(*secret), secrets, text) for text in inputs.get("texts", [])]
)
class _DropsQuotes(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:
return GenericGuardrailAPIInputs(texts=[text.replace('"', "") for text in inputs.get("texts", [])])
class _MasksOverlappingPairs(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:
return GenericGuardrailAPIInputs(
texts=[text.replace("b\nc", "[BC]").replace("a b", "[AB]") for text in inputs.get("texts", [])]
)
class _NullsAMaskedValue(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:
return GenericGuardrailAPIInputs(
texts=[re.sub(r'": ".*"$', '": null', text).replace("b", "[B]") for text in inputs.get("texts", [])]
)
class _ReturnsNoTexts(CustomGuardrail):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:
return GenericGuardrailAPIInputs()
@pytest.mark.parametrize(
("call_type", "handler"),
(
(CallTypes.decisions, DecisionsHandler),
(CallTypes.adecisions, DecisionsHandler),
(CallTypes.systemone, SystemOneHandler),
(CallTypes.asystemone, SystemOneHandler),
),
)
def test_decision_call_types_resolve_to_their_handler(call_type: CallTypes, handler: type[object]) -> None:
assert get_guardrail_translation_mapping(call_type) is handler
@pytest.mark.asyncio
async def test_decision_message_text_parts_and_question_text_are_guardrailed() -> None:
guardrail: Final = _TaggingGuardrail()
image: Final = {"type": "input_image", "image_url": "data:image/png;base64,AA=="}
data: Final = {
"model": "judge",
"input": [
{"role": "user", "content": [{"type": "input_text", "text": "ticket one"}, image]},
{"role": "user", "content": "ticket two"},
],
"questions": [
{"type": "predicate", "name": "urgent", "instructions": "Is it urgent?"},
{"type": "box", "name": "logo", "instructions": "ticket one"},
{
"type": "choice",
"name": "team",
"instructions": "Which team?",
"choices": [{"value": "billing", "description": "Payment issues"}, {"value": True}],
},
{
"type": "score",
"name": "mood",
"instructions": "How upset?",
"levels": [{"label": "calm"}, {"label": "angry", "description": "Uses caps"}],
},
],
}
result: Final = await DecisionsHandler().process_input_messages(data, guardrail)
assert guardrail.seen == [
(
"request",
("ticket one", "ticket two", "Is it urgent?", "Which team?", "Payment issues", "How upset?", "Uses caps"),
)
]
assert {key: value for key, value in result.items() if key != "metadata"} == {
"model": "judge",
"input": [
{"role": "user", "content": [{"type": "input_text", "text": "ticket one [GUARDRAILED]"}, image]},
{"role": "user", "content": "ticket two [GUARDRAILED]"},
],
"questions": [
{"type": "predicate", "name": "urgent", "instructions": "Is it urgent? [GUARDRAILED]"},
{"type": "box", "name": "logo", "instructions": "ticket one [GUARDRAILED]"},
{
"type": "choice",
"name": "team",
"instructions": "Which team? [GUARDRAILED]",
"choices": [{"value": "billing", "description": "Payment issues [GUARDRAILED]"}, {"value": True}],
},
{
"type": "score",
"name": "mood",
"instructions": "How upset? [GUARDRAILED]",
"levels": [{"label": "calm"}, {"label": "angry", "description": "Uses caps [GUARDRAILED]"}],
},
],
}
assert data["input"][0]["content"][0]["text"] == "ticket one"
@pytest.mark.asyncio
async def test_systemone_state_keys_scalars_and_question_text_are_guardrailed() -> None:
guardrail: Final = _TaggingGuardrail()
data: Final = {
"model": "judge",
"state": {"subject": "refund", "lines": ["order 7", "", 3, True], "meta": {"id": None}},
"questions": {
"refund": {"type": "noul", "instructions": {"ask": "Does the customer want a refund?"}},
"routed": {"type": "noul", "criteria": {"true": "The customer asks for money back", "false": None}},
"tier": {"type": "score", "criteria": ["calm", {"level": "angry"}]},
},
}
result: Final = await SystemOneHandler().process_input_messages(data, guardrail)
assert guardrail.seen == [
(
"request",
(
'"subject": "refund"',
"lines",
"order 7",
"3",
"meta",
'"id": null',
'"ask": "Does the customer want a refund?"',
"The customer asks for money back",
"calm",
'"level": "angry"',
),
)
]
assert {key: value for key, value in result.items() if key != "metadata"} == {
"model": "judge",
"state": {
"subject": '"refund" [GUARDRAILED]',
"lines [GUARDRAILED]": ["order 7 [GUARDRAILED]", "", "3 [GUARDRAILED]", True],
"meta [GUARDRAILED]": {"id": "null [GUARDRAILED]"},
},
"questions": {
"refund": {
"type": "noul",
"instructions": {"ask": '"Does the customer want a refund?" [GUARDRAILED]'},
},
"routed": {
"type": "noul",
"criteria": {"true": "The customer asks for money back [GUARDRAILED]", "false": None},
},
"tier": {"type": "score", "criteria": ["calm [GUARDRAILED]", {"level": '"angry" [GUARDRAILED]'}]},
},
}
@pytest.mark.asyncio
async def test_systemone_secret_key_and_card_number_are_masked_before_the_provider() -> None:
data: Final = {
"model": "judge",
"state": {"AKIAIOSFODNN7EXAMPLE": "<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_masks_a_card_number_that_json_escaping_splits() -> None:
guardrail: Final = ContentFilterGuardrail(
guardrail_name="cards",
patterns=[ContentFilterPattern(pattern_type="prebuilt", pattern_name="visa", action=ContentFilterAction.MASK)],
)
data: Final = {
"model": "judge",
"state": {"card": "4111\t1111\t1111\t1111", "4111\n1111\n1111\n1111": "on file", "note": "first\nsecond"},
"questions": {"q": {"type": "noul", "criteria": {}}},
}
result: Final = await SystemOneHandler().process_input_messages(data, guardrail)
assert result["state"] == {"card": "[VISA_REDACTED]", "[VISA_REDACTED]": "on file", "note": "first\nsecond"}
@pytest.mark.asyncio
@pytest.mark.parametrize(
("pattern_names", "state", "masked"),
(
pytest.param(
("visa",),
{"4111111111111111": "4111\t1111\t1111\t1111"},
{"[VISA_REDACTED]": "[VISA_REDACTED]"},
id="key-and-escaped-value",
),
pytest.param(
("email",),
{"note": "Please call me back.\ndana.quill@example.com"},
{"note": "Please call me back.\n[EMAIL_REDACTED]"},
id="match-after-an-escape",
),
pytest.param(
("au_tfn", "email"),
{"TFN": "123456789 contact dana@example.com\n"},
{"TFN": "[AU_TFN_REDACTED] contact [EMAIL_REDACTED]\n"},
id="key-context-match-in-an-escaped-value",
),
pytest.param(
("au_tfn", "visa"),
{"TFN": "123456789\n4111\t1111\t1111\t1111"},
{"TFN": "[AU_TFN_REDACTED]\n[VISA_REDACTED]"},
id="key-context-match-beside-a-split-match",
),
),
)
async def test_systemone_masks_findings_on_both_sides_of_a_json_escape(
pattern_names: tuple[str, ...], state: dict[str, JsonValue], masked: dict[str, JsonValue]
) -> None:
guardrail: Final = ContentFilterGuardrail(
guardrail_name="mask",
patterns=[
ContentFilterPattern(pattern_type="prebuilt", pattern_name=name, action=ContentFilterAction.MASK)
for name in pattern_names
],
)
data: Final = {"model": "judge", "state": state, "questions": {"q": {"type": "noul", "criteria": {}}}}
result: Final = await SystemOneHandler().process_input_messages(data, guardrail)
assert result["state"] == masked
@pytest.mark.asyncio
async def test_systemone_keeps_a_secret_mask_that_covers_an_escaped_value() -> None:
data: Final = {
"model": "judge",
"state": {"password": "hunter2\nAKIAIOSFODNN7EXAMPLE"},
"questions": {"q": {"type": "noul", "criteria": {}}},
}
result: Final = await SystemOneHandler().process_input_messages(
data, _ENTERPRISE_SecretDetection(guardrail_name="hide-secrets")
)
assert result["state"] == {"password": "[REDACTED]"}
@pytest.mark.asyncio
@pytest.mark.parametrize(
("guardrail", "note"),
(
pytest.param(_MasksOverlappingPairs(guardrail_name="pairs"), "a b\nc", id="partly-overlapping-masks"),
pytest.param(_NullsAMaskedValue(guardrail_name="nulls"), "a\nb", id="value-turned-into-null"),
),
)
async def test_systemone_rejects_an_escaped_entry_whose_two_rewrites_cannot_merge(
guardrail: CustomGuardrail, note: str
) -> None:
data: Final = {"model": "judge", "state": {"note": note}, "questions": {"q": {"type": "noul", "criteria": {}}}}
with pytest.raises(UnappliableRequestRewrite):
await SystemOneHandler().process_input_messages(data, guardrail)
@pytest.mark.asyncio
async def test_systemone_keeps_every_state_entry_when_masked_keys_collide() -> None:
data: Final = {
"model": "judge",
"state": {"AKIAIOSFODNN7EXAMPLE": "first", "items": 2, "AKIAI44QH8DHBEXAMPLE": "second"},
"questions": {"q": {"type": "noul", "criteria": {}}},
}
result: Final = await SystemOneHandler().process_input_messages(data, _MasksSecrets(guardrail_name="mask"))
assert result["state"] == {"<AWS_KEY>": "first", "items": 2, "<AWS_KEY>_2": "second"}
@pytest.mark.asyncio
async def test_systemone_rejects_a_rewrite_that_loses_the_state_key() -> None:
data: Final = {
"model": "judge",
"state": {"ticket": "refund order 7"},
"questions": {"q": {"type": "noul", "criteria": {}}},
}
with pytest.raises(UnappliableRequestRewrite):
await SystemOneHandler().process_input_messages(data, _DropsQuotes(guardrail_name="drops"))
@pytest.mark.asyncio
async def test_systemone_rejects_masked_keys_that_stay_ambiguous() -> None:
data: Final = {
"model": "judge",
"state": {"AKIAIOSFODNN7EXAMPLE": "first", "AKIAI44QH8DHBEXAMPLE": "second", "<AWS_KEY>_2": "third"},
"questions": {"q": {"type": "noul", "criteria": {}}},
}
with pytest.raises(UnappliableRequestRewrite):
await SystemOneHandler().process_input_messages(data, _MasksSecrets(guardrail_name="mask"))
@pytest.mark.asyncio
@pytest.mark.parametrize(
("handler", "data"),
(
(
SystemOneHandler,
{"model": "judge", "state": {}, "questions": {"q": {"type": "noul", "criteria": {}}}},
),
(
DecisionsHandler,
{
"model": "judge",
"input": [{"role": "user", "content": ""}],
"questions": [{"type": "predicate", "instructions": ""}],
},
),
),
)
async def test_request_without_text_is_returned_unscanned(
handler: type[DecisionsHandler] | type[SystemOneHandler], data: dict[str, object]
) -> None:
guardrail: Final = _TaggingGuardrail()
result: Final = await handler().process_input_messages(data, guardrail)
assert result is data
assert guardrail.seen == []
@pytest.mark.asyncio
@pytest.mark.parametrize(
("handler", "data"),
(
(
DecisionsHandler,
{"model": "judge", "input": "ticket", "questions": [{"type": "predicate", "instructions": "Is it urgent?"}]},
),
(
SystemOneHandler,
{"model": "judge", "state": {"total": 3, "paid": 2.5}, "questions": {"q": {"type": "noul", "criteria": {}}}},
),
),
)
async def test_guardrail_that_returns_no_texts_leaves_the_request_text_as_sent(
handler: type[DecisionsHandler] | type[SystemOneHandler], data: dict[str, object]
) -> None:
result: Final = await handler().process_input_messages(data, _ReturnsNoTexts(guardrail_name="noop"))
assert result == data
@pytest.mark.asyncio
@pytest.mark.parametrize("handler", (DecisionsHandler, SystemOneHandler))
async def test_decision_answers_are_returned_as_the_provider_sent_them(
handler: type[DecisionsHandler] | type[SystemOneHandler],
) -> None:
guardrail: Final = _TaggingGuardrail()
response: Final = OpenAIDecisionResponse.model_validate(
{
"model": "judge",
"answers": [{"type": "refusal", "name": "q"}],
"usage": {"input_tokens": 1, "output_tokens": 0, "total_tokens": 1},
}
)
result: Final = await handler().process_output_response(response, guardrail)
assert result is response
assert guardrail.seen == []
def test_router_sees_the_image_in_a_decision_input_routed_by_its_proxy_route() -> None:
structured: Final = resolve_structured_messages(
messages=None,
request_kwargs={
"input": [
{
"role": "user",
"content": [
{"type": "input_text", "text": "Is this a cat?"},
{"type": "input_image", "image_url": "https://example.com/cat.png"},
],
}
],
"litellm_metadata": {"user_api_key_request_route": "/v1/decisions"},
},
)
assert structured is not None and request_contains_image_content(structured)

22
tests/unit/proxy/common_utils/test_callback_utils.py
View file

@ -205,6 +205,28 @@ def test_get_logging_caching_headers_merges_metadata_and_litellm_metadata():
assert headers["x-litellm-policy-sources"] == "global-baseline=team_default"
@pytest.mark.parametrize(
("metadata_guardrails", "litellm_metadata_guardrails", "expected"),
(
(["lakera-v1"], ["hide-secrets"], "lakera-v1,hide-secrets"),
(["hide-secrets"], ["hide-secrets", "lakera-v1"], "hide-secrets,lakera-v1"),
),
)
def test_applied_guardrails_header_lists_guardrails_from_both_metadata_buckets(
metadata_guardrails: list[str], litellm_metadata_guardrails: list[str], expected: str
) -> None:
request_data: Final = {
"metadata": {"applied_guardrails": metadata_guardrails},
"litellm_metadata": {"applied_guardrails": litellm_metadata_guardrails, "guardrail_scan_ids": ("scan-1",)},
}
headers: Final = get_logging_caching_headers(request_data)
assert headers is not None
assert headers["x-litellm-applied-guardrails"] == expected
assert headers["x-litellm-guardrail-scan-id"] == "scan-1"
def _record(
request_data: dict[str, object],
scan_id: str | None,

93
tests/unit/proxy/guardrails/guardrail_hooks/dynamoai/test_dynamoai.py
View file

@ -7,12 +7,15 @@ import json
import os
from unittest.mock import AsyncMock, MagicMock, patch
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Final
import pytest
import respx
import litellm
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import CallTypesLiteral
from litellm.proxy.guardrails.guardrail_hooks.dynamoai import DynamoAIGuardrails
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
@ -203,6 +206,96 @@ async def test_post_call_hook_skips_dynamoai_when_the_reply_has_no_text(respx_mo
assert route.called is False
_INPUT_HOOKS = [
pytest.param(
"pre_call",
lambda guardrail, data, call_type: guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type
),
id="pre_call",
),
pytest.param(
"during_call",
lambda guardrail, data, call_type: guardrail.async_moderation_hook(
data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type
),
id="during_call",
),
]
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS)
@pytest.mark.parametrize(
("request_fields", "call_type", "applied"),
[
pytest.param({"messages": [{"role": "user", "content": "hi"}]}, "completion", ["test-dynamoai"], id="messages"),
pytest.param({"state": "hi"}, "aextraction", None, id="no-messages-on-extraction"),
pytest.param({"state": "hi"}, "completion", ["test-dynamoai"], id="no-messages-on-chat"),
],
)
async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text(
respx_mock: respx.MockRouter,
event_hook: str,
run_hook: Callable[[DynamoAIGuardrails, Mapping[str, object], CallTypesLiteral], Awaitable[object]],
request_fields: Mapping[str, object],
call_type: CallTypesLiteral,
applied: Sequence[str] | None,
) -> None:
respx_mock.post(_ANALYZE_URL).respond(json={"finalAction": "NONE", "appliedPolicies": []})
guardrail: Final = DynamoAIGuardrails(
guardrail_name="test-dynamoai",
api_key="test-api-key",
api_base="https://api.dynamo.ai",
event_hook=event_hook,
default_on=True,
)
metadata: Final[dict[str, object]] = {}
await run_hook(guardrail, {**request_fields, "metadata": metadata}, call_type)
assert metadata.get("applied_guardrails") == applied
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.parametrize(
("response", "request_route", "applied"),
[
pytest.param(_assistant_reply("Hello, how are you?"), "/v1/chat/completions", ["test-dynamoai"], id="chat-reply"),
pytest.param(
{"answers": {"aws_key": {"type": "text", "value": None}}},
"/v1/extraction",
None,
id="no-choices-on-extraction",
),
pytest.param(
{"answers": {"aws_key": {"type": "text", "value": None}}},
"/v1/responses",
["test-dynamoai"],
id="no-choices-on-responses",
),
],
)
async def test_post_call_hook_lists_the_guardrail_only_when_it_scanned_text(
respx_mock: respx.MockRouter,
response: litellm.ModelResponse | Mapping[str, object],
request_route: str,
applied: Sequence[str] | None,
) -> None:
respx_mock.post(_ANALYZE_URL).respond(json={"finalAction": "NONE", "appliedPolicies": []})
metadata: Final[dict[str, object]] = {}
await _post_call_guardrail().async_post_call_success_hook(
data={"messages": [{"role": "user", "content": "hi"}], "metadata": metadata},
user_api_key_dict=UserAPIKeyAuth(request_route=request_route),
response=response,
)
assert metadata.get("applied_guardrails") == applied
@pytest.fixture(autouse=True)
def _vcr_outcome_gate(request, vcr):
install_live_call_probe(request, vcr)

109
tests/unit/proxy/guardrails/guardrail_hooks/ibm_guardrails/test_ibm_detector.py
View file

@ -1,10 +1,15 @@
import json
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Final
import pytest
import respx
import litellm
from litellm.caching.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import CallTypesLiteral
from litellm.proxy.guardrails.guardrail_hooks.ibm_guardrails.ibm_detector import IBMGuardrailDetector
pytestmark = pytest.mark.usefixtures("httpx_transport")
@ -122,3 +127,107 @@ async def test_orchestrator_detection_blocks_the_request_and_names_the_detector(
)
assert "- HAP (detector: hap-detector, score: 0.970)\n Text: 'hello'" in str(blocked.value)
_INPUT_HOOKS = [
pytest.param(
"pre_call",
lambda guardrail, data, call_type: guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type
),
id="pre_call",
),
pytest.param(
"during_call",
lambda guardrail, data, call_type: guardrail.async_moderation_hook(
data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type
),
id="during_call",
),
]
_REQUEST_TEXT_CASES = [
pytest.param({"messages": [{"role": "user", "content": "hello there"}]}, "completion", ["ibm-guard"], id="messages"),
pytest.param({"state": "hello there"}, "aextraction", None, id="no-messages-on-extraction"),
pytest.param({"state": "hello there"}, "completion", ["ibm-guard"], id="no-messages-on-chat"),
]
@pytest.mark.asyncio
@pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS)
@pytest.mark.parametrize(("request_fields", "call_type", "applied"), _REQUEST_TEXT_CASES)
async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text(
respx_mock: respx.MockRouter,
event_hook: str,
run_hook: Callable[[IBMGuardrailDetector, Mapping[str, object], CallTypesLiteral], Awaitable[object]],
request_fields: Mapping[str, object],
call_type: CallTypesLiteral,
applied: Sequence[str] | None,
) -> None:
respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[]])
guardrail: Final = IBMGuardrailDetector(
guardrail_name="ibm-guard",
auth_token="ibm-token",
base_url=_BASE_URL,
detector_id="hap-detector",
is_detector_server=True,
event_hook=event_hook,
default_on=True,
)
metadata: Final[dict[str, object]] = {}
await run_hook(guardrail, {**request_fields, "metadata": metadata}, call_type)
assert metadata.get("applied_guardrails") == applied
@pytest.mark.asyncio
@pytest.mark.parametrize(
("response", "request_route", "applied"),
[
pytest.param(
litellm.ModelResponse(
choices=[{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "hi"}}]
),
"/v1/chat/completions",
["ibm-guard"],
id="chat-reply",
),
pytest.param(
{"answers": {"aws_key": {"type": "text", "value": None}}},
"/v1/extraction",
None,
id="no-choices-on-extraction",
),
pytest.param(
{"answers": {"aws_key": {"type": "text", "value": None}}},
"/v1/responses",
["ibm-guard"],
id="no-choices-on-responses",
),
],
)
async def test_post_call_hook_lists_the_guardrail_only_when_it_scanned_text(
respx_mock: respx.MockRouter,
response: litellm.ModelResponse | Mapping[str, object],
request_route: str,
applied: Sequence[str] | None,
) -> None:
respx_mock.post(_DETECTOR_SERVER_URL).respond(json=[[]])
guardrail: Final = IBMGuardrailDetector(
guardrail_name="ibm-guard",
auth_token="ibm-token",
base_url=_BASE_URL,
detector_id="hap-detector",
is_detector_server=True,
event_hook="post_call",
default_on=True,
)
metadata: Final[dict[str, object]] = {}
await guardrail.async_post_call_success_hook(
data={"messages": [{"role": "user", "content": "hello"}], "metadata": metadata},
user_api_key_dict=UserAPIKeyAuth(request_route=request_route),
response=response,
)
assert metadata.get("applied_guardrails") == applied

41
tests/unit/proxy/guardrails/guardrail_hooks/test_aim.py
View file

@ -7,6 +7,9 @@ from copy import deepcopy
from unittest.mock import AsyncMock, patch
import httpx
from collections.abc import Mapping, Sequence
from typing import Final
from litellm.types.utils import CallTypesLiteral
import pytest
import respx
from httpx import Request, Response
@ -316,7 +319,7 @@ async def test_aim_anonymize_action_blocks_malformed_redacted_messages(
)
assert exc_info.value.code == "400"
assert data == request_body
assert {key: value for key, value in data.items() if key != "metadata"} == request_body
_OUTPUT_REQUEST = {
@ -815,3 +818,39 @@ async def test_aim_streaming_block_stops_after_verified_chunks(monkeypatch: pyte
{"choices": [{"delta": {"content": "A"}}]},
{"done": True},
]
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.parametrize(
("call_type", "request_fields", "aim_calls", "applied"),
[
pytest.param(
"adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, 1, ["aim"], id="decisions"
),
pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"),
pytest.param("acompletion", {"messages": [{"role": "user", "content": "Refund?"}]}, 1, None, id="chat"),
],
)
async def test_aim_calls_aim_and_lists_itself_only_when_the_request_has_text(
respx_mock: respx.MockRouter,
call_type: CallTypesLiteral,
request_fields: Mapping[str, object],
aim_calls: int,
applied: Sequence[str] | None,
) -> None:
analyze: Final = respx_mock.post("http://aim.test/fw/v1/analyze").respond(
json={"analysis_result": {"analysis_time_ms": 1, "policy_drill_down": {}}, "required_action": None}
)
guardrail: Final = AimGuardrail(api_key="hs-aim-key", api_base="http://aim.test", guardrail_name="aim")
metadata: Final[dict[str, object]] = {}
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
cache=DualCache(),
data={**request_fields, "metadata": metadata},
call_type=call_type,
)
assert analyze.call_count == aim_calls
assert metadata.get("applied_guardrails") == applied

40
tests/unit/proxy/guardrails/guardrail_hooks/test_cato_networks.py
View file

@ -3,7 +3,11 @@ import json
import ssl
from unittest.mock import AsyncMock, MagicMock, patch
from collections.abc import Mapping, Sequence
from typing import Final
from litellm.types.utils import CallTypesLiteral
import pytest
import respx
from fastapi.exceptions import HTTPException
from httpx import Request, Response
from websockets.exceptions import ConnectionClosed
@ -2785,3 +2789,39 @@ async def test_cato_still_inspects_every_conversational_call_type(call_type: str
)
mock_post.assert_called_once()
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.parametrize(
("call_type", "request_fields", "cato_calls", "applied"),
[
pytest.param(
"adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, 1, ["cato"], id="decisions"
),
pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"),
pytest.param("acompletion", {"messages": [{"role": "user", "content": "Refund?"}]}, 1, None, id="chat"),
],
)
async def test_cato_calls_cato_and_lists_itself_only_when_the_request_has_text(
respx_mock: respx.MockRouter,
call_type: CallTypesLiteral,
request_fields: Mapping[str, object],
cato_calls: int,
applied: Sequence[str] | None,
) -> None:
analyze: Final = respx_mock.post("http://cato.test/fw/v1/analyze").respond(
json={"analysis_result": {}, "required_action": None}
)
guardrail: Final = CatoNetworksGuardrail(api_key="cato-key", api_base="http://cato.test", guardrail_name="cato")
metadata: Final[dict[str, object]] = {}
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
cache=DualCache(),
data={**request_fields, "metadata": metadata},
call_type=call_type,
)
assert analyze.call_count == cato_calls
assert metadata.get("applied_guardrails") == applied

25
tests/unit/proxy/guardrails/guardrail_hooks/test_decision_model.py
View file

@ -3,7 +3,7 @@
from __future__ import annotations
import asyncio
from collections.abc import Callable
from collections.abc import Callable, Sequence
from typing import Final, Literal, cast
from unittest.mock import AsyncMock, MagicMock, patch
@ -1308,3 +1308,26 @@ async def test_tool_call_in_both_tool_calls_and_structured_messages_is_screened_
await guardrail.apply_guardrail(inputs, _request_data(), "request")
assert [call.kwargs["input"] for call in router.adecisions.await_args_list] == ['send_email({"body": "hello"})']
@pytest.mark.asyncio
@pytest.mark.parametrize(
("texts", "unreachable", "applied"),
[
pytest.param(["hello"], False, ["test_decision_model"], id="scanned"),
pytest.param([], False, None, id="no-text"),
pytest.param(["hello"], True, None, id="unreachable-fail-open"),
],
)
async def test_decision_model_lists_itself_only_when_it_returned_verdicts(
texts: list[str], unreachable: bool, applied: Sequence[str] | None
) -> None:
router: Final = _decision_router(probability=0.1)
if unreachable:
router.adecisions.side_effect = httpx.ConnectError("unreachable")
guardrail: Final = _make_guardrail(router_provider=lambda: router, unreachable_fallback="fail_open")
request_data: Final = _request_data()
await guardrail.apply_guardrail({"texts": texts}, request_data, "request")
assert cast(dict[str, object], request_data["metadata"]).get("applied_guardrails") == applied

104
tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py Normal file
View file

@ -0,0 +1,104 @@
from collections.abc import Mapping
from itertools import chain
from typing import Final
import httpx
import pytest
from pydantic import BaseModel, ConfigDict, TypeAdapter
from litellm.caching import DualCache
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation
from litellm.types.guardrails import GuardrailEventHooks
_LAKERA_REQUEST_BODY: Final = TypeAdapter(dict[str, object])
class _Metadata(BaseModel):
model_config = ConfigDict(frozen=True)
applied_guardrails: tuple[str, ...] = ()
def _applied_guardrails(data: Mapping[str, object]) -> tuple[str, ...]:
return tuple(
chain.from_iterable(
_Metadata.model_validate(data.get(bucket) or {}).applied_guardrails
for bucket in ("metadata", "litellm_metadata")
)
)
def _lakera(transport: httpx.MockTransport) -> lakeraAI_Moderation:
return lakeraAI_Moderation(
guardrail_name="lakera-v1",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
api_key="lk-test",
api_base="https://lakera.test",
async_handler=AsyncHTTPHandler(transport=transport),
)
def _refuse_any_request(request: httpx.Request) -> httpx.Response:
raise AssertionError(f"unexpected Lakera request with body {request.content!r}")
def _clean_scan_of(expected_input: str) -> httpx.MockTransport:
def respond(request: httpx.Request) -> httpx.Response:
assert _LAKERA_REQUEST_BODY.validate_json(request.content)["input"] == expected_input
return httpx.Response(200, json={"results": [{"flagged": False}]})
return httpx.MockTransport(respond)
@pytest.mark.asyncio
async def test_request_without_text_lakera_reads_skips_the_vendor_and_is_not_listed() -> None:
guardrail: Final = _lakera(httpx.MockTransport(_refuse_any_request))
data: Final[dict[str, object]] = {
"model": "m",
"state": "paid with sk-live-1234",
"fields": {"total": {"type": "number"}},
}
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), cache=DualCache(), data=data, call_type="aextraction"
)
assert _applied_guardrails(data) == ()
@pytest.mark.asyncio
async def test_clean_scan_of_string_input_lists_the_guardrail() -> None:
guardrail: Final = _lakera(_clean_scan_of("is this a refund request?"))
data: Final[dict[str, object]] = {
"model": "m",
"input": "is this a refund request?",
"questions": [{"type": "predicate", "instructions": "Is it a refund request?"}],
}
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), cache=DualCache(), data=data, call_type="adecisions"
)
assert _applied_guardrails(data) == ("lakera-v1",)
@pytest.mark.asyncio
async def test_clean_scan_of_chat_messages_keeps_the_header_as_before() -> None:
scans: Final[list[httpx.Request]] = []
def respond(request: httpx.Request) -> httpx.Response:
scans.append(request)
return httpx.Response(200, json={"results": [{"flagged": False}]})
guardrail: Final = _lakera(httpx.MockTransport(respond))
data: Final[dict[str, object]] = {"model": "m", "messages": [{"role": "user", "content": "is this a refund?"}]}
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"), cache=DualCache(), data=data, call_type="completion"
)
assert len(scans) == 1
assert _applied_guardrails(data) == ()

19
tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py
View file

@ -319,6 +319,25 @@ class TestAsyncPostCallSuccessHookSkipFlags:
assert "a@b.com" not in result_dict["choices"][0]["message"]["content"]
@pytest.mark.asyncio
async def test_post_call_skips_a_response_with_no_assistant_text() -> None:
guardrail: Final = LakeraAIGuardrail(api_key="test_key")
data: Final = {"model": "clairic", "metadata": {}}
llm_response: Final = MagicMock()
llm_response.model_dump.return_value = {"answer": "Berlin"}
with patch.object(guardrail, "call_v2_guard", new_callable=AsyncMock) as mock_call:
result: Final = await guardrail.async_post_call_success_hook(
data=data,
user_api_key_dict=UserAPIKeyAuth(api_key="test_key"),
response=llm_response,
)
mock_call.assert_not_called()
assert result is llm_response
assert "applied_guardrails" not in data["metadata"]
PII_ONLY_LAKERA_RESPONSE = {
"flagged": True,
"breakdown": [{"detector_type": "pii/email", "detected": True, "message_id": 0}],

43
tests/unit/proxy/guardrails/guardrail_hooks/test_lasso.py
View file

@ -1549,3 +1549,46 @@ async def test_api_error_handling_legacy_messages(
assert "Failed to verify request safety with Lasso API" in str(excinfo.value)
assert error_message in str(excinfo.value)
assert route.called
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.parametrize(
("call_type", "request_fields", "lasso_calls", "applied"),
[
pytest.param(
"adecisions", {"input": [{"role": "user", "content": "Is this a refund?"}]}, 1, ["lasso"], id="decisions"
),
pytest.param("aextraction", {"state": {"invoice": "INV-1"}}, 0, None, id="extraction"),
pytest.param("acompletion", {"messages": [{"role": "user", "content": "Refund?"}]}, 1, None, id="chat"),
],
)
@pytest.mark.parametrize("event_hook", ["pre_call", "during_call"])
async def test_lasso_lists_itself_only_when_it_classified_request_text(
respx_mock: respx.MockRouter,
event_hook: str,
call_type: str,
request_fields: dict[str, object],
lasso_calls: int,
applied: list[str] | None,
) -> None:
classify: Final = respx_mock.post("https://server.lasso.security/gateway/v3/classify").respond(
json={"deputies": {}, "findings": {}, "violations_detected": False}
)
guardrail: Final = LassoGuardrail(
lasso_api_key="test-api-key", guardrail_name="lasso", event_hook=event_hook, default_on=True
)
metadata: Final[dict[str, object]] = {}
data: Final = {**request_fields, "metadata": metadata}
if event_hook == "pre_call":
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type
)
else:
await guardrail.async_moderation_hook(
data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type, cache=DualCache()
)
assert classify.call_count == lasso_calls
assert metadata.get("applied_guardrails") == applied

75
tests/unit/proxy/guardrails/guardrail_hooks/test_llm_shield_proxy.py
View file

@ -1,21 +1,28 @@
import asyncio
import json
from types import SimpleNamespace
from collections.abc import Awaitable, Callable, Mapping
from typing import Final
from unittest.mock import AsyncMock
import pytest
from httpx import Request, Response
from pydantic import TypeAdapter
import litellm
from litellm.caching.caching_handler import _PENDING_CACHE_WRITES
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.exceptions import GuardrailRaisedException
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.llm_shield_proxy.llm_shield_proxy import (
GUARDRAIL_NAME,
LLMShieldProxyGuardrail,
)
from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2
from litellm.proxy.utils import ProxyLogging
from litellm.types.extraction import ExtractionAnswer, ExtractionResponse, ExtractionUsage
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.llms.openai import (
FunctionCallArgumentsDeltaEvent,
@ -1887,7 +1894,75 @@ class TestCompletionsRestoration:
assert "".join(chunk.choices[0].text or "" for chunk in out) == "Mail [EMAIL_1"
_SHIELD_TEXTS: Final = TypeAdapter(tuple[str, ...])
def _pseudonymizing_shield(original: str, placeholder: str) -> Callable[..., Awaitable[Response]]:
async def post(url: str, headers: Mapping[str, str], json: Mapping[str, object], timeout: float) -> Response:
texts: Final = _SHIELD_TEXTS.validate_python(json["texts"])
swap: Final = (original, placeholder) if url.endswith("/redact") else (placeholder, original)
return _response({"texts": [text.replace(*swap) for text in texts]})
return post
class TestProxyWiring:
@pytest.mark.asyncio
async def test_extraction_request_is_redacted_and_the_answer_restored(self, monkeypatch: pytest.MonkeyPatch):
guardrail: Final = _guardrail(event_hook=["pre_call", "post_call"])
guardrail.async_handler.post = _pseudonymizing_shield("alice@example.com", "[EMAIL_1]")
monkeypatch.setattr(litellm, "callbacks", [guardrail])
proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache())
caller: Final = UserAPIKeyAuth(api_key="sk-9876", request_route="/v1/extraction")
sent: Final = await proxy_logging.pre_call_hook(
user_api_key_dict=caller,
data={
"model": "m",
"state": "Write to alice@example.com",
"fields": {"email": {"type": "text", "instructions": "The address after alice@example.com"}},
},
call_type="aextraction",
)
reply: Final = ExtractionResponse(
model="m",
answers={"email": ExtractionAnswer(type="text", value="[EMAIL_1]")},
truncated=False,
usage=ExtractionUsage(input_tokens=1, output_tokens=1),
)
restored: Final = await proxy_logging.post_call_success_hook(
data=sent, response=reply, user_api_key_dict=caller
)
assert (sent["state"], sent["fields"]["email"]["instructions"]) == (
"Write to [EMAIL_1]",
"The address after [EMAIL_1]",
)
assert restored.answers["email"].value == "alice@example.com"
@pytest.mark.asyncio
async def test_decision_request_is_redacted_and_keeps_its_routing_tags(self, monkeypatch: pytest.MonkeyPatch):
guardrail: Final = _guardrail()
guardrail.async_handler.post = _pseudonymizing_shield("alice@example.com", "[EMAIL_1]")
monkeypatch.setattr(litellm, "callbacks", [guardrail])
sent: Final = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"),
data={
"model": "m",
"input": "Refund alice@example.com",
"questions": [{"type": "predicate", "instructions": "Is alice@example.com owed money?"}],
"metadata": {"tags": ["prod"]},
},
call_type="adecisions",
)
assert (sent["input"], sent["questions"][0]["instructions"]) == (
"Refund [EMAIL_1]",
"Is [EMAIL_1] owed money?",
)
assert sent[get_metadata_variable_name_from_kwargs(sent)]["tags"] == ["prod"]
def test_dashboard_config_model_is_exposed(self):
"""The guardrail garden reads the provider's fields from `get_config_model`."""
model = LLMShieldProxyGuardrail.get_config_model()

105
tests/unit/proxy/guardrails/guardrail_hooks/test_model_armor.py
View file

@ -3,7 +3,7 @@ import base64
import io
import json
from collections.abc import Iterator, Sequence
from typing import cast
from typing import Final, cast
from unittest.mock import AsyncMock, MagicMock, Mock, patch
import pytest
@ -5291,6 +5291,109 @@ async def test_logging_only_flagged_prompt_still_scans_response():
assert len(flagged) == 2
class _RedactsPhoneNumbers(AsyncHTTPHandler):
def __init__(self) -> None:
pass
async def post(
self,
url: str,
json: dict[str, object] | None = None,
headers: dict[str, str] | None = None,
**kwargs: object,
) -> httpx.Response:
payload: Final = cast(dict[str, dict[str, str]], json)
text: Final = (payload.get("userPromptData") or payload["modelResponseData"])["text"]
redacted: Final = text.replace("555-123-4567", "[REDACTED]")
result: Final = (
_clean_armor_response()
if redacted == text
else {
"sanitizationResult": {
"filterMatchState": "MATCH_FOUND",
"filterResults": {
"sdp": {
"sdpFilterResult": {
"deidentifyResult": {"matchState": "MATCH_FOUND", "data": {"text": redacted}}
}
}
},
}
}
)
return httpx.Response(200, json=result, request=httpx.Request("POST", url))
@pytest.mark.asyncio
@pytest.mark.parametrize("mask_request_content", (True, False))
async def test_apply_guardrail_blocks_a_masking_finding(mask_request_content: bool) -> None:
guardrail: Final = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-block",
event_hook=GuardrailEventHooks.pre_call,
mask_request_content=mask_request_content,
async_handler=_RedactsPhoneNumbers(),
access_token_provider=_async_token_provider,
)
with pytest.raises(HTTPException) as exc_info:
await guardrail.apply_guardrail(
inputs={"texts": ["call me at 555-123-4567", "refund order 4411"]},
request_data={"metadata": {}},
input_type="request",
)
assert exc_info.value.status_code == 400
@pytest.mark.asyncio
async def test_apply_guardrail_returns_clean_text_unchanged() -> None:
inputs: Final = {"texts": ["refund order 4411"]}
guardrail: Final = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-clean",
event_hook=GuardrailEventHooks.pre_call,
async_handler=_RedactsPhoneNumbers(),
access_token_provider=_async_token_provider,
)
assert await guardrail.apply_guardrail(inputs=inputs, request_data={"metadata": {}}, input_type="request") is inputs
class _UnavailableArmor(AsyncHTTPHandler):
def __init__(self) -> None:
pass
async def post(self, url: str, **kwargs: object) -> httpx.Response:
return httpx.Response(503, text="unavailable", request=httpx.Request("POST", url))
@pytest.mark.asyncio
@pytest.mark.parametrize("fail_on_error", (True, False))
async def test_apply_guardrail_honors_fail_on_error_when_model_armor_is_down(fail_on_error: bool) -> None:
guardrail: Final = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-down",
event_hook=GuardrailEventHooks.pre_call,
fail_on_error=fail_on_error,
async_handler=_UnavailableArmor(),
access_token_provider=_async_token_provider,
)
inputs: Final = {"texts": ["refund order 4411"]}
if fail_on_error:
with pytest.raises(ModelArmorAPIError):
await guardrail.apply_guardrail(inputs=inputs, request_data={"metadata": {}}, input_type="request")
else:
assert await guardrail.apply_guardrail(inputs=inputs, request_data={"metadata": {}}, input_type="request") is inputs
@pytest.mark.asyncio
async def test_apply_guardrail_raises_on_flagged_when_not_logging_only():
"""The /guardrails/apply_guardrail endpoint calls apply_guardrail directly; a

43
tests/unit/proxy/guardrails/guardrail_hooks/test_pangea.py
View file

@ -1,9 +1,14 @@
from unittest.mock import AsyncMock, patch
import httpx
from collections.abc import Mapping, Sequence
from typing import Final
import pytest
import respx
from fastapi import HTTPException
from litellm import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.pangea import initialize_guardrail
from litellm.proxy.guardrails.guardrail_hooks.pangea.pangea import (
PangeaGuardrailMissingSecrets,
@ -368,3 +373,41 @@ async def test_pangea_ai_guard_response_transformed(pangea_guardrail):
)
assert response.choices[0]["message"]["content"] == "Yes, here is an SSN: <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: respx.MockRouter,
call_type: str,
request_fields: Mapping[str, object],
pangea_calls: int,
applied: Sequence[str] | None,
) -> None:
guard: Final = respx_mock.post("https://ai-guard.aws.us.pangea.cloud/v1beta/guard").respond(
json={"result": {"blocked": False, "transformed": False}}
)
guardrail: Final = PangeaHandler(
mode="pre_call", guardrail_name="pangea-ai-guard", api_key="pts_pangeatokenid", default_on=True
)
metadata: Final[dict[str, object]] = {}
await guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(),
cache=DualCache(),
data={**request_fields, "metadata": metadata},
call_type=call_type,
)
assert guard.call_count == pangea_calls
assert metadata.get("applied_guardrails") == applied

59
tests/unit/proxy/guardrails/test_pillar_guardrails.py
View file

@ -16,7 +16,10 @@ from unittest.mock import Mock, patch
import json
from urllib.parse import unquote
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Final
import pytest
import respx
from fastapi.exceptions import HTTPException
from httpx import Request, Response
@ -24,6 +27,7 @@ from httpx import Request, Response
import litellm
from litellm import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.types.utils import CallTypesLiteral
from litellm.proxy.common_utils.callback_utils import get_logging_caching_headers
from litellm.proxy.guardrails.guardrail_hooks.pillar import (
PillarGuardrail,
@ -1513,3 +1517,58 @@ async def test_mcp_call_masking(
if __name__ == "__main__":
# Run the tests
pytest.main([__file__, "-v"])
_INPUT_HOOKS = [
pytest.param(
"pre_call",
lambda guardrail, data, call_type: guardrail.async_pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(), cache=DualCache(), data=data, call_type=call_type
),
id="pre_call",
),
pytest.param(
"during_call",
lambda guardrail, data, call_type: guardrail.async_moderation_hook(
data=data, user_api_key_dict=UserAPIKeyAuth(), call_type=call_type
),
id="during_call",
),
]
@pytest.mark.asyncio
@pytest.mark.usefixtures("httpx_transport")
@pytest.mark.parametrize(("event_hook", "run_hook"), _INPUT_HOOKS)
@pytest.mark.parametrize(
("request_fields", "call_type", "applied"),
[
pytest.param({"messages": [{"role": "user", "content": "hi"}]}, "completion", ["pillar-test"], id="messages"),
pytest.param({"state": "hi"}, "aextraction", None, id="no-messages-on-extraction"),
pytest.param({"state": "hi"}, "completion", ["pillar-test"], id="no-messages-on-chat"),
],
)
async def test_input_hooks_list_the_guardrail_only_when_they_scanned_text(
respx_mock: respx.MockRouter,
event_hook: str,
run_hook: Callable[[PillarGuardrail, Mapping[str, object], CallTypesLiteral], Awaitable[object]],
request_fields: Mapping[str, object],
call_type: CallTypesLiteral,
applied: Sequence[str] | None,
) -> None:
respx_mock.post(url__startswith="https://api.pillar.security/api/v1/protect").respond(
json={"session_id": "s", "flagged": False, "scanners": {}, "evidence": []}
)
guardrail: Final = PillarGuardrail(
guardrail_name="pillar-test",
api_key="test-pillar-key",
api_base="https://api.pillar.security",
on_flagged_action="block",
event_hook=event_hook,
default_on=True,
)
metadata: Final[dict[str, object]] = {}
await run_hook(guardrail, {**request_fields, "metadata": metadata}, call_type)
assert metadata.get("applied_guardrails") == applied

34
tests/unit/proxy/policy_engine/test_pipeline_executor.py
View file

@ -10,6 +10,8 @@ import pickle
from typing import Literal
from unittest.mock import MagicMock
from collections.abc import Mapping
from typing import Final
import pytest
import litellm
@ -30,7 +32,7 @@ from litellm.types.proxy.policy_engine.pipeline_types import (
GuardrailPipeline,
PipelineStep,
)
from litellm.types.utils import CallTypesLiteral
from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs
try:
from fastapi.exceptions import HTTPException
@ -1137,8 +1139,16 @@ class _PolicyOptOutGuardrail(CustomGuardrail):
def __init__(self):
super().__init__(guardrail_name="policy-opt-out", default_on=True)
self.native_pre_call_ran = False
self.scanned: tuple[str, ...] = ()
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: Mapping[str, object],
input_type: Literal["request", "response"],
logging_obj: object = None,
) -> GenericGuardrailAPIInputs:
self.scanned = (*self.scanned, *inputs.get("texts", []))
return inputs
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
@ -1164,6 +1174,26 @@ async def test_pipeline_step_keeps_native_hook_when_opted_out(monkeypatch):
assert "guardrail_to_apply" not in data
@pytest.mark.asyncio
async def test_pipeline_step_scans_extraction_state_through_apply_guardrail_when_opted_out(
monkeypatch: pytest.MonkeyPatch,
) -> None:
guardrail: Final = _PolicyOptOutGuardrail()
monkeypatch.setattr(litellm, "callbacks", [guardrail])
outcome, _, _, _ = await PipelineExecutor._run_step(
step=PipelineStep(guardrail="policy-opt-out", on_fail="block", on_pass="allow"),
mode="pre_call",
data={"model": "m", "state": "invoice from vendor", "fields": {"total": {"type": "number"}}},
user_api_key_dict=UserAPIKeyAuth(),
call_type="aextraction",
)
assert outcome == "pass"
assert guardrail.native_pre_call_ran is False
assert guardrail.scanned == ("invoice from vendor",)
class _TextReturningGuardrail(CustomGuardrail):
def __init__(self, returned_texts):
super().__init__(guardrail_name="masker", event_hook="post_call", default_on=True)

447
tests/unit/proxy/test_proxy_logging_hook_detection.py
View file

@ -1,14 +1,28 @@
from collections.abc import Mapping
from itertools import chain
from typing import Final, Literal
import httpx
import pytest
from fastapi import HTTPException
from pydantic import BaseModel, ConfigDict
import litellm
from litellm.caching import DualCache
from litellm.exceptions import SensitiveDataRouteException
from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.guardrails.guardrail_hooks.model_armor.model_armor import ModelArmorGuardrail
from litellm.proxy.guardrails.guardrail_hooks.tool_permission import ToolPermissionGuardrail
from litellm.proxy.utils import ProxyLogging
from litellm.types.decisions import OpenAIDecisionResponse
from litellm.types.extraction import ExtractionResponse
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import CallTypesLiteral
from litellm.types.utils import CallTypesLiteral, GenericGuardrailAPIInputs
def test_has_post_call_response_headers_callbacks_ignores_empty_callbacks(
@ -1060,3 +1074,434 @@ async def test_parallel_post_call_guardrails_keep_native_hook_when_opted_out(mon
assert opted_out.native_hooks_ran == ["post_call"]
assert routed.native_hooks_ran == []
_SECRET: Final = "sk-live-1234"
class _Metadata(BaseModel):
model_config = ConfigDict(frozen=True)
applied_guardrails: tuple[str, ...] = ()
class _RoutedMetadata(BaseModel):
model_config = ConfigDict(frozen=True)
tags: tuple[str, ...] = ()
requester_ip_address: str | None = None
user_api_key_key_alias: str | None = None
def _applied_guardrails(data: Mapping[str, object]) -> tuple[str, ...]:
return tuple(
chain.from_iterable(
_Metadata.model_validate(data.get(bucket) or {}).applied_guardrails
for bucket in ("metadata", "litellm_metadata")
)
)
class _MasksSecretWithNativeHooks(_KeepsNativeHooks):
def __init__(self, **kwargs: object) -> None:
super().__init__(**kwargs)
self.scanned: list[tuple[str, tuple[str, ...]]] = []
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:
texts: Final = tuple(inputs.get("texts") or ())
self.scanned.append((input_type, texts))
return GenericGuardrailAPIInputs(texts=[text.replace(_SECRET, "[MASKED]") for text in texts])
@pytest.mark.asyncio
@pytest.mark.parametrize(
("call_type", "data", "expected"),
(
(
"aextraction",
{"model": "m", "state": f"paid with {_SECRET}", "fields": {"total": {"type": "number"}}},
{"state": "paid with [MASKED]"},
),
(
"asystemone",
{
"model": "m",
"state": {"ticket": f"my key is {_SECRET}"},
"questions": {"leak": {"type": "noul", "instructions": f"Is {_SECRET} a key?"}},
},
{
"state": {"ticket": "my key is [MASKED]"},
"questions": {"leak": {"type": "noul", "instructions": "Is [MASKED] a key?"}},
},
),
(
"adecisions",
{
"model": "m",
"input": f"my key is {_SECRET}",
"questions": [{"type": "predicate", "instructions": f"Is {_SECRET} a key?"}],
},
{
"input": "my key is [MASKED]",
"questions": [{"type": "predicate", "instructions": "Is [MASKED] a key?"}],
},
),
),
)
async def test_pre_call_hook_scans_decision_and_extraction_text_with_a_native_hook_guardrail(
monkeypatch: pytest.MonkeyPatch,
call_type: CallTypesLiteral,
data: dict[str, object],
expected: dict[str, object],
) -> None:
guardrail: Final = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"),
data=data,
call_type=call_type,
)
assert {key: result[key] for key in expected} == expected
assert guardrail.native_hooks_ran == []
assert _applied_guardrails(result) == ("keeps_native",)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("call_type", "request_fields"),
(
("aextraction", {"state": f"paid with {_SECRET}", "fields": {"total": {"type": "number"}}}),
("asystemone", {"state": {"ticket": _SECRET}, "questions": {"leak": {"type": "noul", "instructions": "Key?"}}}),
("adecisions", {"input": _SECRET, "questions": [{"type": "predicate", "instructions": "Is it a key?"}]}),
),
)
async def test_pre_call_hook_keeps_request_tags_where_the_router_reads_them(
monkeypatch: pytest.MonkeyPatch, call_type: CallTypesLiteral, request_fields: Mapping[str, object]
) -> None:
guardrail: Final = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", key_alias="billing"),
data={"model": "m", **request_fields, "metadata": {"tags": ["prod"], "requester_ip_address": "10.0.0.7"}},
call_type=call_type,
)
routed_metadata: Final = _RoutedMetadata.model_validate(result[get_metadata_variable_name_from_kwargs(result)])
assert routed_metadata == _RoutedMetadata(
tags=("prod",), requester_ip_address="10.0.0.7", user_api_key_key_alias="billing"
)
assert _applied_guardrails(result) == ("keeps_native",)
class _ModelArmorAnswering(AsyncHTTPHandler):
def __init__(self, status_code: int) -> None:
self.status_code = status_code
async def post(self, url: str, **kwargs: object) -> httpx.Response:
clean: Final = {"sanitizationResult": {"filterMatchState": "NO_MATCH_FOUND", "filterResults": {}}}
return httpx.Response(self.status_code, json=clean, request=httpx.Request("POST", url))
async def _model_armor_token() -> tuple[str, str]:
return ("test-token", "test-project")
@pytest.mark.asyncio
@pytest.mark.parametrize(("status_code", "applied"), ((200, ("model-armor",)), (503, ())))
@pytest.mark.parametrize(
("call_type", "request_fields"),
(
("aextraction", {"state": "refund order 4411", "fields": {"total": {"type": "number"}}}),
("asystemone", {"state": {"ticket": "refund order 4411"}, "questions": {"q": {"type": "noul"}}}),
("adecisions", {"input": "refund order 4411", "questions": [{"type": "predicate", "instructions": "Refund?"}]}),
),
)
async def test_pre_call_hook_lists_model_armor_only_when_it_answered(
monkeypatch: pytest.MonkeyPatch,
call_type: CallTypesLiteral,
request_fields: Mapping[str, object],
status_code: int,
applied: tuple[str, ...],
) -> None:
guardrail: Final = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
guardrail_name="model-armor",
event_hook=GuardrailEventHooks.pre_call,
default_on=True,
fail_on_error=False,
async_handler=_ModelArmorAnswering(status_code),
access_token_provider=_model_armor_token,
)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"),
data={"model": "m", **request_fields},
call_type=call_type,
)
assert _applied_guardrails(result) == applied
@pytest.mark.asyncio
async def test_pre_call_hook_does_not_list_a_guardrail_for_a_systemone_request_without_text(
monkeypatch: pytest.MonkeyPatch,
) -> None:
guardrail: Final = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"),
data={"model": "m", "state": {}, "questions": {"routed": {"type": "choice", "criteria": {"refund": None}}}},
call_type="asystemone",
)
assert guardrail.scanned == []
assert _applied_guardrails(result) == ()
@pytest.mark.asyncio
async def test_during_call_hook_scans_extraction_state_with_a_native_hook_guardrail(
monkeypatch: pytest.MonkeyPatch,
) -> None:
guardrail: Final = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.during_call, default_on=True)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
await ProxyLogging(user_api_key_cache=DualCache()).during_call_hook(
data={"model": "m", "state": f"paid with {_SECRET}", "fields": {"total": {"type": "number"}}},
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"),
call_type="aextraction",
)
assert guardrail.native_hooks_ran == []
assert guardrail.scanned == [("request", (f"paid with {_SECRET}",))]
@pytest.mark.asyncio
async def test_pre_call_hook_keeps_native_hook_guardrail_on_its_own_hook_for_chat(
monkeypatch: pytest.MonkeyPatch,
) -> None:
guardrail: Final = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"),
data={"model": "m", "messages": [{"role": "user", "content": f"my key is {_SECRET}"}]},
call_type="acompletion",
)
assert guardrail.native_hooks_ran == ["pre_call"]
assert guardrail.scanned == []
@pytest.mark.asyncio
@pytest.mark.parametrize("run_in_parallel", (False, True))
async def test_post_call_success_hook_scans_extraction_answers_with_a_native_hook_guardrail(
monkeypatch: pytest.MonkeyPatch, run_in_parallel: bool
) -> None:
guardrail: Final = _MasksSecretWithNativeHooks(
event_hook=GuardrailEventHooks.post_call, default_on=True, run_in_parallel=run_in_parallel
)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
response: Final = ExtractionResponse.model_validate(
{
"model": "m",
"answers": {"key": {"type": "text", "value": _SECRET}},
"truncated": False,
"usage": {"input_tokens": 5, "output_tokens": 0},
}
)
await ProxyLogging(user_api_key_cache=DualCache()).post_call_success_hook(
data={"model": "m", "state": "s", "fields": {"key": {"type": "text"}}},
response=response,
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route="/v1/extraction"),
)
assert guardrail.native_hooks_ran == []
assert guardrail.scanned == [("response", (_SECRET,))]
class _ReadsNoDecisionText(CustomGuardrail):
def __init__(
self,
guardrail_name: str,
default_on: bool = False,
event_hook: GuardrailEventHooks = GuardrailEventHooks.pre_call,
) -> None:
super().__init__(guardrail_name=guardrail_name, event_hook=event_hook, default_on=default_on)
self.native_hooks_ran: list[str] = []
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict[str, object],
call_type: CallTypesLiteral,
) -> dict[str, object]:
self.native_hooks_ran.append(call_type)
return data
_REQUEST_WITH_EVERY_TEXT_FIELD: Final = {
"model": "m",
"state": "s",
"input": "hi",
"questions": [{"type": "predicate", "instructions": "Is it a greeting?"}],
"messages": [{"role": "user", "content": "hi"}],
}
def _warning(call_type: str) -> str:
return (
"Guardrails requested-native, default-native do not implement apply_guardrail, "
f"so on {call_type} requests they scan only the fields their own hook reads"
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("call_type", "warning"),
(("aextraction", _warning("aextraction")), ("adecisions", _warning("adecisions")), ("acompletion", None)),
)
async def test_requested_user_guardrail_without_apply_guardrail_stays_listed_and_is_warned_about(
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
call_type: CallTypesLiteral,
warning: str | None,
) -> None:
requested: Final = _ReadsNoDecisionText("requested-native")
default_on: Final = _ReadsNoDecisionText("default-native", default_on=True)
post_call_only: Final = _ReadsNoDecisionText(
"post-call-native", default_on=True, event_hook=GuardrailEventHooks.post_call
)
scanner: Final = _MasksSecretWithNativeHooks(event_hook=GuardrailEventHooks.pre_call, default_on=True)
monkeypatch.setattr(litellm, "callbacks", [requested, default_on, post_call_only, scanner])
with caplog.at_level("WARNING"):
result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"),
data={**_REQUEST_WITH_EVERY_TEXT_FIELD, "metadata": {"guardrails": ["requested-native"]}},
call_type=call_type,
)
assert requested.native_hooks_ran == [call_type]
assert "requested-native" in _applied_guardrails(result)
warnings: Final = [
record.getMessage() for record in caplog.records if "do not implement apply_guardrail" in record.getMessage()
]
assert warnings == ([] if warning is None else [warning])
def _tool_permission_guardrail() -> ToolPermissionGuardrail:
return ToolPermissionGuardrail(
guardrail_name="tool-permission",
event_hook=GuardrailEventHooks.pre_call,
rules=[],
default_action="allow",
on_disallowed_action="block",
)
@pytest.mark.asyncio
@pytest.mark.parametrize(("call_type", "listed"), (("aextraction", False), ("adecisions", False), ("acompletion", True)))
async def test_requested_built_in_guardrail_is_listed_only_where_its_hook_reads_the_text(
monkeypatch: pytest.MonkeyPatch, call_type: CallTypesLiteral, listed: bool
) -> None:
monkeypatch.setattr(litellm, "callbacks", [_tool_permission_guardrail()])
result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"),
data={**_REQUEST_WITH_EVERY_TEXT_FIELD, "metadata": {"guardrails": ["tool-permission"]}},
call_type=call_type,
)
assert ("tool-permission" in _applied_guardrails(result)) is listed
class _ReroutesToPrivateModel(CustomGuardrail):
def __init__(self) -> None:
super().__init__(
guardrail_name="reroute", event_hook=GuardrailEventHooks.pre_call, default_on=True, run_in_parallel=True
)
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
cache: DualCache,
data: dict[str, object],
call_type: CallTypesLiteral,
) -> dict[str, object]:
raise SensitiveDataRouteException(
route_to_model="private-model", session_id="s1", guardrail_name="reroute", sticky_session_routing=False
)
@pytest.mark.asyncio
@pytest.mark.parametrize(("call_type", "listed"), (("aextraction", False), ("acompletion", True)))
async def test_rerouted_request_lists_a_requested_built_in_guardrail_only_where_it_reads_the_text(
monkeypatch: pytest.MonkeyPatch, call_type: CallTypesLiteral, listed: bool
) -> None:
monkeypatch.setattr(litellm, "callbacks", [_tool_permission_guardrail(), _ReroutesToPrivateModel()])
result: Final[dict[str, object]] = await ProxyLogging(user_api_key_cache=DualCache()).pre_call_hook(
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876"),
data={**_REQUEST_WITH_EVERY_TEXT_FIELD, "metadata": {"guardrails": ["tool-permission"]}},
call_type=call_type,
)
assert result["model"] == "private-model"
assert ("tool-permission" in _applied_guardrails(result)) is listed
class _RecordsScans(_AppliesGuardrail):
def __init__(self, **kwargs: object) -> None:
super().__init__(**kwargs)
self.scanned: list[str] = []
async def apply_guardrail(
self,
inputs: GenericGuardrailAPIInputs,
request_data: dict[str, object],
input_type: Literal["request", "response"],
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:
self.scanned.append(input_type)
return inputs
@pytest.mark.asyncio
@pytest.mark.parametrize("route", ("/v1/decisions", "/v1/systemone"))
async def test_post_call_guardrail_leaves_decision_answers_unscanned_and_unlisted(
monkeypatch: pytest.MonkeyPatch, route: str
) -> None:
guardrail: Final = _RecordsScans(event_hook=GuardrailEventHooks.post_call, default_on=True)
monkeypatch.setattr(litellm, "callbacks", [guardrail])
response: Final = OpenAIDecisionResponse.model_validate(
{
"model": "m",
"answers": [{"type": "refusal", "name": "q"}],
"usage": {"input_tokens": 1, "output_tokens": 0, "total_tokens": 1},
}
)
data: Final[dict[str, object]] = {"model": "m", "metadata": {}}
result: Final = await ProxyLogging(user_api_key_cache=DualCache()).post_call_success_hook(
data=data,
response=response,
user_api_key_dict=UserAPIKeyAuth(api_key="sk-9876", request_route=route),
)
assert result is response
assert guardrail.scanned == []
assert _applied_guardrails(data) == ()

8
tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py
View file

@ -278,7 +278,7 @@ def test_process_guardrail_metadata_calls_header_helper(proxy_logging, monkeypat
monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add)
data = {"metadata": {"guardrails": ["g1", "g2"]}}
proxy_logging._process_guardrail_metadata(data)
proxy_logging._process_guardrail_metadata(data, "completion")
snapshot = {
"call_count": len(calls),
"first_name": calls[0]["name"],
@ -303,7 +303,7 @@ def test_process_guardrail_metadata_skips_already_applied(proxy_logging, monkeyp
monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add)
data = {"metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g1"]}}
proxy_logging._process_guardrail_metadata(data)
proxy_logging._process_guardrail_metadata(data, "completion")
assert calls == ["g2"]
@ -315,12 +315,12 @@ def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging, monkeypat
"add_guardrail_to_applied_guardrails_header",
MagicMock(side_effect=AssertionError("should not be called")),
)
proxy_logging._process_guardrail_metadata({})
proxy_logging._process_guardrail_metadata({}, "completion")
def test_process_guardrail_metadata_invalid_data_raises(proxy_logging):
with pytest.raises(AttributeError):
proxy_logging._process_guardrail_metadata(None) # type: ignore[arg-type]
proxy_logging._process_guardrail_metadata(None, "completion") # type: ignore[arg-type]
# ---------------------------------------------------------------------------

2
tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py
View file

@ -170,7 +170,7 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox
proxy_logging.slack_alerting_instance = MagicMock(alerting=None)
invoked = {}
def fake_process(d):
def fake_process(d, call_type):
invoked["data"] = d
proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment]