diff --git a/enterprise/litellm_enterprise/integrations/custom_guardrail.py b/enterprise/litellm_enterprise/integrations/custom_guardrail.py index 58c24219b47..56452bf3e85 100644 --- a/enterprise/litellm_enterprise/integrations/custom_guardrail.py +++ b/enterprise/litellm_enterprise/integrations/custom_guardrail.py @@ -3,6 +3,11 @@ from typing import List, Optional, Union from litellm.types.guardrails import GuardrailEventHooks, Mode +def _names_a_mode_tag(tag: object, mode: Mode) -> bool: + """Whether a request tag, which a client's body can make anything, is one of ``mode``'s tags""" + return isinstance(tag, str) and tag in mode.tags + + class EnterpriseCustomGuardrailHelper: @staticmethod def _should_run_if_mode_by_tag( @@ -55,7 +60,7 @@ class EnterpriseCustomGuardrailHelper: matched_mode = None if request_tags: for tag in request_tags: - if tag in event_hook.tags: + if _names_a_mode_tag(tag, event_hook): matched_mode = event_hook.tags[tag] break diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 68876812c3e..71078eefd31 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -211,46 +211,34 @@ def guardrail_request_data_with_streaming( } -_GUARDRAIL_ENTRY_WITH_SETTINGS: Final = TypeAdapter(Mapping[str, object]) - - def merged_guardrail_entries(*entry_lists: Iterable[object]) -> list[object]: """Merges guardrail lists so each guardrail is named once: the first ``{name: settings}`` entry naming it wins over any plain name in either order, else its first plain name stays""" entries: Final = tuple(chain.from_iterable(entry_lists)) - settings_names: Final = tuple(_settings_entry_names(entry) for entry in entries) - named_with_settings: Final = frozenset[str]().union(*settings_names) - return [ - entry - for index, entry in enumerate(entries) - if not _superseded_guardrail_entry(entries, settings_names, named_with_settings, index) - ] + if all(isinstance(entry, str) for entry in entries): + return list(dict.fromkeys(entries)) + named_with_settings: Final = frozenset( + chain.from_iterable(guardrail_entry_names(entry) for entry in entries if isinstance(entry, dict)) + ) + seen: Final[set[str]] = set() # mutable-ok: local to this walk + merged: Final[list[object]] = [] # mutable-ok: local to this walk + for entry in entries: + names = guardrail_entry_names(entry) + if (names and seen.issuperset(names)) or (isinstance(entry, str) and entry in named_with_settings): + continue + merged.append(entry) + seen.update(names) + return merged def guardrail_entry_names(entry: object) -> tuple[str, ...]: if isinstance(entry, str): return (entry,) if isinstance(entry, dict): - return tuple(_GUARDRAIL_ENTRY_WITH_SETTINGS.validate_python(entry)) + return tuple(cast("Mapping[str, object]", entry)) # cast-ok: a guardrail settings entry is keyed by name return () -def _settings_entry_names(entry: object) -> frozenset[str]: - return frozenset(guardrail_entry_names(entry)) if isinstance(entry, dict) else frozenset() - - -def _superseded_guardrail_entry( - entries: tuple[object, ...], - settings_names: tuple[frozenset[str], ...], - named_with_settings: frozenset[str], - index: int, -) -> bool: - entry: Final = entries[index] - if isinstance(entry, str): - return entry in named_with_settings or entry in entries[:index] - return bool(settings_names[index]) and settings_names[index] <= frozenset[str]().union(*settings_names[:index]) - - def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = None) -> bool: if event_type in _REALTIME_STREAMING_HOOKS: return True diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ce57a9bd4a5..5a57e86617a 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2586,7 +2586,10 @@ class Logging(LiteLLMLoggingBaseClass): prices: Final = declared_cost_map_pricing(model_info["key"], model_info["key"], payload["custom_llm_provider"]) return any(value is not None and ("cost" in key or "pricing" in key) for key, value in prices.items()) - def prepare_response_cost(self, response: object, now: datetime.datetime) -> PreparedCost | None: + def prepare_response_cost( + self, response: object, now: datetime.datetime, settle_unknown_cost_at_zero: bool = False + ) -> PreparedCost | None: + """The response's cost, ready to finalize, or ``None`` when it has no cost result, unless that settles at 0""" if self.prepared_cost is not None: return self.prepared_cost recorded_cache_hit: Final[object] = self.model_call_details.get("cache_hit") @@ -2612,8 +2615,10 @@ class Logging(LiteLLMLoggingBaseClass): publish_logging_payload=False, ) payload: Final = prepared[3] - raw_cost: Final[object] = self.model_call_details.get("response_cost") - if payload is None or not isinstance(raw_cost, (int, float)): + recorded_cost: Final[object] = self.model_call_details.get("response_cost") + unknown_cost: Final = 0.0 if settle_unknown_cost_at_zero else None + raw_cost: Final = recorded_cost if isinstance(recorded_cost, (int, float)) else unknown_cost + if payload is None or raw_cost is None: return None recorded_usage: Final[object] = self.model_call_details.get("combined_usage_object") response_usage: Final[object] = getattr(prepared[2], "usage", None) diff --git a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py index 3514a5e6e5c..7dacd755a7c 100644 --- a/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py +++ b/litellm/llms/bedrock/passthrough/guardrail_translation/handler.py @@ -1516,6 +1516,7 @@ class BedrockPassthroughGuardrailHandler(BaseTranslation): litellm_logging_obj=litellm_logging_obj, user_api_key_dict=user_api_key_dict, request_data=request_data, + read_whole_response=True, ) if not isinstance(response, dict): diff --git a/litellm/llms/pass_through/guardrail_translation/README.md b/litellm/llms/pass_through/guardrail_translation/README.md index db4c0704e75..74bde8e809c 100644 --- a/litellm/llms/pass_through/guardrail_translation/README.md +++ b/litellm/llms/pass_through/guardrail_translation/README.md @@ -24,7 +24,7 @@ litellm/proxy/pass_through_endpoints/ The `PassThroughEndpointHandler` enables guardrails to run on passthrough endpoint requests by: 1. **Field Targeting**: Extracts specific fields from the request/response using JSONPath expressions configured in `request_fields` / `response_fields` -2. **Full Payload Fallback**: If no field targeting is configured, processes the entire payload +2. **Full Payload Fallback**: If no field targeting is configured, processes the entire request payload. A response is read whole, field names and integers included, except a provider's reply, where the ids, timestamps and usage counts around the text are left out 3. **Config Access**: Uses `get_passthrough_guardrails_config()` / `set_passthrough_guardrails_config()` helpers to access the passthrough guardrails configuration stored in request metadata ## Example Config diff --git a/litellm/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index dea8bcb73a4..0af14c23654 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -6,14 +6,15 @@ It uses the field targeting configuration from litellm_logging_obj to extract specific fields for guardrail processing. """ -from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping, Sequence +from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Iterator, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Optional, Protocol, runtime_checkable from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import without_server_streaming_classification -from litellm.litellm_core_utils.core_helpers import has_attributes +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, has_attributes from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, StreamingScanKey, @@ -29,6 +30,95 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +# The fields of a provider's JSON reply that name or encode something rather than carry its text, by the field they sit +# under, and the fields that hold what the model wrote, which are read whole. +_REPLY_TOP_LEVEL_FIELDS: Final = frozenset( + ( + "finish_reason", + "id", + "model", + "modelVersion", + "object", + "responseId", + "role", + "service_tier", + "status", + "stop_reason", + "stop_sequence", + "system_fingerprint", + "type", + ) +) +_REPLY_NESTED_FIELDS: Final = MappingProxyType( + { + "candidates": frozenset(("finishReason",)), + "choices": frozenset(("finish_reason",)), + "content": frozenset(("id", "role", "type")), + "data": frozenset(("b64_json", "embedding", "object")), + "message": frozenset(("role",)), + "output": frozenset(("id", "role", "status", "type")), + "parts": frozenset(("thoughtSignature",)), + "predictions": frozenset(("bytesBase64Encoded",)), + "tool_calls": frozenset(("id", "type")), + } +) +_REPLY_MODEL_OUTPUT_FIELDS: Final = frozenset(("args", "arguments", "input", "parameters")) +_REPLY_LIST_FIELDS: Final = ("candidates", "choices", "output", "predictions") +_IMAGE_REPLY_FIELDS: Final = frozenset(("created", "data")) +_NO_FIELDS: Final[frozenset[str]] = frozenset() + + +def _is_provider_reply(response: Mapping[str, object]) -> bool: + """Whether a JSON response has the shape of a provider's reply, by the type of the field that names it, since a + custom API's JSON can use the same field names.""" + return ( + any(isinstance(response.get(field), list) for field in _REPLY_LIST_FIELDS) + or isinstance(response.get("object"), str) + or (response.get("type") == "message" and "stop_reason" in response) + or ("finish_reason" in response and "message" in response) + or _IMAGE_REPLY_FIELDS.issubset(response) + ) + + +def _reply_field_position(name: str, skipped: frozenset[str] | None) -> frozenset[str] | None: + if skipped is None or name in _REPLY_MODEL_OUTPUT_FIELDS: + return None + return _REPLY_NESTED_FIELDS.get(name, _NO_FIELDS) + + +def _reply_entries( + fields: Mapping[str, object], skipped: frozenset[str] | None +) -> Iterator[tuple[object, frozenset[str] | None]]: + for name, item in fields.items(): + if skipped is None: + yield name, None + if skipped is None or name not in skipped: + yield item, _reply_field_position(name, skipped) + + +def _response_text(response: Mapping[str, object]) -> str: + """The text a guardrail reads from a JSON response. In a provider's reply that is its strings without the ids and + encoded data around them, so a timestamp or a usage count doesn't read as a phone or card number. What the model + wrote into a tool call, and any other API's JSON, is read whole, field names and integers included. Floats are + scores and vectors.""" + top_level: Final = _REPLY_TOP_LEVEL_FIELDS if _is_provider_reply(response) else None + pending: Final[list[tuple[object, frozenset[str] | None]]] = [(response, top_level)] # mutable-ok: stack + texts: Final[list[str]] = [] # mutable-ok: local to this walk + while pending: + value, skipped = pending.pop() + fields = as_str_mapping(value) + items = as_object_list(value) + if isinstance(value, str): + texts.append(value) + elif skipped is None and isinstance(value, int) and not isinstance(value, bool): + texts.append(str(value)) + elif fields is not None: + pending.extend(reversed(tuple(_reply_entries(fields, skipped)))) + elif items is not None: + pending.extend((item, skipped) for item in reversed(items)) + return "\n".join(text for text in texts if text) + + class PassThroughEndpointHandler(BaseTranslation): """ Handler for processing passthrough endpoint requests with guardrails. @@ -143,6 +233,7 @@ class PassThroughEndpointHandler(BaseTranslation): litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None, user_api_key_dict: Optional["UserAPIKeyAuth"] = None, request_data: dict | None = None, + read_whole_response: bool = False, ) -> object: """ Process output response by applying guardrails to targeted fields. @@ -152,7 +243,10 @@ class PassThroughEndpointHandler(BaseTranslation): guardrail_to_apply: The guardrail instance to apply litellm_logging_obj: Optional logging object user_api_key_dict: User API key metadata to pass to guardrails + read_whole_response: Without targeted fields, read the response as JSON, ids and usage included, as a call + with no logging object does, such as a logging_only scan """ + response_fields: Final = as_str_mapping(response) if not isinstance(response, dict): verbose_proxy_logger.debug("PassThroughEndpointHandler: Response is not a dict, skipping") return response @@ -168,7 +262,11 @@ class PassThroughEndpointHandler(BaseTranslation): field_expressions: Final = settings.response_fields if settings else None # Extract text to check - text_to_check: Final = self._extract_text_for_guardrail(response, field_expressions) + text_to_check: Final = ( + self._extract_text_for_guardrail(response, field_expressions) + if field_expressions or read_whole_response or litellm_logging_obj is None + else _response_text(response_fields or {}) + ) if not text_to_check: return response diff --git a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py index 7e052796545..5e47ec88a3e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py +++ b/litellm/proxy/guardrails/guardrail_hooks/generic_guardrail_api/generic_guardrail_api.py @@ -365,7 +365,10 @@ class GenericGuardrailAPI(CustomGuardrail): top_level_metadata: Final = request_data.get("metadata", {}) # Merge both sources, preferring litellm_metadata if both exist - metadata_dict: Final = {**top_level_metadata, **litellm_metadata} + metadata_dict: Final = { + **(top_level_metadata if isinstance(top_level_metadata, dict) else {}), + **(litellm_metadata if isinstance(litellm_metadata, dict) else {}), + } if not metadata_dict: return result_metadata @@ -375,12 +378,13 @@ class GenericGuardrailAPI(CustomGuardrail): # Fields in metadata are already prefixed with 'user_api_key_' for field_name in GenericGuardrailAPIMetadata.__annotations__: value = metadata_dict.get(field_name) - if value is not None: + if isinstance(value, str): result_metadata[field_name] = value # handle user_api_key_token = user_api_key_hash - if metadata_dict.get("user_api_key_token") is not None: - result_metadata["user_api_key_hash"] = metadata_dict.get("user_api_key_token") + token: Final = metadata_dict.get("user_api_key_token") + if isinstance(token, str): + result_metadata["user_api_key_hash"] = token verbose_proxy_logger.debug( "Generic Guardrail API: Extracted user metadata: %s", @@ -389,6 +393,14 @@ class GenericGuardrailAPI(CustomGuardrail): return result_metadata + def _dynamic_request_body_params(self, request_body: object) -> Mapping[str, object]: + """The ``extra_body`` a request's ``body`` asks this guardrail to send. A custom endpoint's client can send a + ``body`` field of any shape, which asks for nothing unless it names this guardrail's params in the usual form.""" + try: + return as_str_mapping(self.get_guardrail_dynamic_request_body_params(request_body)) or {} # pyright: ignore[reportArgumentType, reportUnknownArgumentType] # a client's JSON, validated by the call + except (TypeError, ValidationError): + return {} + def _fail_open_passthrough( self, *, @@ -558,16 +570,16 @@ class GenericGuardrailAPI(CustomGuardrail): tools: Final = inputs.get("tools") structured_messages: Final = inputs.get("structured_messages") tool_calls: Final = inputs.get("tool_calls") - model: Final = inputs.get("model") + model: Final[object] = inputs.get("model") # Use provided request_data or create an empty dict if request_data is None: request_data = {} - request_body: Final = request_data.get("body") or {} + request_body: Final[object] = request_data.get("body") # Merge additional provider specific params from config and dynamic params - dynamic_params: Final[Mapping[str, object]] = self.get_guardrail_dynamic_request_body_params(request_body) or {} + dynamic_params: Final = self._dynamic_request_body_params(request_body) additional_params: Final = {**self.additional_provider_specific_params, **dynamic_params} # Extract user API key metadata @@ -594,7 +606,7 @@ class GenericGuardrailAPI(CustomGuardrail): tool_calls=tool_calls, additional_provider_specific_params=additional_params, input_type=input_type, - model=model, + model=model if isinstance(model, str) else None, ) headers: Final = self._build_request_headers() diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index cb1b223fb47..f93204f1bc9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -15,6 +15,7 @@ from typing import Final, Literal import httpx from fastapi import HTTPException +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_proxy_logger @@ -22,6 +23,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, @@ -45,6 +47,20 @@ INPUT_POSITIONING_MAP: Final = { Role.ASSISTANT.value: 2, } +_CONTENT_PARTS: Final = TypeAdapter(list[object]) + + +def _input_item_text(item: object) -> str: + """The text of one entry of a list ``input``: a string, or a Responses message's string or parts ``content``.""" + if isinstance(item, str): + return item + content: Final = (as_str_mapping(item) or {}).get("content") + if isinstance(content, str): + return content + parts: Final = _CONTENT_PARTS.validate_python(content) if isinstance(content, list) else () + part_texts: Final = ((as_str_mapping(part) or {}).get("text") for part in parts) + return "\n".join(text for text in part_texts if isinstance(text, str)) + class lakeraAI_Moderation(CustomGuardrail): @classmethod @@ -219,7 +235,11 @@ class lakeraAI_Moderation(CustomGuardrail): } ) elif "input" in data and isinstance(data["input"], list): - text = "\n".join(data["input"]) + item_texts: Final = tuple(map(_input_item_text, _CONTENT_PARTS.validate_python(data["input"]))) + text = "\n".join(item_texts) if any(item_texts) else "" + if not text: + verbose_proxy_logger.debug("Skipping lakera prompt injection, input has no text") + return _json_data = json.dumps( { "input": text, @@ -227,6 +247,10 @@ class lakeraAI_Moderation(CustomGuardrail): } ) + if not _json_data and call_type == "pass_through_endpoint": + verbose_proxy_logger.debug("Skipping lakera prompt injection, pass-through body has no messages or input") + return + verbose_proxy_logger.debug("Lakera AI Request Args %s", _json_data) # https://platform.lakera.ai/account/api-keys diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 3e5c54b4501..d127e10e685 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -781,6 +781,9 @@ 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: + verbose_proxy_logger.debug("Lakera AI: not running guardrail. No inspectable text in data") + return response # Call Lakera guardrail lakera_guardrail_response, _ = await self.call_v2_guard( diff --git a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py index 43eb54d1464..71772ad6c46 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/noma/noma_v2.py @@ -21,6 +21,7 @@ from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.llms.custom_httpx.http_handler import ( @@ -63,6 +64,11 @@ class _Action(str, enum.Enum): GUARDRAIL_INTERVENED = "GUARDRAIL_INTERVENED" +def _user_api_key_alias(request_data: Mapping[str, object], metadata_key: str) -> object: + """The key alias recorded under ``metadata_key``, which a pass-through client can set to anything.""" + return (as_str_mapping(request_data.get(metadata_key)) or {}).get("user_api_key_alias") + + class NomaV2Guardrail(CustomGuardrail): reads_stream_marks: ClassVar[bool] = True """Noma keeps one record per streamed reply when every scan carries ``stream`` and only the last one carries the @@ -336,8 +342,8 @@ class NomaV2Guardrail(CustomGuardrail): # (ports v1 fallback from PR #16832). if application_id is None: application_id = self._get_non_empty_str( - request_data.get("litellm_metadata", {}).get("user_api_key_alias") - ) or self._get_non_empty_str(request_data.get("metadata", {}).get("user_api_key_alias")) + _user_api_key_alias(request_data, "litellm_metadata") + ) or self._get_non_empty_str(_user_api_key_alias(request_data, "metadata")) try: payload: Final = self._build_scan_payload( diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 4444d6d25f9..501c8352b10 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -4377,11 +4377,12 @@ async def add_guardrails_from_policy_engine( all_tags: Final = _policy_match_tags(data, metadata_variable_name) _team_alias: Final = user_api_key_dict.team_alias _key_alias: Final = user_api_key_dict.key_alias + _model: Final = data.get("model") context: Final = PolicyMatchContext.model_validate( { "team_alias": _team_alias if isinstance(_team_alias, str) else None, "key_alias": _key_alias if isinstance(_key_alias, str) else None, - "model": data.get("model"), + "model": _model if isinstance(_model, str) else None, "tags": all_tags, } ) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 627164abb93..b6c330daabe 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -443,6 +443,7 @@ async def gemini_proxy_route( custom_llm_provider="gemini", is_streaming_request=is_streaming_request, query_params=merged_params, + inherit_auth_guardrails=True, ) # dynamically construct pass-through endpoint based on incoming path received_value: Final = await endpoint_func( request, @@ -1108,6 +1109,7 @@ async def anthropic_proxy_route( target=str(updated_url), custom_headers=_upstream_headers_for_anthropic_route(request, user_api_key_dict, auth_header), is_streaming_request=is_streaming_request, + inherit_auth_guardrails=True, ) # dynamically construct pass-through endpoint based on incoming path received_value: Final = await endpoint_func( request, @@ -3230,6 +3232,7 @@ async def _base_vertex_proxy_route( upstream_headers_filter=partial( _without_dropped_anthropic_betas, _vertex_dropped_anthropic_betas(general_settings, request.headers) ), + inherit_auth_guardrails=True, ) # dynamically construct pass-through endpoint based on incoming path try: @@ -3832,6 +3835,7 @@ class BaseOpenAIPassThroughHandler: if custom_llm_provider else None ), + inherit_auth_guardrails=True, ) # dynamically construct pass-through endpoint based on incoming path received_value: Final = await endpoint_func( request, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 7f5983d5799..a2b77b3168e 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -6,10 +6,19 @@ import posixpath import re import traceback from base64 import b64encode -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Iterable, Mapping, MutableMapping, Sequence +from collections.abc import ( + AsyncGenerator, + AsyncIterator, + Awaitable, + Callable, + Iterable, + Mapping, + MutableMapping, + Sequence, +) from dataclasses import dataclass from datetime import datetime -from functools import reduce +from functools import partial, reduce from itertools import chain, count, groupby, takewhile from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, TypeAlias, TypedDict, TypeVar, cast @@ -57,7 +66,9 @@ from litellm.constants import ( ) from litellm.integrations.custom_guardrail import ( CustomGuardrail, + guardrail_entry_names, guardrail_request_data_with_streaming, + merged_guardrail_entries, without_server_streaming_classification, ) from litellm.integrations.custom_logger import CustomLogger @@ -113,6 +124,7 @@ from litellm.proxy.common_request_processing import ( open_sse_before_first_byte, resolve_litellm_call_id, ) +from litellm.proxy.common_utils.callback_utils import strip_callback_config from litellm.proxy.common_utils.error_body_call_id import JSON_OBJECT, error_body_call_id, with_call_id from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -136,6 +148,7 @@ from litellm.proxy.litellm_pre_call_utils import ( # noqa: F401 # legacy modul _key_or_team_allows_client_pricing_override, # pyright: ignore[reportPrivateUsage] # reuse the proxy's pricing trust policy _strip_client_pricing_overrides, # pyright: ignore[reportPrivateUsage] # sanitize before trusted hooks add guardrail costs get_dynamic_logging_metadata, # pyright: ignore[reportPrivateUsage] # shared proxy helper, same import style as _read_request_body above + move_guardrails_to_metadata, # pyright: ignore[reportUnknownVariableType] # untyped data param; the same merge the native routes run redact_credential_headers, strip_client_guardrail_information, ) @@ -145,10 +158,11 @@ from litellm.proxy.route_priority import ( configured_pass_through_routes_first, shadowed_by_builtin_claim, ) -from litellm.proxy.utils import normalize_route_for_root_path +from litellm.proxy.utils import normalize_route_for_root_path, pipeline_managed_guardrail_names from litellm.repositories.team_repository import TeamRepository from litellm.secret_managers.main import get_secret_str from litellm.types import utils as types_utils +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.litellm_params import ProxyRequestState, wire_names from litellm.types.llms.custom_http import httpxSpecialProvider from litellm.types.passthrough_endpoints.pass_through_endpoints import ( @@ -169,7 +183,7 @@ from .llm_provider_handlers.tinyfish_passthrough_logging_handler import ( ) from .llm_provider_handlers.vertex_passthrough_logging_handler import VertexPassthroughLoggingHandler from .streaming_handler import PassThroughStreamingHandler -from .success_handler import PassThroughEndpointLogging +from .success_handler import PassThroughEndpointLogging, finalize_cost_without_success_logs from .upstream_usage_headers import ( UpstreamReportedUsage, apply_upstream_reported_model, @@ -180,12 +194,27 @@ from .upstream_usage_headers import ( if TYPE_CHECKING: from litellm.proxy.proxy_server import ProxyConfig +_GuardedResponse = TypeVar("_GuardedResponse") + router: Final = APIRouter() pass_through_endpoint_logging: Final = PassThroughEndpointLogging() +_blocked_response_cost_logging: Final = PassThroughEndpointLogging(log_dispatch=finalize_cost_without_success_logs) _METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata")) _DECODED_JSON: Final = TypeAdapter(object) +_GUARDRAIL_ENTRIES: Final = TypeAdapter(list[object]) +_FIELDS_GUARDRAILS_READ_BY_SHAPE: Final = ( + ("guardrails", (list, dict)), + ("guardrail_config", (dict,)), + ("proxy_server_request", (dict,)), +) +_POST_CALL_SELECTION_FIELDS: Final = ( + "_guardrail_pipelines", + "_pipeline_managed_guardrails", + "spend_logs_metadata", + "tags", +) _KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState)) _PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",)) @@ -867,7 +896,208 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): _update_stream_param_based_on_request_body = update_stream_param_based_on_request_body -def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | None) -> None: +def _drop_unmergeable_guardrail_fields( + parsed_body: dict[str, object], # mutable-ok: strips the request body in place before the guardrail merge +) -> None: + """Drop a ``guardrails`` that isn't a list or an object, and a ``guardrail_config`` or ``proxy_server_request`` that + isn't an object, which the merge and the guardrails can't read. The provider receives none of these fields.""" + for field, shape in _FIELDS_GUARDRAILS_READ_BY_SHAPE: + if field in parsed_body and not isinstance(parsed_body[field], shape): + parsed_body.pop(field) + + +async def _inherit_auth_guardrails( + parsed_body: dict[str, object] | None, # mutable-ok: merges the guardrails into the live request body + user_api_key_dict: UserAPIKeyAuth, + configured_guardrails: Sequence[object] | None, + inherit: bool, +) -> Sequence[object] | None: + """Return the guardrails a pass-through request runs after the provider responds. + + The endpoint's configured guardrails win. Without any, and with ``inherit`` and a JSON body, the key, team, + project, policy and request-body guardrails are merged into ``parsed_body["metadata"]`` for the pre-call hook, + with the ones the client named in ``metadata.guardrails`` kept after them, and returned only when at least one + of them runs post_call. + """ + if configured_guardrails or not inherit or not parsed_body: + return configured_guardrails + _drop_unmergeable_guardrail_fields(parsed_body) + client_sent_metadata: Final = "metadata" in parsed_body + if not isinstance(parsed_body.get("metadata"), dict): + parsed_body["metadata"] = {} # rebind-ok: the merge below writes into the live request body's metadata + client_named: Final = (as_str_mapping(parsed_body["metadata"]) or {}).get("guardrails") + await move_guardrails_to_metadata( + data=parsed_body, + _metadata_variable_name="metadata", + user_api_key_dict=user_api_key_dict, + ) + metadata: Final = parsed_body["metadata"] + if not metadata and not client_sent_metadata: + parsed_body.pop("metadata") + return None + merged: Final = (as_str_mapping(metadata) or {}).get("guardrails") + if not isinstance(merged, list): + return None + entries: Final = merged_guardrail_entries( + _GUARDRAIL_ENTRIES.validate_python(merged), _client_named_guardrail_entries(client_named) + ) + if entries and isinstance(metadata, dict): + metadata["guardrails"] = entries + admin_metadata: Final = _admin_guardrail_metadata(user_api_key_dict) + admin_scoped_body: Final = { + **parsed_body, + **{key: {**(as_str_mapping(parsed_body.get(key)) or {}), **admin_metadata} for key in sorted(_METADATA_KEYS)}, + } + return entries if entries and _named_guardrail_runs_post_call(admin_scoped_body, entries) else None + + +def _client_named_guardrail_entries(client_named: object) -> Sequence[object]: + """The guardrails ``metadata.guardrails`` names: its entries, or the keys of the dict form the guardrails read""" + if isinstance(client_named, list): + return _GUARDRAIL_ENTRIES.validate_python(client_named) + return tuple(as_str_mapping(client_named) or ()) + + +def _admin_guardrail_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object]: + """The key and team metadata ``should_run_guardrail`` reads a global guardrail opt-out from, as the native routes set it""" + return MappingProxyType( + { + "user_api_key_metadata": strip_callback_config(user_api_key_dict.metadata), + "user_api_key_team_metadata": strip_callback_config(user_api_key_dict.team_metadata), + } + ) + + +def _holds_resolved_pipelines(pipelines: object) -> bool: + """Whether a pipeline slot holds the policy engine's (name, pipeline) tuples, which a JSON body can't""" + entries: Final = as_object_list(pipelines) + return entries is not None and all(isinstance(entry, tuple) for entry in entries) + + +def _post_call_selection_metadata(kwargs: Mapping[str, object], inherited: bool) -> Mapping[str, object]: + """The policy pipelines resolved for a request that inherits auth guardrails, the tags a tag-scoped guardrail mode + reads and the spend_logs_metadata a blocked response's row keeps, from the logging metadata that request setup moved + the body's ``metadata`` into""" + if not inherited: + return MappingProxyType({}) + litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} + metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {} + selection: Final = {field: metadata[field] for field in _POST_CALL_SELECTION_FIELDS if field in metadata} + if "_guardrail_pipelines" in selection and not _holds_resolved_pipelines(selection["_guardrail_pipelines"]): + selection.pop("_guardrail_pipelines") + return selection + + +def _starts_post_call_pass(guardrails: Sequence[object] | None, policy_pipelines: Mapping[str, object]) -> bool: + return bool(guardrails) or bool(pipeline_managed_guardrail_names({"metadata": dict(policy_pipelines)}, "post_call")) + + +def _post_call_hook_data( + parsed_body: Mapping[str, object] | None, + guardrails: Sequence[object] | None, + logging_obj: LiteLLMLoggingObj, + user_api_key_dict: UserAPIKeyAuth, + policy_pipelines: Mapping[str, object], + inherited: bool, +) -> dict[str, object]: # mutable-ok: the guardrails record their results in this request payload + """The body with what the post-call guardrails read and the kwargs split stripped: ``metadata`` naming the guardrails, + which ``should_run_guardrail`` checks, the policy pipelines to run, and the logging object, which unified + (apply_guardrail) guardrails resolve a pass-through response's call type from when the route has no + API_ROUTE_TO_CALL_TYPES entry. A custom endpoint whose guardrails don't run post_call gets no logging object, so a + global unified guardrail still can't resolve its response.""" + body: Final = parsed_body or {} + admin_metadata: Final = _admin_guardrail_metadata(user_api_key_dict) if inherited else MappingProxyType({}) + data: Final = { + **body, + "metadata": { + **(as_str_mapping(body.get("metadata")) or {}), + **policy_pipelines, + "guardrails": guardrails or (), + **admin_metadata, + }, + } + if not inherited and not _named_guardrail_runs_post_call(data, guardrails or ()): + return data + return {**data, "litellm_logging_obj": logging_obj} + + +def _provider_charge_settlement( + response: httpx.Response, + response_body: Mapping[str, object] | None, + url_route: str, + start_time: datetime, + logging_obj: LiteLLMLoggingObj, + parsed_body: Mapping[str, object] | None, + custom_llm_provider: str | None, + handler_kwargs: Mapping[str, object], +) -> Callable[[], Awaitable[None]]: + return partial( + _blocked_response_cost_logging.pass_through_async_success_handler, + finalize_before_dispatch=True, + httpx_response=response, + response_body=response_body, + url_route=url_route, + result="", + start_time=start_time, + end_time=datetime.now(start_time.tzinfo), + logging_obj=logging_obj, + cache_hit=False, + request_body=parsed_body or {}, + custom_llm_provider=custom_llm_provider, + **handler_kwargs, + ) + + +async def _run_post_call_guardrails_keeping_the_charge( + run_guardrails: Callable[[], Awaitable[_GuardedResponse]], + settle_provider_charge: Callable[[], Awaitable[None]], + settles_the_charge: bool = True, +) -> _GuardedResponse: + """A guardrail that blocks the response, or fails, does so after the provider answered and billed the call, so the + charge is settled before the failure path writes its row, as on the native routes. A custom endpoint runs every + hook here, not only guardrails, and doesn't settle, so its failure row stays uncharged.""" + try: + return await run_guardrails() + except Exception: + if not settles_the_charge: + raise + try: + await settle_provider_charge() + except Exception: # noqa: BLE001 - the guardrail's outcome still reaches the client + verbose_proxy_logger.exception("pass_through_endpoint: could not settle a blocked response's cost") + raise + + +def _response_after_post_call_guardrails( + guarded_body: object, + response_body: dict[str, object] | None, # mutable-ok: pass_through_request's own response dict, returned as is + content: bytes, +) -> tuple[dict[str, object] | None, bytes, bool]: # mutable-ok: pass_through_request keeps a dict for the id rewrite + """The response body, its bytes and whether they changed, once the guardrails returned ``guarded_body``. A hook that + returns anything but a JSON object leaves the response as the provider sent it.""" + if isinstance(guarded_body, dict): + guarded: Final = cast("dict[str, object]", guarded_body) # cast-ok: the hook's JSON response + return guarded, json.dumps(guarded).encode("utf-8"), True + verbose_proxy_logger.debug( + "pass_through_endpoint: post_call_success_hook returned %s, expected dict — using original response", + type(guarded_body).__name__, + ) + return response_body, content, False + + +def _named_guardrail_runs_post_call(data: Mapping[str, object], entries: Sequence[object]) -> bool: + names: Final = frozenset(chain.from_iterable(map(guardrail_entry_names, entries))) + return any( + guardrail.guardrail_name in names + and guardrail.should_run_guardrail( # pyright: ignore[reportUnknownMemberType] # takes the untyped request dict + data=data, event_type=GuardrailEventHooks.post_call + ) + for guardrail in litellm.logging_callback_manager.get_custom_loggers_for_type(callback_type=CustomGuardrail) + if isinstance(guardrail, CustomGuardrail) + ) + + +def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict[str, object] | None) -> None: """Copy guardrail logging entries from ``guardrail_data`` onto ``request_data``. Post-call guardrails run against a throwaway ``hook_data`` dict (its @@ -876,18 +1106,30 @@ def _carry_guardrail_logging_info(request_data: dict, guardrail_data: dict | Non ``standard_logging_guardrail_information`` there and not on the dict the failure handler forwards to ``post_call_failure_hook``. Without this the otel guardrail span is emitted on allow but missing on block. Carry the - entries over so the failure path matches the unified path. + entries over so the failure path matches the unified path. The failure row + reads its ``metadata`` here too, so it keeps the ``spend_logs_metadata`` a + built-in route's post-call pass was given, as its success row would. """ if guardrail_data is None: return + spend_logs_metadata: Final = (as_str_mapping(guardrail_data.get("metadata")) or {}).get("spend_logs_metadata") + if spend_logs_metadata is not None: + request_fields: Final = cast("Mapping[str, object]", request_data) # cast-ok: the JSON body plus logging kwargs + request_data["metadata"] = { # rebind-ok: the failure path hands this dict to post_call_failure_hook + **(as_str_mapping(request_fields.get("metadata")) or {}), + "spend_logs_metadata": spend_logs_metadata, + } source_key: Final = get_metadata_variable_name_from_kwargs(guardrail_data) - source_metadata: Final = guardrail_data.get(source_key) or {} + source_metadata: Final = as_str_mapping(guardrail_data.get(source_key)) or {} entries: Final = source_metadata.get("standard_logging_guardrail_information") if not entries: return _, metadata = get_or_create_metadata_bucket(request_data) - metadata.setdefault("standard_logging_guardrail_information", list(entries)) + metadata.setdefault( + "standard_logging_guardrail_information", + list(cast("Iterable[object]", entries)), # cast-ok: guardrails record a list of entries here + ) _GOOGLE_PASSTHROUGH_PROVIDERS: Final = frozenset( @@ -1255,6 +1497,7 @@ async def pass_through_request( timeout: float | None = None, upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers, operation: GenAIOperation | None = None, + inherit_auth_guardrails: bool = False, ): """ Pass through endpoint handler, makes the httpx request for pass-through endpoints and ensures logging hooks are called @@ -1277,6 +1520,8 @@ async def pass_through_request( general_settings.pass_through_request_timeout, then 600s. upstream_headers_filter: Optional last rewrite of the headers sent upstream, after x-pass- expansion operation: Optional GenAI operation the route reports, unless the upstream names its own + inherit_auth_guardrails: Run the key, team, project and policy guardrails on a JSON body, as the + native routes do, when the pass-through guardrails config resolves to none """ from litellm.exceptions import ModifyResponseException from litellm.litellm_core_utils.litellm_logging import Logging @@ -1385,6 +1630,13 @@ async def pass_through_request( _parsed_body["metadata"]["guardrails"] = guardrails_to_run verbose_proxy_logger.debug("Added guardrails to passthrough request metadata: %s", guardrails_to_run) + post_call_guardrails: Final = await _inherit_auth_guardrails( + parsed_body=cast(dict[str, object] | None, _parsed_body), # cast-ok: json + user_api_key_dict=user_api_key_dict, + configured_guardrails=guardrails_to_run, + inherit=inherit_auth_guardrails, + ) + ## LOGGING OBJECT ## - initialize before pre_call_hook so guardrails can access it # Surface the requested model (when the body carries one) so logging/spans # read e.g. ``chat gpt-4o`` instead of ``chat unknown``. @@ -1466,7 +1718,7 @@ async def pass_through_request( ) hooked_body: Final = _json_object(_parsed_body) client_user_id: Final = _anthropic_user_id_to_forward( - hooked_body, str(url), proxy_logging_obj.request_has_pre_call_guardrails + hooked_body, str(url), proxy_logging_obj.request_is_scanned_before_the_provider ) kwargs = HttpPassThroughEndpointHelpers.init_kwargs_for_pass_through_endpoint( # rebind-ok: pre-existing rebinding on a rename-only line user_api_key_dict=user_api_key_dict, @@ -1925,33 +2177,48 @@ async def pass_through_request( logging_obj=logging_obj, ) - if response.status_code < 400 and response_body is not None and guardrails_to_run: - # Build an enriched data dict: _parsed_body has been stripped of - # `metadata` by both pre_call_hook and _init_kwargs_for_pass_through_endpoint, - # so we re-attach the configured guardrails here so should_run_guardrail - # sees them. - hook_data: Final = dict(_parsed_body or {}) - existing_metadata = hook_data.get("metadata") - if not isinstance(existing_metadata, dict): - existing_metadata = {} - hook_data["metadata"] = { - **existing_metadata, - "guardrails": guardrails_to_run, - } - post_call_guardrail_data = hook_data - response_body = await proxy_logging_obj.post_call_success_hook( - data=hook_data, - user_api_key_dict=user_api_key_dict, - response=response_body, + policy_pipelines: Final = _post_call_selection_metadata( + cast(dict[str, object], kwargs), # cast-ok: the logging fields init_kwargs built above + inherit_auth_guardrails, + ) + if ( + response.status_code < 400 + and response_body is not None + and _starts_post_call_pass(post_call_guardrails, policy_pipelines) + ): + json_body: Final = cast(dict[str, object] | None, _parsed_body) # cast-ok: json + hook_data: Final = _post_call_hook_data( + json_body, + post_call_guardrails, + logging_obj, + user_api_key_dict, + policy_pipelines, + inherit_auth_guardrails, + ) + post_call_guardrail_data = hook_data + guarded_body: Final = await _run_post_call_guardrails_keeping_the_charge( + partial( + proxy_logging_obj.post_call_success_hook, + data=hook_data, + user_api_key_dict=user_api_key_dict, + response=response_body, + guardrails_only=inherit_auth_guardrails, + ), + _provider_charge_settlement( + response, + cast(dict[str, object] | None, response_body), # cast-ok: json + str(url), + start_time, + logging_obj, + json_body, + custom_llm_provider, + cast(dict[str, object], kwargs), # cast-ok: the logging fields init_kwargs built above + ), + settles_the_charge=inherit_auth_guardrails, + ) + response_body, content, _content_modified = _response_after_post_call_guardrails( + guarded_body, response_body, content ) - if isinstance(response_body, dict): - content = json.dumps(response_body).encode("utf-8") - _content_modified = True - else: - verbose_proxy_logger.debug( - "pass_through_endpoint: post_call_success_hook returned %s, expected dict — using original response", - type(response_body).__name__, - ) elif response_body is None: verbose_proxy_logger.debug( "pass_through_endpoint: response body not JSON-parseable, skipping post-call guardrails" @@ -2320,6 +2587,7 @@ def create_pass_through_route( timeout: float | None = None, upstream_headers_filter: Callable[[Mapping[str, str]], Mapping[str, str]] = _unfiltered_headers, operation: GenAIOperation | None = None, + inherit_auth_guardrails: bool = False, ): # check if target is an adapter.py or a url from litellm._uuid import uuid @@ -2469,6 +2737,7 @@ def create_pass_through_route( timeout=cast(float | None, param_timeout), upstream_headers_filter=upstream_headers_filter, operation=operation, + inherit_auth_guardrails=inherit_auth_guardrails, ) finally: if hasattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY): diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 8f55906498b..335b4889037 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -2,7 +2,7 @@ import asyncio import json from collections.abc import Mapping, Sequence from contextlib import suppress -from datetime import datetime +from datetime import datetime, timezone from types import MappingProxyType from typing import Any, Final from urllib.parse import urlparse @@ -79,6 +79,22 @@ def _websocket_frames( ) +async def finalize_cost_without_success_logs( + *, + logging_obj: LiteLLMLoggingObj, + standard_logging_response_object: object, + **_success_log_fields: object, # kwargs-ok: accepts every field a log dispatch receives; none are needed here +) -> None: + """Settles the provider's charge on the logging object for a response that won't be logged as a success. A response + with no cost result settles at 0, the spend its success row would carry, so the failure row keeps its usage.""" + from litellm.proxy.spend_tracking.cost_finalization import finalize_proxy_response_cost, prepare_response_cost + + prepare_response_cost( + logging_obj, standard_logging_response_object, datetime.now(timezone.utc), settle_unknown_cost_at_zero=True + ) + await finalize_proxy_response_cost(logging_obj, standard_logging_response_object) + + class PassThroughEndpointLogging: def __init__( self, diff --git a/litellm/proxy/spend_tracking/cost_finalization.py b/litellm/proxy/spend_tracking/cost_finalization.py index c42579714e8..2db4b05f1e6 100644 --- a/litellm/proxy/spend_tracking/cost_finalization.py +++ b/litellm/proxy/spend_tracking/cost_finalization.py @@ -11,9 +11,11 @@ from litellm.proxy.hooks.search_free_tier import SharedCounter, allocate_search_ from litellm.types.utils import CostBreakdown, StandardLoggingPayload -def prepare_response_cost(logging_obj: Logging, response: object, now: datetime) -> PreparedCost | None: +def prepare_response_cost( + logging_obj: Logging, response: object, now: datetime, settle_unknown_cost_at_zero: bool = False +) -> PreparedCost | None: try: - return logging_obj.prepare_response_cost(response, now) + return logging_obj.prepare_response_cost(response, now, settle_unknown_cost_at_zero) except Exception: # noqa: BLE001 # preparation must not prevent serving an already completed response verbose_proxy_logger.exception("Unable to prepare completed response cost for %s", logging_obj.litellm_call_id) return None diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 71ecae74058..77162cb92a4 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -157,6 +157,7 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.integrations.SlackAlerting.utils import add_langfuse_trace_id_to_alert from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route from litellm.litellm_core_utils.core_helpers import ( + as_object_list, coerce_token_limit, get_or_create_metadata_bucket, independent_snapshot, @@ -774,6 +775,20 @@ def pipeline_managed_guardrail_names( ) +def _runs_after_the_provider(pipeline_entry: object) -> bool: + if not isinstance(pipeline_entry, tuple): + return False + pipeline: Final = cast("tuple[object, ...]", pipeline_entry)[-1] # cast-ok: the policy engine's (name, pipeline) + return getattr(pipeline, "mode", None) == "post_call" + + +def _holds_pipeline_before_the_provider(pipelines: object) -> bool: + if not pipelines: + return False + entries: Final = as_object_list(pipelines) + return entries is None or not all(map(_runs_after_the_provider, entries)) + + def _ran_or_runs_outside_pipelines( callback: CustomGuardrail | None, data: Mapping[str, object], @@ -875,6 +890,29 @@ def _drop_response_steps_that_did_not_run( _drop_applied_guardrails_in_place(data.get(bucket_key), not_run) +def _add_applied_guardrails_in_place( + data: dict[str, object], # mutable-ok: the names are recorded in the live request body's metadata + names: Sequence[str], +) -> None: + """Adds each name the way add_guardrail_to_applied_guardrails_header does, in time linear in the names""" + from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header + + if not names: + return + metadata: Final = cast("dict[str, object]", get_or_create_metadata_bucket(data)[1]) # cast-ok: str-keyed bucket + if "applied_guardrails" not in metadata: + metadata["applied_guardrails"] = list(dict.fromkeys(names)) + return + applied: Final = metadata["applied_guardrails"] + if not isinstance(applied, list): + for name in names: + add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=name) + return + entries: Final = cast("list[object]", applied) # cast-ok: entries are only compared with names + present: Final = frozenset(name for name in entries if isinstance(name, str)) + entries.extend(dict.fromkeys(name for name in names if name not in present)) + + def _drop_applied_guardrails_in_place(bucket: object, names: frozenset[str]) -> None: if not isinstance(bucket, dict): return @@ -2904,10 +2942,6 @@ class ProxyLogging: policy pipeline steps that have not run. With ``skip_guardrails`` the pipelines run on another request, so every step is added. """ - from litellm.proxy.common_utils.callback_utils import ( - add_guardrail_to_applied_guardrails_header, - ) - metadata_standard: Final = data.get("metadata") or {} metadata_litellm: Final = data.get("litellm_metadata") or {} @@ -2927,18 +2961,21 @@ class ProxyLogging: if not isinstance(applied_guardrails, list): applied_guardrails = [] - not_run: Final = ( - frozenset[str]() - if skip_guardrails - else _request_steps_that_have_not_run(cast(dict[str, object], data)) # cast-ok: request payload + request: Final = cast("dict[str, object]", data) # cast-ok: request payload + not_run: Final = frozenset[str]() if skip_guardrails else _request_steps_that_have_not_run(request) + already_applied: Final = frozenset( + name + for name in cast("Sequence[object]", applied_guardrails) # cast-ok: checked to be a list above + if isinstance(name, str) + ) + _add_applied_guardrails_in_place( + request, + [ + name + for name in guardrails_in_metadata + if isinstance(name, str) and name not in already_applied and name not in not_run + ], ) - for guardrail_name in guardrails_in_metadata: - if ( - isinstance(guardrail_name, str) - and guardrail_name not in applied_guardrails - and guardrail_name not in not_run - ): - add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=guardrail_name) async def _maybe_execute_pipelines( self, @@ -3111,8 +3148,18 @@ class ProxyLogging: request_data: Mapping[str, object], event_type: GuardrailEventHooks = GuardrailEventHooks.pre_call, ) -> bool: - if _policy_state_metadata(request_data).get("_guardrail_pipelines"): - return True + return bool(_policy_state_metadata(request_data).get("_guardrail_pipelines")) or ( + self._request_has_guardrail_callback(request_data, event_type) + ) + + def request_is_scanned_before_the_provider(self, request_data: Mapping[str, object]) -> bool: + """``request_has_pre_call_guardrails``, except that a resolved policy pipeline running post_call doesn't count""" + return _holds_pipeline_before_the_provider( + _policy_state_metadata(request_data).get("_guardrail_pipelines") + ) or self._request_has_guardrail_callback(request_data, GuardrailEventHooks.pre_call) + + @staticmethod + def _request_has_guardrail_callback(request_data: Mapping[str, object], event_type: GuardrailEventHooks) -> bool: caps: Final = ProxyLogging.callback_capabilities() if caps.has_content_enforcer: return True @@ -4349,6 +4396,7 @@ class ProxyLogging: data: dict, response: LLMResponseTypes, user_api_key_dict: UserAPIKeyAuth, + guardrails_only: bool = False, ): """ Allow user to modify outgoing data @@ -4358,6 +4406,10 @@ class ProxyLogging: 2. /embeddings 3. /image/generation 4. /files + + ``guardrails_only`` is for a relayed response that never reached a router deployment and isn't in the shape the + other hooks expect: only the guardrails the request selected run, not the guardrails of a model that merely + shares the body's ``model`` name, nor the callbacks that aren't guardrails. """ from litellm.proxy.proxy_server import llm_router @@ -4375,7 +4427,9 @@ class ProxyLogging: guardrail_callbacks, other_callbacks = _partition_post_call_callbacks() 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) + guardrail_data: Final = ( + data if guardrails_only else check_and_merge_model_level_guardrails(data=data, llm_router=llm_router) + ) parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = tuple( callback @@ -4443,7 +4497,7 @@ class ProxyLogging: ############ Handle CustomLogger ############################### ################################################################# - for callback in other_callbacks: + for callback in () if guardrails_only else other_callbacks: callback_response: LLMResponseTypes | None = await callback.async_post_call_success_hook( user_api_key_dict=user_api_key_dict, data=data, response=response ) diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py new file mode 100644 index 00000000000..b3d711c3e9d --- /dev/null +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -0,0 +1,1044 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from contextlib import ExitStack +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal + +import httpx +import pytest +import yaml +from integration._support.anthropic_sse import message_json +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually +from integration._support.openai_wire import chat_reply +from integration._support.process import OwnedProxy, owned_proxy_process +from integration._support.safeguards_rig import key_hash +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +GUARDRAIL: Final = "passthrough-auth-guardrail" +GLOBAL_POST_CALL: Final = "passthrough-global-post-call" +KEY_POST_CALL: Final = "passthrough-key-post-call" +GENERIC_POST_CALL: Final = "passthrough-generic-post-call" +BLOCKING_POST_CALL: Final = "passthrough-post-call-blocks-every-reply" +TAG_SCOPED_POST_CALL: Final = "passthrough-tag-scoped-post-call" +POST_CALL_TAG: Final = "scan-the-reply" +MODEL_POST_CALL: Final = "passthrough-model-post-call" +GUARDED_MODEL: Final = "passthrough-guarded-model" +POST_CALL_GUARDED_MODEL: Final = "passthrough-post-call-guarded-model" +BLOCK_MARKER: Final = "block-this-prompt" +RAISING_HOOK_MODEL: Final = "passthrough-model-a-logger-hook-cannot-read" +_RAISING_HOOK_MODULE: Final = f""" +from litellm.integrations.custom_logger import CustomLogger + + +class RaisingHook(CustomLogger): + async def async_post_call_success_hook(self, data, user_api_key_dict, response): + if data.get("model") == "{RAISING_HOOK_MODEL}": + raise ValueError("this hook only reads native responses") + + +proxy_handler_instance = RaisingHook() +""" +OPENAI_ROUTE: Final = "/openai/v1/chat/completions" +ANTHROPIC_ROUTE: Final = "/anthropic/v1/messages" +GEMINI_ROUTE: Final = "/gemini/v1beta/models/gemini-2.5-flash:generateContent" +CUSTOM_ROUTE: Final = "/passthrough-auth-custom" +CUSTOM_GUARDED_ROUTE: Final = "/passthrough-auth-custom-guarded" +CUSTOM_PRE_CALL_ROUTE: Final = "/passthrough-auth-custom-pre-call" +GLOBAL_UNIFIED_POST_CALL: Final = "passthrough-global-unified-post-call" +UNPRICED_MODEL: Final = "passthrough-auth-unpriced-model" +MONITOR_POST_CALL: Final = "passthrough-auth-monitor-post-call" +MONITORED_PROMPT: Final = "reply-with-the-monitored-id" +MONITORED_REPLY_ID: Final = "msg_passthrough_auth_monitored" +PHONE_POST_CALL: Final = "passthrough-auth-phone-post-call" +PIPELINE_BLOCKS: Final = "passthrough-pipeline-blocks-the-response" +PIPELINE_ALLOWS: Final = "passthrough-pipeline-allows-the-response" +PIPELINE_BEFORE_THE_CALL: Final = "passthrough-pipeline-before-the-call" +PIPELINE_MODIFIES: Final = "passthrough-pipeline-modifies-the-response" +_CONFIG: Final = Path(__file__).resolve().parents[1] / "proxy_config.yaml" +_JSON: Final = TypeAdapter(dict[str, JsonValue]) +_JSON_LIST: Final = TypeAdapter(list[JsonValue]) + +Level = Literal["key", "team", "project"] + + +def _route_body(route: str, prompt: str) -> dict[str, JsonValue]: + match route: + case "/anthropic/v1/messages": + return {"model": "claude-sonnet-5-5", "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]} + case "/gemini/v1beta/models/gemini-2.5-flash:generateContent": + return {"contents": [{"role": "user", "parts": [{"text": prompt}]}]} + case _: + return {"model": "gpt-5.6", "messages": [{"role": "user", "content": prompt}]} + + +def _pipeline_policy(mode: str, guardrail: str, on_fail: str) -> dict[str, JsonValue]: + return { + "guardrails": {"add": [guardrail]}, + "pipeline": {"mode": mode, "steps": [{"guardrail": guardrail, "on_pass": "allow", "on_fail": on_fail}]}, + } + + +def _verdict(request: Request) -> Reply: + action: Final = "BLOCKED" if BLOCK_MARKER.encode() in request.body else "NONE" + return Reply(body=json.dumps({"action": action, "blocked_reason": "synthetic block"}).encode()) + + +def _block_every_reply(_request: Request) -> Reply: + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": "synthetic block"}).encode()) + + +_OPENAI_BODY: Final = json.dumps( + { + **_JSON.validate_json(chat_reply("chatcmpl-passthrough-auth", "gpt-5.6", "Hi", stream=False).body), + "created": 1793600123, + }, + indent=2, +).encode() + + +def _openai_reply(_request: Request) -> Reply: + return Reply(body=_OPENAI_BODY) + + +def _lakera_pii_only(_request: Request) -> Reply: + return Reply( + body=json.dumps( + { + "flagged": True, + "breakdown": [{"detector_type": "pii/email", "detected": True, "message_id": 0}], + "payload": [], + } + ).encode() + ) + + +def _lakera_clean(_request: Request) -> Reply: + return Reply(body=json.dumps({"flagged": False, "breakdown": [], "payload": []}).encode()) + + +def _lakera_flagged(_request: Request) -> Reply: + return Reply( + body=json.dumps( + { + "flagged": True, + "breakdown": [{"detector_type": "prompt_attack", "detected": True, "message_id": 0}], + "payload": [], + } + ).encode() + ) + + +_ANTHROPIC_BODY: Final = json.dumps( + _JSON.validate_json(message_json("msg_passthrough_auth", "claude-sonnet-5-5", "Hi")), indent=2 +).encode() +_GEMINI_BODY: Final = json.dumps( + { + "candidates": [{"content": {"role": "model", "parts": [{"text": "Hi"}]}, "finishReason": "STOP"}], + "usageMetadata": {"promptTokenCount": 4, "candidatesTokenCount": 1, "totalTokenCount": 5}, + }, + indent=2, +).encode() +_UPSTREAM_BODIES: Final = MappingProxyType( + {OPENAI_ROUTE: _OPENAI_BODY, ANTHROPIC_ROUTE: _ANTHROPIC_BODY, GEMINI_ROUTE: _GEMINI_BODY} +) + + +_MONITORED_ANTHROPIC_BODY: Final = json.dumps( + _JSON.validate_json(message_json(MONITORED_REPLY_ID, "claude-sonnet-5-5", "Hi")) +).encode() + + +_UNPRICED_ANTHROPIC_BODY: Final = json.dumps( + { + **JSON_OBJECT.validate_json(message_json("msg_passthrough_auth_unpriced", UNPRICED_MODEL, "Hi")), + "usage": {"input_tokens": 300, "output_tokens": 200}, + } +).encode() + + +def _anthropic_reply(request: Request) -> Reply: + if UNPRICED_MODEL.encode() in request.body: + return Reply(body=_UNPRICED_ANTHROPIC_BODY) + return Reply(body=_MONITORED_ANTHROPIC_BODY if MONITORED_PROMPT.encode() in request.body else _ANTHROPIC_BODY) + + +def _gemini_reply(_request: Request) -> Reply: + return Reply(body=_GEMINI_BODY) + + +@dataclass(frozen=True, slots=True) +class AuthGuardrailRig: + owned: OwnedProxy + vendor: Wire + post_call_vendor: Wire + key_post_call_vendor: Wire + generic_post_call_vendor: Wire + global_unified_vendor: Wire + model_post_call_vendor: Wire + upstreams: Mapping[str, Wire] + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[AuthGuardrailRig]: + directory: Final = tmp_path_factory.mktemp("passthrough-auth-guardrails") + with ExitStack() as stack: + vendor: Final = stack.enter_context(wire_server(_verdict)) + post_call_vendor: Final = stack.enter_context(wire_server(_lakera_clean)) + key_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_flagged)) + generic_post_call_vendor: Final = stack.enter_context(wire_server(_verdict)) + global_unified_vendor: Final = stack.enter_context(wire_server(_verdict)) + blocking_post_call_vendor: Final = stack.enter_context(wire_server(_block_every_reply)) + monitor_post_call_vendor: Final = stack.enter_context(wire_server(_lakera_pii_only)) + model_post_call_vendor: Final = stack.enter_context(wire_server(_verdict)) + upstreams: Final = MappingProxyType( + { + OPENAI_ROUTE: stack.enter_context(wire_server(_openai_reply)), + ANTHROPIC_ROUTE: stack.enter_context(wire_server(_anthropic_reply)), + GEMINI_ROUTE: stack.enter_context(wire_server(_gemini_reply)), + } + ) + config: Final = _JSON.validate_python(yaml.safe_load(_CONFIG.read_text())) + general_settings: Final = _JSON.validate_python(config.get("general_settings") or {}) + (directory / "passthrough_raising_hook.py").write_text(_RAISING_HOOK_MODULE) + path: Final = directory / "passthrough-auth-guardrails.yaml" + path.write_text( + yaml.safe_dump( + { + **config, + "litellm_settings": { + **_JSON.validate_python(config.get("litellm_settings") or {}), + "callbacks": ["passthrough_raising_hook.proxy_handler_instance"], + }, + "general_settings": { + **general_settings, + "pass_through_endpoints": [ + { + "path": CUSTOM_ROUTE, + "target": f"{upstreams[OPENAI_ROUTE].url}/v1/chat/completions", + "auth": True, + }, + { + "path": CUSTOM_GUARDED_ROUTE, + "target": f"{upstreams[OPENAI_ROUTE].url}/v1/chat/completions", + "auth": True, + "cost_per_request": 0.25, + "guardrails": {BLOCKING_POST_CALL: None}, + }, + { + "path": CUSTOM_PRE_CALL_ROUTE, + "target": f"{upstreams[OPENAI_ROUTE].url}/v1/chat/completions", + "auth": True, + "guardrails": {GUARDRAIL: None}, + }, + ], + }, + "policies": { + PIPELINE_BLOCKS: _pipeline_policy("post_call", KEY_POST_CALL, "block"), + PIPELINE_ALLOWS: _pipeline_policy("post_call", KEY_POST_CALL, "allow"), + PIPELINE_BEFORE_THE_CALL: _pipeline_policy("pre_call", GUARDRAIL, "block"), + PIPELINE_MODIFIES: _pipeline_policy("post_call", KEY_POST_CALL, "modify_response"), + }, + "policy_attachments": [ + {"policy": policy, "tags": [policy]} + for policy in (PIPELINE_BLOCKS, PIPELINE_ALLOWS, PIPELINE_BEFORE_THE_CALL, PIPELINE_MODIFIES) + ], + "model_list": [ + *_JSON_LIST.validate_python(config.get("model_list") or []), + { + "model_name": GUARDED_MODEL, + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "synthetic-openai-key", + "guardrails": [GUARDRAIL], + }, + }, + { + "model_name": POST_CALL_GUARDED_MODEL, + "litellm_params": { + "model": "openai/gpt-5.6", + "api_key": "synthetic-openai-key", + "guardrails": [MODEL_POST_CALL], + }, + }, + ], + "guardrails": [ + { + "guardrail_name": GLOBAL_POST_CALL, + "litellm_params": { + "guardrail": "lakera_v2", + "mode": "post_call", + "default_on": True, + "api_base": post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": GUARDRAIL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": KEY_POST_CALL, + "litellm_params": { + "guardrail": "lakera_v2", + "mode": "post_call", + "default_on": False, + "api_base": key_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": BLOCKING_POST_CALL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": False, + "api_base": blocking_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": TAG_SCOPED_POST_CALL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": {"tags": {POST_CALL_TAG: "post_call"}, "default": "logging_only"}, + "default_on": False, + "api_base": blocking_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": MONITOR_POST_CALL, + "litellm_params": { + "guardrail": "lakera_v2", + "mode": "post_call", + "default_on": False, + "on_flagged": "monitor", + "api_base": monitor_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": GLOBAL_UNIFIED_POST_CALL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": True, + "api_base": global_unified_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": GENERIC_POST_CALL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": False, + "api_base": generic_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, + { + "guardrail_name": PHONE_POST_CALL, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": "post_call", + "default_on": False, + "patterns": [ + {"pattern_type": "prebuilt", "pattern_name": "us_phone", "action": "BLOCK"} + ], + }, + }, + { + "guardrail_name": MODEL_POST_CALL, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "post_call", + "default_on": False, + "api_base": model_post_call_vendor.url, + "api_key": "synthetic-guardrail-key", + }, + }, + ], + } + ) + ) + environment: Final = { + "OPENAI_API_BASE": upstreams[OPENAI_ROUTE].url, + "OPENAI_API_KEY": "synthetic-openai-key", + "ANTHROPIC_API_BASE": upstreams[ANTHROPIC_ROUTE].url, + "ANTHROPIC_API_KEY": "synthetic-anthropic-key", + "GEMINI_API_BASE": upstreams[GEMINI_ROUTE].url, + "GEMINI_API_KEY": "synthetic-gemini-key", + } + owned: Final = stack.enter_context(owned_proxy_process(session_gateway, directory, environment, config=path)) + yield AuthGuardrailRig( + owned, + vendor, + post_call_vendor, + key_post_call_vendor, + generic_post_call_vendor, + global_unified_vendor, + model_post_call_vendor, + upstreams, + ) + + +def _key_with_guardrail(scenario: Scenario, level: Level) -> str: + attached: Final[dict[str, JsonValue]] = {"metadata": {"guardrails": [GUARDRAIL]}} + team: Final = scenario.team(**(attached if level == "team" else {})) + match level: + case "project": + return scenario.key(team_id=team, project_id=scenario.project(team, **attached)) + case "key": + return scenario.key(team_id=team, **attached) + case _: + return scenario.key(team_id=team) + + +def _send(gateway: Gateway, route: str, body: Mapping[str, JsonValue], key: str) -> httpx.Response: + return gateway.request("POST", route, body, key=key, headers={"x-goog-api-key": key}) + + +def _drain(rig: AuthGuardrailRig) -> None: + rig.vendor.drain() + rig.post_call_vendor.drain() + rig.key_post_call_vendor.drain() + rig.generic_post_call_vendor.drain() + rig.model_post_call_vendor.drain() + for upstream in rig.upstreams.values(): + upstream.drain() + + +@pytest.mark.parametrize("level", ["key", "team", "project"]) +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_auth_guardrail_blocks_the_prompt_before_the_provider(rig: AuthGuardrailRig, route: str, level: Level) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_guardrail(scenario, level) + _drain(rig) + + response: Final = _send(gateway, route, _route_body(route, BLOCK_MARKER), key) + + assert response.status_code == 400, response.text + assert [BLOCK_MARKER.encode() in request.body for request in rig.vendor.drain()] == [True] + assert rig.upstreams[route].drain() == () + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_auth_guardrail_lets_a_clean_prompt_reach_the_provider(rig: AuthGuardrailRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_guardrail(scenario, "key") + marker: Final = f"clean-{uuid.uuid4().hex}" + _drain(rig) + + response: Final = _send(gateway, route, _route_body(route, marker), key) + + assert response.status_code == 200, response.text + assert [marker.encode() in request.body for request in rig.vendor.drain()] == [True] + assert [marker.encode() in request.body for request in rig.upstreams[route].drain()] == [True] + + +def test_empty_request_guardrails_list_does_not_drop_the_key_guardrail(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_guardrail(scenario, "key") + _drain(rig) + + response: Final = gateway.request( + "POST", OPENAI_ROUTE, {**_route_body(OPENAI_ROUTE, BLOCK_MARKER), "guardrails": []}, key=key + ) + + assert response.status_code == 400, response.text + assert rig.upstreams[OPENAI_ROUTE].drain() == () + + +def test_request_without_a_body_is_not_sent_to_the_guardrail(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_guardrail(scenario, "key") + _drain(rig) + + response: Final = gateway.request("GET", "/openai/v1/models", key=key) + + assert response.status_code == 200, response.text + assert rig.vendor.drain() == () + assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 + + +def test_pre_call_guardrail_relays_the_provider_reply_untouched(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_guardrail(scenario, "key") + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _route_body(OPENAI_ROUTE, "What is the capital of France?"), key) + + assert response.status_code == 200, response.text + assert response.content == _OPENAI_BODY + assert len(rig.vendor.drain()) == 1 + assert rig.post_call_vendor.drain() == () + + +def test_key_post_call_guardrail_runs_on_a_model_with_its_own_guardrails(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [KEY_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, + OPENAI_ROUTE, + {**_route_body(OPENAI_ROUTE, "What is the capital of France?"), "model": GUARDED_MODEL}, + key, + ) + + assert response.status_code == 400, response.text + assert len(rig.key_post_call_vendor.drain()) == 1 + assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 + + +@pytest.mark.parametrize("route", [ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_pre_call_guardrail_relays_every_provider_reply_untouched(rig: AuthGuardrailRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_guardrail(scenario, "team") + _drain(rig) + + response: Final = _send(gateway, route, _route_body(route, "What is the capital of France?"), key) + + assert response.status_code == 200, response.text + assert response.content == _UPSTREAM_BODIES[route] + assert len(rig.vendor.drain()) == 1 + assert rig.post_call_vendor.drain() == () + + +def test_custom_pass_through_endpoint_keeps_its_opt_in(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key( + team_id=scenario.team(), + metadata={"guardrails": [GUARDRAIL], "allowed_passthrough_routes": [CUSTOM_ROUTE]}, + ) + _drain(rig) + + response: Final = gateway.request("POST", CUSTOM_ROUTE, _route_body(OPENAI_ROUTE, BLOCK_MARKER), key=key) + + assert response.status_code == 200, response.text + assert rig.vendor.drain() == () + assert [BLOCK_MARKER.encode() in request.body for request in rig.upstreams[OPENAI_ROUTE].drain()] == [True] + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_key_apply_guardrail_post_call_scans_the_provider_reply(rig: AuthGuardrailRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]}) + _drain(rig) + + response: Final = _send(gateway, route, _route_body(route, "What is the capital of France?"), key) + + assert response.status_code == 200, response.text + assert [b'"Hi' in request.body for request in rig.generic_post_call_vendor.drain()] == [True] + + +def test_post_call_guardrail_does_not_read_a_reply_timestamp_as_a_phone_number(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [PHONE_POST_CALL]}) + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _route_body(OPENAI_ROUTE, "What is the capital of France?"), key) + + assert response.status_code == 200, response.text + assert response.json()["created"] == 1793600123 + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_guardrail_the_client_names_in_metadata_still_runs_for_a_guarded_key(rig: AuthGuardrailRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, route, {**_route_body(route, BLOCK_MARKER), "metadata": {"guardrails": [GUARDRAIL]}}, key + ) + + assert response.status_code == 400, response.text + assert [BLOCK_MARKER.encode() in request.body for request in rig.vendor.drain()] == [True] + assert rig.upstreams[route].drain() == () + + +@pytest.mark.parametrize("settings", [{}, True]) +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_guardrail_the_client_names_in_the_dict_form_still_runs_for_a_guarded_key( + rig: AuthGuardrailRig, route: str, settings: JsonValue +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, route, {**_route_body(route, BLOCK_MARKER), "metadata": {"guardrails": {GUARDRAIL: settings}}}, key + ) + + assert response.status_code == 400, response.text + assert [BLOCK_MARKER.encode() in request.body for request in rig.vendor.drain()] == [True] + assert rig.upstreams[route].drain() == () + + +def _key_with_opt_out(scenario: Scenario, opt_out: str) -> str: + guarded: Final[dict[str, JsonValue]] = {"guardrails": [GENERIC_POST_CALL]} + match opt_out: + case "key_disable": + return scenario.key(metadata=guarded, disable_global_guardrails=True) + case "key_opted_out": + return scenario.key(metadata={**guarded, "opted_out_global_guardrails": [GLOBAL_POST_CALL]}) + case "team_disable": + return scenario.key(team_id=scenario.team(disable_global_guardrails=True), metadata=guarded) + case "team_opted_out": + return scenario.key( + team_id=scenario.team(metadata={"opted_out_global_guardrails": [GLOBAL_POST_CALL]}), metadata=guarded + ) + case _: + return scenario.key(metadata=guarded) + + +@pytest.mark.parametrize( + ("opt_out", "global_vendor_calls"), + [("key_disable", 0), ("key_opted_out", 0), ("team_disable", 0), ("team_opted_out", 0), ("none", 1)], +) +@pytest.mark.parametrize("route", [OPENAI_ROUTE, "/v1/chat/completions"]) +def test_post_call_pass_honors_a_global_guardrail_opt_out_as_the_native_route_does( + rig: AuthGuardrailRig, route: str, opt_out: str, global_vendor_calls: int +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = _key_with_opt_out(scenario, opt_out) + _drain(rig) + + response: Final = _send( + gateway, + route, + {**_route_body(OPENAI_ROUTE, "What is the capital of France?"), "model": GUARDED_MODEL}, + key, + ) + + assert response.status_code == 200, response.text + assert len(rig.generic_post_call_vendor.drain()) == 1 + assert len(rig.post_call_vendor.drain()) == global_vendor_calls + + +def test_model_guardrail_does_not_run_on_pass_through_traffic_that_names_the_model(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, + OPENAI_ROUTE, + {**_route_body(OPENAI_ROUTE, "What is the capital of France?"), "model": POST_CALL_GUARDED_MODEL}, + key, + ) + + assert response.status_code == 200, response.text + assert len(rig.generic_post_call_vendor.drain()) == 1 + assert rig.model_post_call_vendor.drain() == () + + +def test_post_call_pass_leaves_a_logger_hook_for_native_responses_out(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, + OPENAI_ROUTE, + {**_route_body(OPENAI_ROUTE, "What is the capital of France?"), "model": RAISING_HOOK_MODEL}, + key, + ) + + assert response.status_code == 200, response.text + assert len(rig.generic_post_call_vendor.drain()) == 1 + assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 + + +def _charged_rows(gateway: Gateway, key: str) -> tuple[Mapping[str, object], ...]: + return tuple( + eventually( + lambda: gateway.rows( + 'SELECT status, spend, total_tokens FROM "LiteLLM_SpendLogs" WHERE api_key = %s', (key_hash(key),) + ), + bool, + ) + ) + + +def _tagged_body(route: str, tag: str) -> dict[str, JsonValue]: + return {**_route_body(route, "What is the capital of France?"), "metadata": {"tags": [tag]}} + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE]) +def test_post_call_policy_pipeline_blocks_a_flagged_pass_through_response(rig: AuthGuardrailRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key() + _drain(rig) + + response: Final = _send(gateway, route, _tagged_body(route, PIPELINE_BLOCKS), key) + + assert response.status_code == 400, response.text + assert len(rig.key_post_call_vendor.drain()) == 1 + assert len(rig.upstreams[route].drain()) == 1 + assert [row["status"] for row in _charged_rows(gateway, key)] == ["failure"] + + +def test_post_call_policy_pipeline_decides_what_a_flagged_step_does(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key() + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _tagged_body(OPENAI_ROUTE, PIPELINE_ALLOWS), key) + + assert response.status_code == 200, response.text + assert len(rig.key_post_call_vendor.drain()) == 1 + + +def test_pre_call_policy_pipeline_relays_the_provider_bytes(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key() + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _tagged_body(OPENAI_ROUTE, PIPELINE_BEFORE_THE_CALL), key) + + assert response.status_code == 200, response.text + assert response.content == _OPENAI_BODY + assert len(rig.vendor.drain()) == 1 + + +def test_lakera_post_call_has_nothing_to_read_on_a_gemini_body(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [KEY_POST_CALL]}) + _drain(rig) + + response: Final = _send(gateway, GEMINI_ROUTE, _route_body(GEMINI_ROUTE, "What is the capital of France?"), key) + + assert response.status_code == 200, response.text + assert rig.key_post_call_vendor.drain() == () + assert len(rig.upstreams[GEMINI_ROUTE].drain()) == 1 + + +def test_post_call_policy_pipeline_lakera_step_has_nothing_to_read_on_a_gemini_body(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key() + _drain(rig) + + response: Final = _send(gateway, GEMINI_ROUTE, _tagged_body(GEMINI_ROUTE, PIPELINE_BLOCKS), key) + + assert response.status_code == 200, response.text + assert rig.key_post_call_vendor.drain() == () + assert len(rig.upstreams[GEMINI_ROUTE].drain()) == 1 + + +@pytest.mark.parametrize(("route", "guardrail"), [(ANTHROPIC_ROUTE, KEY_POST_CALL), (GEMINI_ROUTE, BLOCKING_POST_CALL)]) +def test_response_a_post_call_guardrail_blocks_still_charges_the_provider_call( + rig: AuthGuardrailRig, route: str, guardrail: str +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [guardrail]}) + _drain(rig) + + response: Final = _send(gateway, route, _route_body(route, "What is the capital of France?"), key) + + assert response.status_code == 400, response.text + assert len(rig.upstreams[route].drain()) == 1 + rows: Final = _charged_rows(gateway, key) + assert [row["status"] for row in rows] == ["failure"], rows + assert float(str(rows[0]["spend"])) > 0, rows + assert int(str(rows[0]["total_tokens"])) > 0, rows + + +def _spend_logs_metadata_of_the_failure_row(gateway: Gateway, key: str) -> JsonValue: + rows: Final = eventually( + lambda: gateway.rows('SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s', (key_hash(key),)), + bool, + ) + assert [row["status"] for row in rows] == ["failure"], rows + return _JSON.validate_python(rows[0]["metadata"]).get("spend_logs_metadata") + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +def test_a_blocked_response_keeps_the_spend_logs_metadata_on_its_row(rig: AuthGuardrailRig, route: str) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"spend_logs_metadata": {"team_field": "team"}}) + key: Final = scenario.key( + team_id=team, + metadata={"guardrails": [BLOCKING_POST_CALL], "spend_logs_metadata": {"key_field": "key"}}, + ) + _drain(rig) + body: Final = { + **_route_body(route, "What is the capital of France?"), + "metadata": {"spend_logs_metadata": {"body_field": "body"}}, + } + + response: Final = _send(gateway, route, body, key) + + assert response.status_code == 400, response.text + assert _spend_logs_metadata_of_the_failure_row(gateway, key) == { + "body_field": "body", + "key_field": "key", + "team_field": "team", + } + + +def test_a_response_a_pipeline_replaces_keeps_the_spend_logs_metadata_on_its_row(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"spend_logs_metadata": {"key_field": "key"}}) + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, _tagged_body(OPENAI_ROUTE, PIPELINE_MODIFIES), key) + + assert response.status_code == 200, response.text + assert b"content_filter" in response.content, response.text + assert _spend_logs_metadata_of_the_failure_row(gateway, key) == {"key_field": "key"} + + +def test_a_custom_endpoints_blocked_response_row_gets_no_spend_logs_metadata(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key( + team_id=scenario.team(), + metadata={ + "allowed_passthrough_routes": [CUSTOM_GUARDED_ROUTE], + "spend_logs_metadata": {"key_field": "key"}, + }, + ) + _drain(rig) + + response: Final = gateway.request("POST", CUSTOM_GUARDED_ROUTE, _route_body(OPENAI_ROUTE, "Hello"), key=key) + + assert response.status_code == 400, response.text + assert _spend_logs_metadata_of_the_failure_row(gateway, key) is None + + +def _row_fields(gateway: Gateway, key: str) -> tuple[Mapping[str, object], ...]: + return tuple( + eventually( + lambda: gateway.rows( + 'SELECT status, spend, total_tokens, api_base, request_id FROM "LiteLLM_SpendLogs" WHERE api_key = %s', + (key_hash(key),), + ), + bool, + ) + ) + + +@pytest.mark.parametrize("route", [OPENAI_ROUTE, ANTHROPIC_ROUTE, GEMINI_ROUTE]) +@pytest.mark.parametrize("value", [None, False, 7, "not-a-list"]) +def test_a_body_guardrails_that_is_not_a_list_is_dropped_for_a_guarded_key( + rig: AuthGuardrailRig, route: str, value: JsonValue +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [GUARDRAIL]}) + _drain(rig) + + response: Final = _send(gateway, route, {**_route_body(route, "Hello"), "guardrails": value}, key) + + assert response.status_code == 200, response.text + assert len(rig.vendor.drain()) == 1 + assert [b'"guardrails"' in request.body for request in rig.upstreams[route].drain()] == [False] + + +def test_a_blocked_response_with_no_cost_result_keeps_its_usage_on_the_failure_row(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [BLOCKING_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, ANTHROPIC_ROUTE, {**_route_body(ANTHROPIC_ROUTE, "Hello"), "model": UNPRICED_MODEL}, key + ) + + assert response.status_code == 400, response.text + assert len(rig.upstreams[ANTHROPIC_ROUTE].drain()) == 1 + rows: Final = _row_fields(gateway, key) + assert [row["status"] for row in rows] == ["failure"], rows + assert float(str(rows[0]["spend"])) == 0, rows + assert int(str(rows[0]["total_tokens"])) > 0, rows + assert rows[0]["api_base"], rows + + +def test_a_tpm_limit_counts_blocked_responses_with_no_cost_result(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [BLOCKING_POST_CALL]}, tpm_limit=600) + _drain(rig) + body: Final = {**_route_body(ANTHROPIC_ROUTE, "Hello"), "model": UNPRICED_MODEL} + + limited: Final = eventually( + lambda: _send(gateway, ANTHROPIC_ROUTE, body, key), + lambda response: response.status_code != 400, + seconds=70, + ) + + assert limited.status_code == 429, limited.text + + +def test_a_tag_scoped_post_call_guardrail_scans_a_reply_for_a_request_with_its_tag(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [TAG_SCOPED_POST_CALL]}) + _drain(rig) + + response: Final = _send( + gateway, + OPENAI_ROUTE, + {**_route_body(OPENAI_ROUTE, "Hello"), "metadata": {"tags": [POST_CALL_TAG]}}, + key, + ) + + assert response.status_code == 400, response.text + assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 + + +def test_custom_endpoint_writes_no_charge_for_a_response_its_post_call_guardrail_blocks(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key( + team_id=scenario.team(), metadata={"allowed_passthrough_routes": [CUSTOM_GUARDED_ROUTE]} + ) + _drain(rig) + + response: Final = gateway.request("POST", CUSTOM_GUARDED_ROUTE, _route_body(OPENAI_ROUTE, "Hello"), key=key) + + assert response.status_code == 400, response.text + assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 + assert [float(str(row["spend"])) for row in _row_fields(gateway, key)] == [0], response.text + + +def test_a_custom_endpoint_with_only_a_pre_call_guardrail_never_calls_a_global_unified_guardrail( + rig: AuthGuardrailRig, +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + custom_key: Final = scenario.key(metadata={"allowed_passthrough_routes": [CUSTOM_PRE_CALL_ROUTE]}) + built_in_key: Final = scenario.key(metadata={"guardrails": [GENERIC_POST_CALL]}) + _drain(rig) + rig.global_unified_vendor.drain() + + custom_response: Final = gateway.request( + "POST", CUSTOM_PRE_CALL_ROUTE, _route_body(OPENAI_ROUTE, "Hello"), key=custom_key + ) + + assert custom_response.status_code == 200, custom_response.text + assert len(rig.vendor.drain()) == 1 + assert rig.global_unified_vendor.drain() == () + + built_in_response: Final = _send(gateway, OPENAI_ROUTE, _route_body(OPENAI_ROUTE, "Hello"), built_in_key) + + assert built_in_response.status_code == 200, built_in_response.text + assert len(rig.global_unified_vendor.drain()) == 1 + + +@pytest.mark.parametrize( + "field", ["Hello there", 7, ["a", "b"], {"guardrails": 7}], ids=["text", "int", "list", "dict"] +) +def test_a_custom_endpoints_own_body_field_still_gets_its_reply_scanned( + rig: AuthGuardrailRig, field: JsonValue +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"allowed_passthrough_routes": [CUSTOM_GUARDED_ROUTE]}) + + response: Final = gateway.request( + "POST", CUSTOM_GUARDED_ROUTE, {**_route_body(OPENAI_ROUTE, "Hello"), "body": field}, key=key + ) + + assert response.status_code == 400, response.text + assert "synthetic block" in response.text + + +@pytest.mark.parametrize( + ("guardrail", "fields"), + [ + (GUARDRAIL, {"proxy_server_request": 5}), + (GUARDRAIL, {"metadata": {"user_api_key_alias": 5}}), + (TAG_SCOPED_POST_CALL, {"metadata": {"tags": [1, {"a": 1}]}}), + ( + None, + { + "metadata": {"_pipeline_managed_guardrails": []}, + "litellm_metadata": {"_guardrail_pipelines": [["p", "q"]]}, + }, + ), + ], + ids=["proxy-server-request", "key-alias", "tags", "pipeline-slots"], +) +def test_a_malformed_proxy_field_in_a_body_still_gets_the_provider_reply( + rig: AuthGuardrailRig, guardrail: str | None, fields: Mapping[str, JsonValue] +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [guardrail]} if guardrail else {}) + _drain(rig) + + response: Final = _send(gateway, OPENAI_ROUTE, {**_route_body(OPENAI_ROUTE, "Hello"), **fields}, key) + + assert response.status_code == 200, response.text + assert len(rig.upstreams[OPENAI_ROUTE].drain()) == 1 + + +def test_a_monitored_lakera_flag_keeps_the_provider_id_on_the_row(rig: AuthGuardrailRig) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"guardrails": [MONITOR_POST_CALL]}) + _drain(rig) + + response: Final = _send(gateway, ANTHROPIC_ROUTE, _route_body(ANTHROPIC_ROUTE, MONITORED_PROMPT), key) + + assert response.status_code == 200, response.text + assert response.content == _MONITORED_ANTHROPIC_BODY + assert [row["request_id"] for row in _row_fields(gateway, key)] == [MONITORED_REPLY_ID] + + +@pytest.mark.parametrize( + ("policy", "forwarded"), [(PIPELINE_ALLOWS, True), (PIPELINE_BEFORE_THE_CALL, False)], ids=["post-call", "pre-call"] +) +def test_only_a_pipeline_that_runs_before_the_provider_withholds_the_client_user_id( + rig: AuthGuardrailRig, policy: str, forwarded: bool +) -> None: + gateway: Final = rig.owned.gateway + with gateway.scenario() as scenario: + key: Final = scenario.key() + _drain(rig) + body: Final = { + **_route_body(ANTHROPIC_ROUTE, "Hello"), + "metadata": {"tags": [policy], "user_id": "client-user-1"}, + } + + response: Final = _send(gateway, ANTHROPIC_ROUTE, body, key) + + assert response.status_code == 200, response.text + assert [b"client-user-1" in request.body for request in rig.upstreams[ANTHROPIC_ROUTE].drain()] == [forwarded] diff --git a/tests/unit/enterprise/integrations/test_custom_guardrail.py b/tests/unit/enterprise/integrations/test_custom_guardrail.py index 84f34144d99..6c6fb13ff59 100644 --- a/tests/unit/enterprise/integrations/test_custom_guardrail.py +++ b/tests/unit/enterprise/integrations/test_custom_guardrail.py @@ -330,3 +330,29 @@ async def test_baggage_never_changes_the_guardrail_mode_key_team_and_deployment_ assert await runs_post_call_guardrail(()) is runs_post_call assert await runs_post_call_guardrail(("usecase:dev",)) is runs_post_call + + +@pytest.mark.parametrize("bucket", ["metadata", "litellm_metadata"]) +def test_a_request_tag_that_cant_match_falls_back_to_the_default_mode(monkeypatch, bucket): + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + cg = CustomGuardrail( + guardrail_name="test_guardrail", + supported_event_hooks=[GuardrailEventHooks.pre_call, GuardrailEventHooks.post_call], + event_hook=Mode(tags={"test_tag": "post_call"}, default="pre_call"), + default_on=True, + ) + + assert ( + cg.should_run_guardrail( + data={bucket: {"tags": [1, {"a": 1}, ["b"]]}}, + event_type=GuardrailEventHooks.pre_call, + ) + is True + ) + assert ( + cg.should_run_guardrail( + data={bucket: {"tags": [{"a": 1}, "test_tag"]}}, + event_type=GuardrailEventHooks.pre_call, + ) + is False + ) diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index c32c1b71f75..1e93db413b1 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -5,7 +5,7 @@ import json import pickle import threading from typing import TYPE_CHECKING, ClassVar, Final, Literal, Optional -from unittest.mock import AsyncMock +from unittest.mock import AsyncMock, patch import pytest @@ -3802,6 +3802,30 @@ def test_default_on_guardrail() -> None: ) +def test_merged_guardrail_entries_dedupes_an_all_plain_list(): + from litellm.integrations.custom_guardrail import merged_guardrail_entries + + names: Final = [f"g{index % 500}" for index in range(1000)] + + assert merged_guardrail_entries(names) == [f"g{index}" for index in range(500)] + + +def test_merged_guardrail_entries_reads_each_entry_of_a_mixed_list_a_fixed_number_of_times(): + from litellm.integrations import custom_guardrail + + entries: Final = [*(f"g{index}" for index in range(1000)), {"g1": {"threshold": 1}}] + + with patch.object( + custom_guardrail, "guardrail_entry_names", wraps=custom_guardrail.guardrail_entry_names + ) as entry_names: + merged: Final = custom_guardrail.merged_guardrail_entries(entries) + + assert "g1" not in merged + assert merged[-1] == {"g1": {"threshold": 1}} + assert len(merged) == 1000 + assert entry_names.call_count <= 2 * len(entries) + + def test_merged_guardrail_entries_names_each_guardrail_once_preferring_its_first_entry_with_settings(): from litellm.integrations.custom_guardrail import merged_guardrail_entries diff --git a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py index 1c0ce5a334a..ab3cf63cfe0 100644 --- a/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py +++ b/tests/unit/llms/bedrock/passthrough/guardrail_translation/test_handler.py @@ -832,6 +832,28 @@ class TestBedrockPassthroughGuardrailHandlerOutput: sent_texts = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"] assert "blocked model output" in sent_texts[0] + @pytest.mark.asyncio + async def test_non_converse_response_is_read_whole_with_its_envelope(self): + response = { + "id": "msg_marker", + "type": "message", + "role": "assistant", + "stop_reason": "end_turn", + "stop_sequence": "stop-marker", + "content": [{"type": "text", "text": "model output"}], + "usage": {"input_tokens": 5551234}, + } + guardrail = _make_guardrail({"texts": []}) + + await BedrockPassthroughGuardrailHandler().process_output_response( + response=response, + guardrail_to_apply=guardrail, + request_data={"endpoint": "model/anthropic.claude-3-sonnet/invoke"}, + ) + + sent_text = guardrail.apply_guardrail.call_args.kwargs["inputs"]["texts"][0] + assert all(part in sent_text for part in ("msg_marker", "stop-marker", "5551234", "model output")) + def _build_event_stream_frame(event_type: str, payload: dict) -> bytes: import json diff --git a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index 7cbabb95931..1ba79b704da 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -1,5 +1,6 @@ import json from typing import Final, Literal +from unittest.mock import AsyncMock, MagicMock import pytest from typing_extensions import ReadOnly, TypedDict @@ -226,3 +227,135 @@ def test_passthrough_route_block_before_any_invoke_frame_is_a_whole_message_in_c "message_stop", ] assert events[2]["delta"]["text"] == "blocked" + + +async def _scanned_response_texts(response: dict[str, object]) -> tuple[str, ...]: + guardrail: Final = MagicMock(guardrail_name="post-call", apply_guardrail=AsyncMock()) + await PassThroughEndpointHandler().process_output_response( + response, guardrail, MagicMock(passthrough_guardrails_config=None), request_data={} + ) + return tuple(text for call in guardrail.apply_guardrail.await_args_list for text in call.kwargs["inputs"]["texts"]) + + +@pytest.mark.asyncio +async def test_a_scan_with_no_logging_object_reads_the_response_as_json(): + response: Final = {"response": {"id": "chatcmpl-1", "object": "chat.completion", "choices": [{"text": "Paris."}]}} + guardrail: Final = MagicMock(guardrail_name="logging-only", apply_guardrail=AsyncMock()) + + await PassThroughEndpointHandler().process_output_response(response, guardrail, request_data={}) + + assert json.loads(guardrail.apply_guardrail.await_args.kwargs["inputs"]["texts"][0]) == response + + +@pytest.mark.asyncio +async def test_post_call_scan_reads_the_reply_text_and_not_the_numbers_around_it(): + response: Final = { + "id": "chatcmpl-1793600123", + "object": "chat.completion", + "created": 1793600123, + "model": "gpt-5.6", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "Paris."}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 4111111111111111, "total_tokens": 5551234567}, + } + + assert await _scanned_response_texts(response) == ("Paris.",) + + +@pytest.mark.asyncio +async def test_post_call_scan_still_reads_every_string_a_reply_carries(): + response: Final = { + "choices": [ + {"message": {"content": "call 555-123-4567", "tool_calls": [{"function": {"arguments": '{"a": "b"}'}}]}} + ], + "candidates": [{"content": {"parts": [{"text": "card 4111 1111 1111 1111"}]}}], + } + + assert await _scanned_response_texts(response) == ('call 555-123-4567\n{"a": "b"}\ncard 4111 1111 1111 1111',) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + {"candidates": [{"content": {"parts": [{"functionCall": {"args": {"id": "alice@example.com"}}}]}}]}, + {"content": [{"input": {"id": "alice@example.com", "message": {"role": "alice@example.com"}}}]}, + {"choices": [{"message": {"tool_calls": [{"function": {"nested": {"type": "alice@example.com"}}}]}}]}, + {"id": "alice@example.com", "result": {"status": "alice@example.com"}}, + ], +) +async def test_post_call_scan_reads_a_field_named_like_the_envelope_in_what_the_model_or_api_wrote(response): + assert all("alice@example.com" in text for text in await _scanned_response_texts(response)) + assert await _scanned_response_texts(response) != () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response, expected", + [ + ({"result": {"credit_card": 4111111111111111, "score": 0.4111111111111111, "ok": True}}, "4111111111111111"), + ( + {"candidates": [{"content": {"parts": [{"functionCall": {"args": {"card": 4111111111111111}}}]}}]}, + "4111111111111111", + ), + ({"choices": [{"message": {"tool_calls": [{"function": {"arguments": {"n": 5551234567}}}]}}]}, "5551234567"), + ({"usage": {}, "customer": {"phone": 5551234567}}, "5551234567"), + ({"output": {"card": 4111111111111111}}, "4111111111111111"), + ({"type": "message", "phone": 5551234567}, "5551234567"), + ], +) +async def test_post_call_scan_reads_an_integer_in_what_the_model_or_api_wrote(response, expected): + assert any(expected in text for text in await _scanned_response_texts(response)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + {"results": {"alice@example.com": "active"}}, + {"candidates": [{"content": {"parts": [{"functionCall": {"args": {"alice@example.com": "active"}}}]}}]}, + ], +) +async def test_post_call_scan_reads_a_field_name_in_what_the_model_or_api_wrote(response): + assert any("alice@example.com" in text for text in await _scanned_response_texts(response)) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "response", + [ + {"created": 1793600123, "data": [{"url": "https://x.example/a.png"}]}, + {"id": "batch_1", "object": "batch", "created_at": 1793600123, "request_counts": {"total": 5551234567}}, + {"predictions": [{"embeddings": {"values": [0.0123456789]}, "total": 5551234567}]}, + {"finish_reason": "COMPLETE", "message": {"role": "assistant"}, "usage": {"total": 5551234567}}, + ], +) +async def test_post_call_scan_leaves_out_the_numbers_of_a_provider_reply(response): + assert all( + "1793600123" not in text and "5551234567" not in text for text in await _scanned_response_texts(response) + ) + + +@pytest.mark.asyncio +async def test_post_call_scan_reads_only_the_text_of_an_anthropic_message(): + response: Final = { + "id": "msg_5551234567", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [ + {"type": "text", "text": "Paris."}, + {"type": "tool_use", "id": "toolu_5551234567", "name": "lookup", "input": {"card": 4111111111111111}}, + ], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5551234567, "output_tokens": 1}, + } + + assert await _scanned_response_texts(response) == ("Paris.\nlookup\ncard\n4111111111111111",) + + +@pytest.mark.asyncio +async def test_post_call_scan_skips_a_reply_with_no_text(): + response: Final = {"object": "list", "data": [{"object": "embedding", "embedding": [0.4111111111111111]}]} + + assert await _scanned_response_texts(response) == () diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py index b5f794cb136..c09eb4ed5c5 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_generic_guardrail_api.py @@ -2692,3 +2692,63 @@ async def test_during_mcp_call_blocked_raises() -> None: ) assert exc_info.value.message == "credential detected in tool input" + + +@pytest.mark.parametrize("value", (None, 0, False, "", [], "text", 7)) +def test_extract_user_api_key_metadata_reads_past_a_litellm_metadata_that_is_not_an_object(value: object) -> None: + guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic") + + metadata: Final = guardrail._extract_user_api_key_metadata( + {"metadata": {"user_api_key_hash": "hash"}, "litellm_metadata": value} + ) + + assert metadata["user_api_key_hash"] == "hash" + + +def test_extract_user_api_key_metadata_leaves_out_a_value_that_isnt_a_string() -> None: + guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic") + + metadata: Final = guardrail._extract_user_api_key_metadata( + {"metadata": {"user_api_key_alias": 5, "user_api_key_token": 7, "user_api_key_team_id": "team"}} + ) + + assert metadata == {"user_api_key_team_id": "team"} + + +async def _vendor_payload(request_data: dict[str, object], inputs: dict[str, object]) -> dict[str, object]: + guardrail: Final = GenericGuardrailAPI(api_base="http://vendor.invalid", guardrail_name="generic") + vendor_reply: Final = MagicMock(json=MagicMock(return_value={"action": "NONE"})) + + with patch.object(guardrail.async_handler, "post", return_value=vendor_reply) as post: + await guardrail.apply_guardrail(inputs=inputs, request_data=request_data, input_type="response") + + return post.call_args.kwargs["json"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "body", + ( + "Hello there", + 7, + True, + ["a", "b"], + {"guardrails": 7}, + {"metadata": {"guardrails": 7}}, + {"guardrails": [{"generic": 7}]}, + ), +) +async def test_a_body_field_a_custom_endpoints_client_sent_asks_the_vendor_for_nothing(body: object) -> None: + payload: Final = await _vendor_payload({"body": body}, {"texts": ["reply"]}) + + assert payload["texts"] == ["reply"] + assert not payload.get("additional_provider_specific_params") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ({"id": 1}, ["gpt"], 7)) +async def test_a_reply_model_that_isnt_a_string_is_left_out_of_the_vendor_request(model: object) -> None: + payload: Final = await _vendor_payload({}, {"texts": ["reply"], "model": model}) + + assert payload["texts"] == ["reply"] + assert payload.get("model") is None diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py new file mode 100644 index 00000000000..bacb576d6c9 --- /dev/null +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai.py @@ -0,0 +1,65 @@ +"""The legacy Lakera guardrail reads client JSON that request setup hasn't normalized on pass-through routes.""" + +import json +from typing import Final +from unittest.mock import AsyncMock, patch + +import httpx +import pytest + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.lakera_ai import lakeraAI_Moderation + +VENDOR: Final = "http://vendor.invalid" +CLEAN: Final = httpx.Response(200, json={"results": [{"flagged": False}]}, request=httpx.Request("POST", VENDOR)) + + +async def _sent_to_lakera(data: dict[str, object], call_type: str = "pass_through_endpoint") -> tuple[str, ...]: + guardrail: Final = lakeraAI_Moderation(api_key="key", api_base=VENDOR) + with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=CLEAN)) as post: + await guardrail._check(data=data, user_api_key_dict=UserAPIKeyAuth(api_key="key"), call_type=call_type) + return tuple(call.kwargs["data"] for call in post.await_args_list) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("token_arrays", [[1, 2, 3], [[1, 2], [3, 4]]]) +async def test_token_arrays_have_no_text_to_send(token_arrays: list[object]) -> None: + assert await _sent_to_lakera({"input": token_arrays}) == () + + +@pytest.mark.asyncio +async def test_string_entries_of_a_list_input_are_sent() -> None: + sent: Final = await _sent_to_lakera({"input": ["hello", "there"]}) + + assert [json.loads(body)["input"] for body in sent] == ["hello\nthere"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "content", + ["card 4111", [{"type": "input_text", "text": "card 4111"}, {"type": "input_image", "image_url": "http://x"}]], +) +async def test_a_responses_message_in_a_list_input_is_sent(content: object) -> None: + sent: Final = await _sent_to_lakera({"input": [{"role": "user", "content": content}]}) + + assert [json.loads(body)["input"] for body in sent] == ["card 4111"] + + +@pytest.mark.asyncio +async def test_a_pass_through_body_with_no_messages_or_input_sends_nothing() -> None: + contents: Final = {"contents": [{"role": "user", "parts": [{"text": "hello"}]}]} + + assert await _sent_to_lakera(contents) == () + + +@pytest.mark.asyncio +async def test_a_native_body_with_no_messages_or_input_still_reaches_lakera() -> None: + assert await _sent_to_lakera({"prompt": "hello"}, call_type="text_completion") == ("",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("strings", "text"), [(["a", "", "b"], "a\n\nb"), (["a", ""], "a\n")]) +async def test_empty_strings_of_a_list_input_keep_their_place(strings: list[object], text: str) -> None: + sent: Final = await _sent_to_lakera({"input": strings}) + + assert [json.loads(body)["input"] for body in sent] == [text] diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py index 23288a6136f..52c3c9eb92d 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_lakera_ai_v2.py @@ -2413,6 +2413,35 @@ async def test_lakera_post_call_allows_clean_content(): assert result is llm_response +@pytest.mark.asyncio +async def test_post_call_with_no_messages_and_no_reply_text_returns_the_response(): + guardrail: Final = LakeraAIGuardrail(api_key="test_key") + response: Final = {"candidates": [{"content": {"role": "model", "parts": [{"text": "hi"}]}}]} + rejected: Final = HTTPException(status_code=400, detail="either content, messages, or tools must be provided") + with patch.object(guardrail, "call_v2_guard", new=AsyncMock(side_effect=rejected)): + result: Final = await guardrail.async_post_call_success_hook( + data={"contents": [{"role": "user", "parts": [{"text": "hello"}]}]}, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + response=response, + ) + assert result == response + + +@pytest.mark.asyncio +async def test_post_call_with_messages_still_reaches_lakera(): + guardrail: Final = LakeraAIGuardrail(api_key="test_key") + rejected: Final = HTTPException(status_code=400, detail="rejected") + with ( + patch.object(guardrail, "call_v2_guard", new=AsyncMock(side_effect=rejected)), + pytest.raises(HTTPException), + ): + await guardrail.async_post_call_success_hook( + data={"messages": [{"role": "user", "content": "hello"}]}, + user_api_key_dict=UserAPIKeyAuth(api_key="test_key"), + response={}, + ) + + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_lakera_post_call_masks_pii_and_allows(): diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py index b145b990539..8d430d83787 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_noma_v2.py @@ -1,7 +1,9 @@ import asyncio import os +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from litellm.proxy.guardrails.guardrail_hooks.noma import NomaV2Guardrail @@ -747,3 +749,21 @@ class TestNomaV2ApplicationIdResolution: payload = call_mock.call_args.kwargs["payload"] assert "application_id" not in payload + + +@pytest.mark.asyncio +@pytest.mark.parametrize("value", (None, 0, False, "", [], "text", 7)) +async def test_noma_v2_scans_a_request_whose_litellm_metadata_is_not_an_object(value: object) -> None: + guardrail: Final = NomaV2Guardrail(api_key="key", api_base="http://vendor.invalid", guardrail_name="noma") + verdict: Final = httpx.Response( + 200, json={"action": "NONE"}, request=httpx.Request("POST", "http://vendor.invalid") + ) + + with patch.object(guardrail.async_handler, "post", new=AsyncMock(return_value=verdict)): + result: Final = await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data={"metadata": {}, "litellm_metadata": value}, + input_type="request", + ) + + assert result["texts"] == ["hello"] diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 11b07472eaa..9410cbab87b 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -344,6 +344,7 @@ class TestVertexAIPassThroughHandler: custom_headers={"Authorization": f"Bearer {test_token}"}, is_streaming_request=False, upstream_headers_filter=mock.ANY, + inherit_auth_guardrails=True, ) @pytest.mark.asyncio @@ -444,6 +445,7 @@ class TestVertexAIPassThroughHandler: custom_headers={"Authorization": f"Bearer {test_token}"}, is_streaming_request=False, upstream_headers_filter=mock.ANY, + inherit_auth_guardrails=True, ) @pytest.mark.parametrize( @@ -540,6 +542,7 @@ class TestVertexAIPassThroughHandler: custom_headers={"Authorization": f"Bearer {default_credentials}"}, is_streaming_request=False, upstream_headers_filter=mock.ANY, + inherit_auth_guardrails=True, ) @pytest.mark.asyncio diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index cd6c1ff11b5..2d8ba198cb0 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -56,6 +56,11 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( HttpPassThroughEndpointHelpers, InitPassThroughEndpointHelpers, SafeRouteAdder, + _anthropic_user_id_to_forward, + _inherit_auth_guardrails, + _post_call_hook_data, + _post_call_selection_metadata, + _starts_post_call_pass, WebsocketFrameModelAllowed, WebsocketFrameModelGate, WebsocketFrameModelRefused, @@ -82,9 +87,10 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( ) from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_tracking_utils import get_spend_logs_id -from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token +from litellm.proxy.utils import InternalUsageCache, ProxyLogging, _merge_guardrails_with_existing, hash_token from litellm.types import utils as types_utils from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep from litellm.types.utils import CallTypes from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, @@ -13242,3 +13248,495 @@ async def test_passthrough_header_and_callback_share_finalized_cost( finally: cleanup() await client.aclose() + + +class TestInheritAuthGuardrails: + """Built-in provider routes run the caller's key, team and project guardrails on pass-through.""" + + @staticmethod + async def _run( + body: dict[str, object], + *, + inherit: bool, + user_api_key_dict: UserAPIKeyAuth, + guardrails: tuple[CustomGuardrail, ...] = (), + ) -> tuple[list[dict[str, object]], AsyncMock]: + _, pre_call_data, post_call = await TestInheritAuthGuardrails._run_with_response( + body, inherit=inherit, user_api_key_dict=user_api_key_dict, guardrails=guardrails + ) + return pre_call_data, post_call + + @staticmethod + async def _run_with_response( + body: dict[str, object], + *, + inherit: bool, + user_api_key_dict: UserAPIKeyAuth, + guardrails: tuple[CustomGuardrail, ...] = (), + ) -> tuple[Response, list[dict[str, object]], AsyncMock]: + pre_call_data: Final[list[dict[str, object]]] = [] + + async def _record_pre_call(**kwargs: object) -> object: + pre_call_data.append(json.loads(json.dumps(kwargs["data"], default=str))) + return kwargs["data"] + + proxy_logging: Final = MagicMock() + proxy_logging.pre_call_hook = AsyncMock(side_effect=_record_pre_call) + proxy_logging.post_call_success_hook = AsyncMock(return_value={"id": "chatcmpl-1"}) + proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) + upstream: Final = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + content=b'{\n "id":"chatcmpl-1"\n}', + request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions"), + ) + request: Final = MagicMock() + request.method = "POST" + request.query_params = {} + request.headers = MagicMock() + request.headers.copy.return_value = {} + module: Final = "litellm.proxy.pass_through_endpoints.pass_through_endpoints" + with ExitStack() as stack: + stack.enter_context( + patch( + f"{module}.HttpPassThroughEndpointHelpers.non_streaming_http_request_handler", + new_callable=AsyncMock, + return_value=upstream, + ) + ) + stack.enter_context(patch(f"{module}._is_streaming_response", return_value=False)) + stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging)) + stack.enter_context(patch("litellm.proxy.proxy_server.premium_user", True)) + stack.enter_context( + patch( + f"{module}.pass_through_endpoint_logging", + MagicMock(pass_through_async_success_handler=AsyncMock()), + ) + ) + stack.enter_context(patch(f"{module}.read_request_body", new_callable=AsyncMock, return_value=body)) + stack.enter_context(patch(f"{module}.safe_get_request_headers", return_value={})) + stack.enter_context(patch.object(litellm, "callbacks", list(guardrails))) + response: Final = await pass_through_request( + request=request, + target="https://api.openai.com/v1/chat/completions", + custom_headers={}, + user_api_key_dict=user_api_key_dict, + stream=False, + inherit_auth_guardrails=inherit, + ) + return response, pre_call_data, proxy_logging.post_call_success_hook + + @staticmethod + def _guardrail(name: str, event_hook: GuardrailEventHooks, *, default_on: bool = False) -> CustomGuardrail: + return CustomGuardrail(guardrail_name=name, event_hook=event_hook, default_on=default_on) + + @staticmethod + def _chat_body(**extra: object) -> dict[str, object]: + return {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}], **extra} + + @pytest.mark.asyncio + @pytest.mark.parametrize("value", [None, False, True, 7, "name"]) + @pytest.mark.parametrize( + ("key_metadata", "body_metadata"), + [({"guardrails": ["key-guard"]}, None), ({}, {"guardrails": ["body-guard"]}), ({}, None)], + ids=["key-guardrails", "metadata-guardrails", "no-guardrails"], + ) + async def test_a_top_level_guardrails_the_merge_cant_add_to_is_dropped(self, value, key_metadata, body_metadata): + body: Final[dict[str, object]] = { + **self._chat_body(guardrails=value), + **({"metadata": body_metadata} if body_metadata is not None else {}), + } + + with patch("litellm.proxy.proxy_server.premium_user", True): + await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test", metadata=key_metadata), None, True) + + assert "guardrails" not in body + assert (body.get("metadata", {}).get("guardrails") or []) == [ + *key_metadata.get("guardrails", ()), + *(body_metadata or {}).get("guardrails", ()), + ] + + @pytest.mark.asyncio + async def test_a_body_with_nothing_to_merge_gets_no_metadata(self): + body: Final[dict[str, object]] = self._chat_body() + + assert await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True) is None + + assert body == self._chat_body() + + @pytest.mark.asyncio + async def test_a_body_with_nothing_to_merge_keeps_the_metadata_it_sent(self): + body: Final[dict[str, object]] = self._chat_body(metadata={}) + + await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True) + + assert body["metadata"] == {} + + @pytest.mark.asyncio + @pytest.mark.parametrize("value", [5, "text", ["list"], None]) + async def test_a_proxy_server_request_that_isnt_an_object_never_reaches_the_guardrails(self, value): + body: Final[dict[str, object]] = self._chat_body(proxy_server_request=value) + + await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True) + + assert "proxy_server_request" not in body + + @pytest.mark.parametrize("slot", [[["p", "q"]], [7], 7, "text", [{"mode": "post_call"}]]) + def test_a_pipeline_slot_a_body_sent_does_not_start_the_post_call_pass(self, slot): + kwargs: Final = { + "litellm_params": { + "metadata": {"_guardrail_pipelines": slot, "_pipeline_managed_guardrails": [], "tags": ["t"]} + } + } + + selection: Final = _post_call_selection_metadata(kwargs, inherited=True) + + assert selection == {"_pipeline_managed_guardrails": [], "tags": ["t"]} + assert _starts_post_call_pass(None, selection) is False + + def test_a_pipeline_the_policy_engine_resolved_starts_the_post_call_pass(self): + pipeline: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="guard")]) + kwargs: Final = {"litellm_params": {"metadata": {"_guardrail_pipelines": [("p", pipeline)]}}} + + selection: Final = _post_call_selection_metadata(kwargs, inherited=True) + + assert selection == {"_guardrail_pipelines": [("p", pipeline)]} + assert _starts_post_call_pass(None, selection) is True + + @pytest.mark.parametrize( + ("event_hook", "inherited", "gets_the_logging_object"), + [ + (GuardrailEventHooks.pre_call, False, False), + (GuardrailEventHooks.post_call, False, True), + (GuardrailEventHooks.pre_call, True, True), + ], + ids=["custom-pre-call", "custom-post-call", "built-in-route"], + ) + def test_a_custom_endpoint_gets_the_logging_object_only_for_a_post_call_guardrail( + self, event_hook, inherited, gets_the_logging_object + ): + logging_obj: Final = MagicMock() + + with patch.object( + litellm.logging_callback_manager, + "get_custom_loggers_for_type", + return_value=[self._guardrail("guard", event_hook)], + ): + data: Final = _post_call_hook_data( + self._chat_body(), ["guard"], logging_obj, UserAPIKeyAuth(api_key="sk-test"), {}, inherited + ) + + assert data.get("litellm_logging_obj") is (logging_obj if gets_the_logging_object else None) + + @pytest.mark.asyncio + @pytest.mark.parametrize("value", [None, False, 7, "config", ["config"]]) + async def test_a_guardrail_config_the_merge_cant_add_to_is_dropped(self, value): + body: Final[dict[str, object]] = self._chat_body( + guardrail_config=value, metadata={"guardrail_config": {"kept": True}} + ) + + await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True) + + assert "guardrail_config" not in body + assert body["metadata"]["guardrail_config"] == {"kept": True} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("body_guardrails", "merged"), + [(["body-guard"], ["key-guard", "body-guard"]), ({"body-guard": {}}, ["key-guard", "body-guard"])], + ids=["list", "object"], + ) + async def test_a_top_level_guardrails_list_or_object_still_merges(self, body_guardrails, merged): + body: Final[dict[str, object]] = self._chat_body(guardrails=body_guardrails) + + with patch("litellm.proxy.proxy_server.premium_user", True): + await _inherit_auth_guardrails( + body, UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}), None, True + ) + + assert body["metadata"]["guardrails"] == merged + + @pytest.mark.parametrize(("mode", "forwarded"), [("post_call", "user-1"), ("pre_call", None)]) + def test_only_a_pipeline_that_runs_before_the_provider_withholds_the_user_id(self, mode, forwarded): + pipeline: Final = GuardrailPipeline(mode=mode, steps=[PipelineStep(guardrail="guard")]) + body: Final = {"model": "claude", "metadata": {"user_id": "user-1", "_guardrail_pipelines": [("p", pipeline)]}} + + with patch.object(litellm, "callbacks", []): + user_id: Final = _anthropic_user_id_to_forward( + body, + "https://api.anthropic.com/v1/messages", + ProxyLogging(user_api_key_cache=DualCache()).request_is_scanned_before_the_provider, + ) + + assert user_id == forwarded + + @pytest.mark.asyncio + async def test_key_team_and_project_guardrails_reach_pre_and_post_call(self): + auth: Final = UserAPIKeyAuth( + api_key="sk-test", + metadata={"guardrails": ["key-guard"]}, + team_metadata={"guardrails": ["team-guard"]}, + project_metadata={"guardrails": ["project-guard"]}, + ) + + pre_call, post_call = await self._run( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=( + self._guardrail("key-guard", GuardrailEventHooks.pre_call), + self._guardrail("team-guard", GuardrailEventHooks.post_call), + ), + ) + + assert sorted(pre_call[0]["metadata"]["guardrails"]) == ["key-guard", "project-guard", "team-guard"] + post_call.assert_awaited_once() + assert sorted(post_call.call_args.kwargs["data"]["metadata"]["guardrails"]) == [ + "key-guard", + "project-guard", + "team-guard", + ] + + @pytest.mark.asyncio + async def test_post_call_guardrails_stay_named_after_a_model_guardrail_merge(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + _, post_call = await self._run( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("key-guard", GuardrailEventHooks.post_call),), + ) + + merged: Final = _merge_guardrails_with_existing(post_call.call_args.kwargs["data"], ["model-guard"]) + assert merged["metadata"]["guardrails"] == ["key-guard", "model-guard"] + + @pytest.mark.asyncio + @pytest.mark.parametrize("settings", [{}, True]) + async def test_guardrail_a_client_names_in_the_dict_form_stays_after_the_keys(self, settings: object): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + pre_call, _ = await self._run( + {**self._chat_body(), "metadata": {"guardrails": {"client-guard": settings}}}, + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("key-guard", GuardrailEventHooks.pre_call),), + ) + + assert pre_call[0]["metadata"]["guardrails"] == ["key-guard", "client-guard"] + + @staticmethod + def _body_with_policy_pipeline(mode: str) -> dict[str, object]: + pipeline: Final = GuardrailPipeline(mode=mode, steps=[PipelineStep(guardrail="step-guard")]) + return { + **TestInheritAuthGuardrails._chat_body(), + "metadata": { + "_guardrail_pipelines": [("policy", pipeline)], + "_pipeline_managed_guardrails": {"step-guard"}, + }, + } + + @pytest.mark.asyncio + async def test_post_call_policy_pipeline_starts_the_post_call_pass_with_its_state(self): + _, post_call = await self._run( + self._body_with_policy_pipeline("post_call"), + inherit=True, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + post_call.assert_awaited_once() + metadata: Final = post_call.await_args.kwargs["data"]["metadata"] + assert metadata["_pipeline_managed_guardrails"] == {"step-guard"} + assert [policy for policy, _ in metadata["_guardrail_pipelines"]] == ["policy"] + + @pytest.mark.asyncio + @pytest.mark.parametrize(("mode", "inherit"), [("pre_call", True), ("post_call", False)]) + async def test_policy_pipeline_leaves_the_response_alone_without_a_post_call_run(self, mode: str, inherit: bool): + response, _, post_call = await self._run_with_response( + self._body_with_policy_pipeline(mode), + inherit=inherit, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + post_call.assert_not_awaited() + assert response.body == b'{\n "id":"chatcmpl-1"\n}' + + @pytest.mark.asyncio + async def test_dict_form_stays_as_sent_for_a_key_without_guardrails(self): + pre_call, post_call = await self._run( + {**self._chat_body(), "metadata": {"guardrails": {"client-guard": {}}}}, + inherit=True, + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + guardrails=(self._guardrail("client-guard", GuardrailEventHooks.post_call),), + ) + + assert pre_call[0]["metadata"]["guardrails"] == {"client-guard": {}} + post_call.assert_not_awaited() + + @pytest.mark.asyncio + async def test_post_call_pass_reads_opt_outs_from_the_key_and_team_not_the_body(self): + auth: Final = UserAPIKeyAuth( + api_key="sk-test", + metadata={"guardrails": ["key-guard"], "disable_global_guardrails": True}, + team_metadata={"opted_out_global_guardrails": ["global-post"]}, + ) + forged: Final = {"user_api_key_metadata": {"disable_global_guardrails": False}} + + _, post_call = await self._run( + {**self._chat_body(), "metadata": forged}, + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("key-guard", GuardrailEventHooks.post_call),), + ) + + metadata: Final = post_call.call_args.kwargs["data"]["metadata"] + assert metadata["user_api_key_metadata"]["disable_global_guardrails"] is True + assert metadata["user_api_key_team_metadata"] == {"opted_out_global_guardrails": ["global-post"]} + + @pytest.mark.asyncio + async def test_a_body_opt_out_in_litellm_metadata_does_not_stop_the_post_call_pass(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["global-post"]}) + + _, post_call = await self._run( + { + **self._chat_body(), + "litellm_metadata": {"user_api_key_metadata": {"disable_global_guardrails": True}}, + }, + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("global-post", GuardrailEventHooks.post_call, default_on=True),), + ) + + post_call.assert_awaited_once() + + @pytest.mark.asyncio + async def test_inherited_post_call_pass_runs_only_the_guardrails_the_request_selected(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + _, post_call = await self._run( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("key-guard", GuardrailEventHooks.post_call),), + ) + + assert post_call.call_args.kwargs["guardrails_only"] is True + + @pytest.mark.asyncio + async def test_named_global_guardrail_the_key_opted_out_of_does_not_start_the_post_call_pass(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"opted_out_global_guardrails": ["global-post"]}) + + response, _, post_call = await self._run_with_response( + {**self._chat_body(), "metadata": {"guardrails": ["global-post"]}}, + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("global-post", GuardrailEventHooks.post_call, default_on=True),), + ) + + post_call.assert_not_awaited() + assert response.body == b'{\n "id":"chatcmpl-1"\n}' + + @pytest.mark.asyncio + async def test_pre_call_only_guardrails_relay_the_response_untouched(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + response, pre_call, post_call = await self._run_with_response( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=( + self._guardrail("key-guard", GuardrailEventHooks.pre_call), + self._guardrail("global-post", GuardrailEventHooks.post_call, default_on=True), + ), + ) + + assert pre_call[0]["metadata"]["guardrails"] == ["key-guard"] + post_call.assert_not_awaited() + assert response.body == b'{\n "id":"chatcmpl-1"\n}' + + @pytest.mark.asyncio + async def test_a_named_default_on_post_call_guardrail_still_checks_the_response(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["global-post"]}) + + _, post_call = await self._run( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("global-post", GuardrailEventHooks.post_call, default_on=True),), + ) + + post_call.assert_awaited_once() + assert post_call.call_args.kwargs["data"]["metadata"]["guardrails"] == ["global-post"] + + @pytest.mark.asyncio + async def test_post_call_guardrails_get_the_request_logging_object(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-post"]}) + + _, post_call = await self._run( + self._chat_body(), + inherit=True, + user_api_key_dict=auth, + guardrails=(self._guardrail("key-post", GuardrailEventHooks.post_call),), + ) + + post_call.assert_awaited_once() + assert isinstance(post_call.call_args.kwargs["data"]["litellm_logging_obj"], LiteLLMLoggingObj) + + @pytest.mark.asyncio + async def test_request_body_guardrails_are_added_to_the_key_guardrails(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + pre_call, _ = await self._run( + self._chat_body(guardrails=["body-guard"]), inherit=True, user_api_key_dict=auth + ) + + assert pre_call[0]["metadata"]["guardrails"] == ["key-guard", "body-guard"] + assert "guardrails" not in pre_call[0] + + @pytest.mark.asyncio + async def test_guardrails_the_client_names_in_metadata_still_run_for_a_guarded_key(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + pre_call, post_call = await self._run( + self._chat_body(metadata={"guardrails": ["client-pre", "client-post"]}), + inherit=True, + user_api_key_dict=auth, + guardrails=( + self._guardrail("key-guard", GuardrailEventHooks.pre_call), + self._guardrail("client-post", GuardrailEventHooks.post_call), + ), + ) + + assert pre_call[0]["metadata"] == {"guardrails": ["key-guard", "client-pre", "client-post"]} + post_call.assert_awaited_once() + assert post_call.call_args.kwargs["data"]["metadata"]["guardrails"] == [ + "key-guard", + "client-pre", + "client-post", + ] + + @pytest.mark.asyncio + async def test_empty_request_body_guardrails_keep_the_key_guardrails(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + pre_call, _ = await self._run(self._chat_body(guardrails=[]), inherit=True, user_api_key_dict=auth) + + assert pre_call[0]["metadata"]["guardrails"] == ["key-guard"] + + @pytest.mark.asyncio + async def test_routes_without_inheritance_keep_key_guardrails_off(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + pre_call, post_call = await self._run(self._chat_body(), inherit=False, user_api_key_dict=auth) + + assert "guardrails" not in pre_call[0].get("metadata", {}) + post_call.assert_not_awaited() + + @pytest.mark.asyncio + async def test_request_without_a_body_gets_no_guardrails(self): + auth: Final = UserAPIKeyAuth(api_key="sk-test", metadata={"guardrails": ["key-guard"]}) + + pre_call, post_call = await self._run({}, inherit=True, user_api_key_dict=auth) + + assert "guardrails" not in pre_call[0].get("metadata", {}) + post_call.assert_not_awaited() diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py index 7b06905c8aa..667625cdb7d 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_post_call_guardrails.py @@ -18,6 +18,7 @@ from litellm.integrations.custom_guardrail import ( ModifyResponseException, ) from litellm.proxy._types import ProxyException +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import _run_post_call_guardrails_keeping_the_charge _PT_MOD = "litellm.proxy.pass_through_endpoints.pass_through_endpoints" _COLLECT = "litellm.proxy.pass_through_endpoints.passthrough_guardrails.PassthroughGuardrailHandler.collect_guardrails" @@ -130,6 +131,29 @@ class TestPassthroughPostCallGuardrails: call_kwargs = mock_proxy_logging.post_call_success_hook.call_args assert call_kwargs.kwargs["response"] == _GEMINI_RESPONSE + @patch(_COLLECT, return_value=["rubrik"]) + async def test_configured_endpoint_keeps_every_post_call_hook_and_reads_no_key_opt_outs( + self, + mock_collect, + ): + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_success_hook = AsyncMock(return_value=_GEMINI_RESPONSE) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) + + with _common_patches(mock_proxy_logging, _make_httpx_response(_GEMINI_RESPONSE)): + await pass_through_request( + request=_make_mock_request(), + target="https://example.com/v1/generateContent", + custom_headers={"Content-Type": "application/json"}, + user_api_key_dict=_make_user_api_key_dict(), + stream=False, + ) + + call_kwargs = mock_proxy_logging.post_call_success_hook.call_args.kwargs + assert call_kwargs["guardrails_only"] is False + assert call_kwargs["data"]["metadata"].keys() == {"guardrails"} + @patch(_COLLECT, return_value=[]) async def test_post_call_success_hook_skipped_when_no_guardrails( self, @@ -204,6 +228,54 @@ class TestPassthroughPostCallGuardrails: assert body["error"]["guardrail_name"] == "rubrik" assert body["error"]["model"] == "gemini-2.0-flash" + @patch(_COLLECT, return_value=["rubrik"]) + async def test_configured_endpoint_writes_no_charge_for_a_failed_post_call_hook(self, mock_collect): + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_success_hook = AsyncMock(side_effect=HTTPException(status_code=418, detail="no")) + mock_proxy_logging.post_call_failure_hook = AsyncMock() + + with ( + _common_patches(mock_proxy_logging, _make_httpx_response(_GEMINI_RESPONSE)), + patch(f"{_PT_MOD}._blocked_response_cost_logging") as settlement, + pytest.raises(ProxyException), + ): + await pass_through_request( + request=_make_mock_request(), + target="https://example.com/v1/generateContent", + custom_headers={"Content-Type": "application/json"}, + user_api_key_dict=_make_user_api_key_dict(), + stream=False, + ) + + settlement.pass_through_async_success_handler.assert_not_called() + + @patch(_COLLECT, return_value=["rubrik"]) + async def test_a_hook_result_that_is_not_a_dict_leaves_the_body_the_logs_read(self, mock_collect): + mock_proxy_logging = MagicMock() + mock_proxy_logging.pre_call_hook = AsyncMock(return_value={}) + mock_proxy_logging.post_call_success_hook = AsyncMock(return_value=MagicMock(id="rebuilt")) + mock_proxy_logging.post_call_response_headers_hook = AsyncMock(return_value={}) + success_handler = AsyncMock() + + with ( + _common_patches(mock_proxy_logging, _make_httpx_response(_GEMINI_RESPONSE)), + patch( + f"{_PT_MOD}.pass_through_endpoint_logging", + MagicMock(pass_through_async_success_handler=success_handler), + ), + ): + result = await pass_through_request( + request=_make_mock_request(), + target="https://example.com/v1/generateContent", + custom_headers={"Content-Type": "application/json"}, + user_api_key_dict=_make_user_api_key_dict(), + stream=False, + ) + + assert success_handler.await_args.kwargs["response_body"] == _GEMINI_RESPONSE + assert json.loads(result.body) == _GEMINI_RESPONSE + @patch(_COLLECT, return_value=["rubrik"]) async def test_deny_forwards_guardrail_logging_info_to_failure_hook( self, @@ -216,7 +288,7 @@ class TestPassthroughPostCallGuardrails: block path dropping it.""" mock_response = _make_httpx_response(_GEMINI_RESPONSE) - def _block(*, data, user_api_key_dict, response): + def _block(*, data, user_api_key_dict, response, guardrails_only): metadata = data.setdefault("metadata", {}) metadata.setdefault("standard_logging_guardrail_information", []).append( {"guardrail_name": "rubrik", "guardrail_status": "guardrail_intervened"} @@ -312,3 +384,37 @@ def test_modify_response_exception_importable_from_both_paths(): ) assert FromExceptions is FromGuardrail + + +@pytest.mark.asyncio +async def test_a_blocked_response_settles_the_provider_charge_and_keeps_the_block(): + blocked = HTTPException(status_code=400, detail="blocked") + settle = AsyncMock() + + with pytest.raises(HTTPException) as caught: + await _run_post_call_guardrails_keeping_the_charge(AsyncMock(side_effect=blocked), settle) + + assert caught.value is blocked + settle.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_a_failed_settlement_does_not_replace_the_block(): + blocked = HTTPException(status_code=400, detail="blocked") + + with pytest.raises(HTTPException) as caught: + await _run_post_call_guardrails_keeping_the_charge( + AsyncMock(side_effect=blocked), AsyncMock(side_effect=ValueError("no price")) + ) + + assert caught.value is blocked + + +@pytest.mark.asyncio +async def test_an_allowed_response_leaves_the_charge_to_the_success_path(): + settle = AsyncMock() + + response = await _run_post_call_guardrails_keeping_the_charge(AsyncMock(return_value={"id": "kept"}), settle) + + assert response == {"id": "kept"} + settle.assert_not_awaited() diff --git a/tests/unit/proxy/spend_tracking/test_cost_finalization.py b/tests/unit/proxy/spend_tracking/test_cost_finalization.py index ef70a193280..677d68fd5f6 100644 --- a/tests/unit/proxy/spend_tracking/test_cost_finalization.py +++ b/tests/unit/proxy/spend_tracking/test_cost_finalization.py @@ -1,6 +1,7 @@ import asyncio from datetime import datetime, timezone from typing import Final, Literal +from unittest.mock import patch import pytest from pydantic import TypeAdapter @@ -612,6 +613,22 @@ async def test_guardrail_json_rejection_retains_completed_provider_usage( ProxyLogging._callback_capabilities_cache.clear() +def test_a_response_with_no_cost_result_settles_at_zero_only_when_asked(logging_obj: Logging) -> None: + from litellm.proxy.spend_tracking.cost_finalization import prepare_response_cost + + unpriced: Final = response() + del unpriced._hidden_params["response_cost"] + + with patch.object(logging_obj, "response_cost_calculator", return_value=None): + assert prepare_response_cost(logging_obj, unpriced, NOW) is None + prepared: Final = prepare_response_cost(logging_obj, unpriced, NOW, settle_unknown_cost_at_zero=True) + + assert prepared is not None + assert prepared.amount == 0.0 + assert prepared.payload["response_cost"] == 0.0 + assert (prepared.usage.prompt_tokens, prepared.usage.completion_tokens) == (10, 5) + + @pytest.mark.asyncio @pytest.mark.parametrize("price", [None, 0.0, 0.02], ids=["unpriced", "free", "priced"]) @pytest.mark.parametrize("discount", [0.0, 1.0], ids=["no-discount", "fully-discounted"]) diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 91c87c63067..eeffba68969 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -5327,6 +5327,33 @@ async def test_embedding_header_forwarding_without_model_group_config(): litellm.model_group_settings = original_model_group_settings +@pytest.mark.asyncio +@pytest.mark.parametrize("model", [0, False, [], {}, 7, True]) +async def test_policy_engine_leaves_a_model_that_is_not_a_string_alone( + model: object, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import Policy, PolicyAttachment, PolicyGuardrails + + policy_registry = get_policy_registry() + monkeypatch.setattr( + policy_registry, "_policies", {"baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker"]))} + ) + monkeypatch.setattr(policy_registry, "_initialized", True) + attachment_registry = get_attachment_registry() + monkeypatch.setattr(attachment_registry, "_attachments", [PolicyAttachment(policy="baseline", scope="*")]) + monkeypatch.setattr(attachment_registry, "_initialized", True) + data = {"model": model, "metadata": {}} + + await add_guardrails_from_policy_engine( + data=data, metadata_variable_name="metadata", user_api_key_dict=UserAPIKeyAuth(api_key="test-key") + ) + + assert data["model"] == model + assert data["metadata"]["guardrails"] == ["pii_blocker"] + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine(): """ diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index 5fd40696224..cc84ff9f247 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -275,54 +275,45 @@ async def test_process_guardrail_callback_enriches_and_reraises_http_exception( # --------------------------------------------------------------------------- -def test_process_guardrail_metadata_calls_header_helper(proxy_logging, monkeypatch): - calls: List[Dict[str, Any]] = [] - - def fake_add(request_data, guardrail_name): - calls.append({"data": request_data, "name": guardrail_name}) - - from litellm.proxy.common_utils import callback_utils - - monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) - data = {"metadata": {"guardrails": ["g1", "g2"]}} +def test_process_guardrail_metadata_adds_each_guardrail_once_in_order(proxy_logging): + data = {"metadata": {"guardrails": ["g1", "g2", "g1"]}} proxy_logging._process_guardrail_metadata(data) - snapshot = { - "call_count": len(calls), - "first_name": calls[0]["name"], - "second_name": calls[1]["name"], - "data_passed_is_input": all(c["data"] is data for c in calls), - } - assert snapshot == { - "call_count": 2, - "first_name": "g1", - "second_name": "g2", - "data_passed_is_input": True, - } + assert data["metadata"]["applied_guardrails"] == ["g1", "g2"] -def test_process_guardrail_metadata_skips_already_applied(proxy_logging, monkeypatch): - calls: List[str] = [] - - def fake_add(request_data, guardrail_name): - calls.append(guardrail_name) - - from litellm.proxy.common_utils import callback_utils - - monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) +def test_process_guardrail_metadata_skips_already_applied(proxy_logging): data = {"metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g1"]}} proxy_logging._process_guardrail_metadata(data) - assert calls == ["g2"] + assert data["metadata"]["applied_guardrails"] == ["g1", "g2"] -def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging, monkeypatch): - from litellm.proxy.common_utils import callback_utils +def test_process_guardrail_metadata_records_in_the_litellm_metadata_bucket(proxy_logging): + data = {"litellm_metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g2"]}} + proxy_logging._process_guardrail_metadata(data) + assert data == {"litellm_metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g2", "g1"]}} - monkeypatch.setattr( - callback_utils, - "add_guardrail_to_applied_guardrails_header", - MagicMock(side_effect=AssertionError("should not be called")), - ) - proxy_logging._process_guardrail_metadata({}) + +def test_process_guardrail_metadata_does_not_search_the_applied_list_per_name(proxy_logging): + class _NoMembershipChecks(list): + def __contains__(self, item): + raise AssertionError("searched the applied list") + + names = [f"guard-{i}" for i in range(100)] + data = {"metadata": {"guardrails": names, "applied_guardrails": _NoMembershipChecks()}} + proxy_logging._process_guardrail_metadata(data) + assert data["metadata"]["applied_guardrails"] == names + + +def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging): + data: Dict[str, Any] = {} + proxy_logging._process_guardrail_metadata(data) + assert data == {} + + +def test_process_guardrail_metadata_with_only_settings_entries_leaves_data_unchanged(proxy_logging): + data = {"metadata": {"guardrails": [{"g1": {"extra_body": {}}}]}} + proxy_logging._process_guardrail_metadata(data) + assert data == {"metadata": {"guardrails": [{"g1": {"extra_body": {}}}]}} def test_process_guardrail_metadata_invalid_data_raises(proxy_logging): diff --git a/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py b/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py index 53d8948869f..61e7e47f7b8 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_post_call_success_hook.py @@ -59,6 +59,47 @@ async def test_post_call_success_hook_runs_other_callback_and_replaces_response( assert out == new_response +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("guardrails_only", "guardrails_seen"), [(False, ["key-guard", "model-guard"]), (True, ["key-guard"])] +) +async def test_post_call_success_hook_merges_model_level_guardrails_unless_guardrails_only( + proxy_logging, make_user_api_key_auth, monkeypatch, guardrails_only, guardrails_seen +): + router = MagicMock() + router.registry.return_value.deployments.return_value = [{"litellm_params": {"guardrails": ["model-guard"]}}] + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + g = _make_guardrail(should_run=False) + monkeypatch.setattr(litellm, "callbacks", [g]) + + await proxy_logging.post_call_success_hook( + data={"model": "alias", "metadata": {"guardrails": ["key-guard"]}}, + response=MagicMock(), + user_api_key_dict=make_user_api_key_auth(), + guardrails_only=guardrails_only, + ) + + assert g.should_run_guardrail.call_args.kwargs["data"]["metadata"]["guardrails"] == guardrails_seen + + +@pytest.mark.asyncio +async def test_post_call_success_hook_guardrails_only_leaves_other_callbacks_out( + proxy_logging, make_user_api_key_auth, monkeypatch +): + other = MagicMock(spec=CustomLogger) + other.async_post_call_success_hook = AsyncMock(return_value={"replaced": True}) + g = _make_guardrail(override={"guarded": True}) + monkeypatch.setattr(litellm, "callbacks", [other, g]) + + out = await proxy_logging.post_call_success_hook( + data={}, response={"original": True}, user_api_key_dict=make_user_api_key_auth(), guardrails_only=True + ) + + other.async_post_call_success_hook.assert_not_called() + g.async_post_call_success_hook.assert_awaited_once() + assert out == {"guarded": True} + + @pytest.mark.asyncio async def test_post_call_success_hook_guardrail_should_not_run_skipped( proxy_logging, make_user_api_key_auth, monkeypatch diff --git a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py index 26ee9f04d01..77542f1c3f4 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import Any, Dict +from typing import Any, Dict, Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -16,6 +16,7 @@ from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline, PipelineStep from litellm.types.utils import CallTypesLiteral @@ -1014,3 +1015,39 @@ def test_request_has_pre_call_guardrails_reads_pipelines_from_both_metadata_buck monkeypatch.setattr(litellm, "callbacks", []) assert proxy_logging.request_has_pre_call_guardrails(request_data) is expected + + +@pytest.mark.parametrize( + ("pipelines", "scanned"), + [ + ([("policy", GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g")]))], False), + ([("policy", GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g")]))], True), + ( + [ + ("after", GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="g")])), + ("before", GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="g")])), + ], + True, + ), + (["p1"], True), + (5, True), + ([], False), + ], + ids=["post-call", "pre-call", "both", "unresolved-entry", "non-list", "empty"], +) +def test_only_a_pipeline_running_before_the_provider_scans_the_request(proxy_logging, monkeypatch, pipelines, scanned): + monkeypatch.setattr(litellm, "callbacks", []) + request_data: Final = {"metadata": {"_guardrail_pipelines": pipelines}} + + assert proxy_logging.request_is_scanned_before_the_provider(request_data) is scanned + assert proxy_logging.request_has_pre_call_guardrails(request_data) is bool(pipelines) + + +def test_a_pre_call_guardrail_scans_the_request_beside_a_post_call_pipeline(proxy_logging, monkeypatch): + guardrail: Final = CustomGuardrail(guardrail_name="g", default_on=True, event_hook=GuardrailEventHooks.pre_call) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + pipeline: Final = GuardrailPipeline(mode="post_call", steps=[PipelineStep(guardrail="other")]) + + assert proxy_logging.request_is_scanned_before_the_provider( + {"metadata": {"_guardrail_pipelines": [("p", pipeline)]}} + )