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/llms/pass_through/guardrail_translation/handler.py b/litellm/llms/pass_through/guardrail_translation/handler.py index 712e51f7afc..0af14c23654 100644 --- a/litellm/llms/pass_through/guardrail_translation/handler.py +++ b/litellm/llms/pass_through/guardrail_translation/handler.py @@ -243,7 +243,8 @@ 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 + 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): @@ -263,7 +264,7 @@ class PassThroughEndpointHandler(BaseTranslation): # Extract text to check text_to_check: Final = ( self._extract_text_for_guardrail(response, field_expressions) - if field_expressions or read_whole_response + if field_expressions or read_whole_response or litellm_logging_obj is None else _response_text(response_fields or {}) ) 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 f3309b98462..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 @@ -378,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", @@ -392,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, *, @@ -561,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 @@ -597,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_v2.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py index 5ca64b3c0bd..90afa69dcf7 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai_v2.py @@ -1,9 +1,9 @@ import copy import json import os -from collections.abc import Mapping, Sequence +from collections.abc import Iterator, Mapping, Sequence from datetime import datetime -from itertools import chain, product +from itertools import chain from string import Formatter from types import MappingProxyType from typing import Final, Literal @@ -261,14 +261,20 @@ def _gemini_text(content: object) -> str: return "\n".join(text for text in texts if isinstance(text, str) and text) +def _gemini_tool_part_values(part: object) -> Iterator[object]: + fields: Final = as_str_mapping(part) or {} + return ( + (as_str_mapping(fields[name]) or {}).get(_GEMINI_TOOL_PART_FIELDS[name]) + for name in fields + if name in _GEMINI_TOOL_PART_FIELDS + ) + + def _gemini_tool_text(content: object) -> str: """What one Gemini ``Content`` carries in its function call, function response and code execution parts, with structured values as JSON.""" parts: Final = _gemini_items((as_str_mapping(content) or {}).get("parts")) - values: Final = ( - (as_str_mapping((as_str_mapping(part) or {}).get(part_name)) or {}).get(field) - for part, (part_name, field) in product(parts, _GEMINI_TOOL_PART_FIELDS.items()) - ) + values: Final = chain.from_iterable(map(_gemini_tool_part_values, parts)) return "\n".join( value if isinstance(value, str) else json.dumps(value, default=str, ensure_ascii=False) for value in values diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 1484ec744d2..2bbd74314d9 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -210,6 +210,11 @@ _METADATA_KEYS: Final = frozenset(("litellm_metadata", "metadata")) _DECODED_JSON: Final = TypeAdapter(object) _GUARDRAIL_ENTRIES: Final = TypeAdapter(list[object]) _RESPONSE_HEADERS: Final = TypeAdapter(Mapping[str, str]) +_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", "tags") _KEPT_OUT_OF_LITELLM_PARAMS: Final = _METADATA_KEYS | frozenset(wire_names(ProxyRequestState)) _PRICING_MODEL_OVERRIDE_PARAMS: Final = frozenset(("base_model",)) @@ -1069,12 +1074,11 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): 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`` that isn't an object, which the - merge can't combine with the metadata's own. The provider receives neither field.""" - if "guardrails" in parsed_body and not isinstance(parsed_body["guardrails"], list | dict): - parsed_body.pop("guardrails") - if "guardrail_config" in parsed_body and not isinstance(parsed_body["guardrail_config"], dict): - parsed_body.pop("guardrail_config") + """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( @@ -1093,6 +1097,7 @@ async def _inherit_auth_guardrails( 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") @@ -1102,6 +1107,9 @@ async def _inherit_auth_guardrails( 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 @@ -1135,6 +1143,12 @@ def _admin_guardrail_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, ) +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, and the tags a tag-scoped guardrail mode reads, from the logging metadata that request setup moved the body's ``metadata`` into""" @@ -1142,7 +1156,10 @@ def _post_call_selection_metadata(kwargs: Mapping[str, object], inherited: bool) return MappingProxyType({}) litellm_params: Final = as_str_mapping(kwargs.get("litellm_params")) or {} metadata: Final = as_str_mapping(litellm_params.get("metadata")) or {} - return {field: metadata[field] for field in _POST_CALL_SELECTION_FIELDS if field in metadata} + 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: @@ -1160,10 +1177,11 @@ def _post_call_hook_data( """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""" + 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({}) - return { + data: Final = { **body, "metadata": { **(as_str_mapping(body.get("metadata")) or {}), @@ -1171,8 +1189,10 @@ def _post_call_hook_data( "guardrails": guardrails or (), **admin_metadata, }, - "litellm_logging_obj": logging_obj, } + 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( diff --git a/tests/integration/authorization/test_passthrough_auth_guardrails.py b/tests/integration/authorization/test_passthrough_auth_guardrails.py index 7786c3e5f4e..86792e1ff28 100644 --- a/tests/integration/authorization/test_passthrough_auth_guardrails.py +++ b/tests/integration/authorization/test_passthrough_auth_guardrails.py @@ -47,6 +47,8 @@ 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" @@ -175,6 +177,7 @@ class AuthGuardrailRig: 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] @@ -187,6 +190,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I 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)) @@ -224,6 +228,12 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "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": { @@ -317,6 +327,16 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I "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": { @@ -367,6 +387,7 @@ def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> I post_call_vendor, key_post_call_vendor, generic_post_call_vendor, + global_unified_vendor, model_post_call_vendor, upstreams, ) @@ -866,6 +887,78 @@ def test_custom_endpoint_writes_no_charge_for_a_response_its_post_call_guardrail 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: 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/llms/pass_through/guardrail_translation/test_handler.py b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py index e028ff2a7bd..1ba79b704da 100644 --- a/tests/unit/llms/pass_through/guardrail_translation/test_handler.py +++ b/tests/unit/llms/pass_through/guardrail_translation/test_handler.py @@ -231,10 +231,22 @@ def test_passthrough_route_block_before_any_invoke_frame_is_a_whole_message_in_c 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, request_data={}) + 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 = { 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 daa9945e406..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 @@ -2703,3 +2703,52 @@ def test_extract_user_api_key_metadata_reads_past_a_litellm_metadata_that_is_not ) 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/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index fa83297bb05..e2c73f0ac5d 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 @@ -59,6 +59,9 @@ from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( SafeRouteAdder, _anthropic_user_id_to_forward, _inherit_auth_guardrails, + _post_call_hook_data, + _post_call_selection_metadata, + _starts_post_call_pass, WebsocketFrameModelAllowed, WebsocketFrameModelGate, WebsocketFrameModelRefused, @@ -13416,11 +13419,83 @@ class TestInheritAuthGuardrails: await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test", metadata=key_metadata), None, True) assert "guardrails" not in body - assert (body["metadata"].get("guardrails") or []) == [ + 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):