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