diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 501c8352b10..c7421408a04 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -3986,6 +3986,11 @@ def add_guardrails_from_key_or_team_metadata( _add_guardrails_from_key_or_team_metadata: Final = add_guardrails_from_key_or_team_metadata +def _string_policy_names(policies: Sequence[object]) -> tuple[str, ...]: + """The policy names in a key, team or project ``policies`` list, less an entry that isn't a name""" + return tuple(name for name in policies if isinstance(name, str)) + + def _add_guardrails_from_policies_in_metadata( key_metadata: dict | None, team_metadata: dict | None, @@ -4021,19 +4026,25 @@ def _add_guardrails_from_policies_in_metadata( if key_metadata and "policies" in key_metadata: if isinstance(key_metadata["policies"], list) and len(key_metadata["policies"]) > 0: premium_user_check() - policy_names.update(key_metadata["policies"]) + policy_names.update( + _string_policy_names(cast("Sequence[object]", key_metadata["policies"])) # cast-ok: checked list + ) # Add team-level policies if team_metadata and "policies" in team_metadata: if isinstance(team_metadata["policies"], list) and len(team_metadata["policies"]) > 0: premium_user_check() - policy_names.update(team_metadata["policies"]) + policy_names.update( + _string_policy_names(cast("Sequence[object]", team_metadata["policies"])) # cast-ok: checked list + ) # Add project-level policies if project_metadata and "policies" in project_metadata: if isinstance(project_metadata["policies"], list) and len(project_metadata["policies"]) > 0: premium_user_check() - policy_names.update(project_metadata["policies"]) + policy_names.update( + _string_policy_names(cast("Sequence[object]", project_metadata["policies"])) # cast-ok: checked list + ) if not policy_names: return diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index a2b77b3168e..69105bcbb71 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -209,6 +209,13 @@ _FIELDS_GUARDRAILS_READ_BY_SHAPE: Final = ( ("guardrail_config", (dict,)), ("proxy_server_request", (dict,)), ) +_PROXY_FILLED_FIELD_SHAPES: Final = ( + ("litellm_call_id", (str,)), + ("litellm_trace_id", (str,)), + ("custom_llm_provider", (str,)), + ("api_base", (str,)), + ("api_key", (str,)), +) _POST_CALL_SELECTION_FIELDS: Final = ( "_guardrail_pipelines", "_pipeline_managed_guardrails", @@ -223,6 +230,19 @@ def _is_mergeable_client_metadata(key: str, value: object) -> bool: return not key.startswith("user_api_key_") and (key != "model_info" or isinstance(value, Mapping)) +def _holds_string_headers(value: object) -> bool: + headers: Final = as_str_mapping(value) + return headers is not None and all(isinstance(header, str) for header in headers.values()) + + +def _is_guardrail_readable_client_metadata(key: str, value: object) -> bool: + if key == "headers": + return _holds_string_headers(value) + if key == "requester_metadata": + return isinstance(value, Mapping) + return _is_mergeable_client_metadata(key, value) + + class _RegisteredPassThroughRoute(TypedDict): endpoint_id: ReadOnly[str] path: ReadOnly[str | None] @@ -925,24 +945,43 @@ async def _inherit_auth_guardrails( 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") + client_metadata: Final = as_str_mapping(parsed_body["metadata"]) or {} + client_body_guardrails: Final = {"guardrails": parsed_body["guardrails"]} if "guardrails" in parsed_body else {} + client_metadata_guardrails: Final = ( + {"guardrails": client_metadata["guardrails"]} if "guardrails" in client_metadata else {} + ) + client_named: Final = client_metadata.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"] + merged: Final = (as_str_mapping(metadata) or {}).get("guardrails") + entries: Final = ( + merged_guardrail_entries( + _GUARDRAIL_ENTRIES.validate_python(merged), _client_named_guardrail_entries(client_named) + ) + if isinstance(merged, list) + else () + ) + runs_guardrails: Final = bool(entries) or bool((as_str_mapping(metadata) or {}).get("_guardrail_pipelines")) + if entries and isinstance(metadata, dict): + metadata["guardrails"] = entries + elif isinstance(merged, list) and isinstance(metadata, dict): + _restore_unmerged_guardrail_fields( + parsed_body, + cast("dict[str, object]", metadata), # cast-ok: the request body's JSON metadata object + client_body_guardrails, + client_metadata_guardrails, + ) + if runs_guardrails: + _drop_client_values_guardrails_cannot_read(parsed_body) 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): + if not entries: 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, @@ -951,11 +990,54 @@ async def _inherit_auth_guardrails( return entries if entries and _named_guardrail_runs_post_call(admin_scoped_body, entries) else None +def _drop_client_values_guardrails_cannot_read( + parsed_body: dict[str, object], # mutable-ok: strips the request body in place before its guardrails run +) -> None: + """Drop what a client sent in a field the proxy fills on the native routes when the guardrails can't read it there: + a ``litellm_call_id``, ``litellm_trace_id``, ``custom_llm_provider``, ``api_base`` or ``api_key`` that isn't a + string, a ``proxy_server_request`` or ``metadata.headers`` whose headers aren't strings, a ``requester_metadata`` + that isn't an object, and any ``user_api_key_*`` metadata field, which only the proxy sets. The provider receives + none of these fields.""" + for field, shape in _PROXY_FILLED_FIELD_SHAPES: + if field in parsed_body and not isinstance(parsed_body[field], shape): + parsed_body.pop(field) + server_request: Final = as_str_mapping(parsed_body.get("proxy_server_request")) or {} + if "headers" in server_request and not _holds_string_headers(server_request["headers"]): + parsed_body.pop("proxy_server_request") + for bucket in (parsed_body.get(key) for key in sorted(_METADATA_KEYS)): + if isinstance(bucket, dict): + _drop_metadata_guardrails_cannot_read(cast("dict[str, object]", bucket)) # cast-ok: a JSON object + + +def _drop_metadata_guardrails_cannot_read( + metadata: dict[str, object], # mutable-ok: strips the request body's metadata in place +) -> None: + unreadable: Final = tuple( + field for field, value in metadata.items() if not _is_guardrail_readable_client_metadata(field, value) + ) + for field in unreadable: + metadata.pop(field) + + +def _restore_unmerged_guardrail_fields( + parsed_body: dict[str, object], # mutable-ok: puts the client's own fields back on the live request body + metadata: dict[str, object], # mutable-ok: the live request body's metadata + client_body_guardrails: Mapping[str, object], + client_metadata_guardrails: Mapping[str, object], +) -> None: + """Put back the client's ``guardrails`` and ``metadata.guardrails`` when the merge found no guardrail to run, so + the body keeps main's shape""" + metadata.pop("guardrails", None) + metadata.update(client_metadata_guardrails) + parsed_body.update(client_body_guardrails) + + 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""" + """The guardrails ``metadata.guardrails`` names: its entries, or the keys of the dict form the guardrails read that + it doesn't turn off with ``false``""" if isinstance(client_named, list): return _GUARDRAIL_ENTRIES.validate_python(client_named) - return tuple(as_str_mapping(client_named) or ()) + return tuple(name for name, run in (as_str_mapping(client_named) or {}).items() if run is not False) def _admin_guardrail_metadata(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, object]: diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index df5e891af5c..1eb59e41a54 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -892,9 +892,8 @@ 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 - + """Adds each name the way add_guardrail_to_applied_guardrails_header does, in time linear in the names. A record a + guardrail left that isn't a list stays as it is, so the request doesn't fail on it.""" if not names: return metadata: Final = cast("dict[str, object]", get_or_create_metadata_bucket(data)[1]) # cast-ok: str-keyed bucket @@ -903,8 +902,6 @@ def _add_applied_guardrails_in_place( 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)) 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 2d8ba198cb0..0f3428ccffd 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 @@ -13372,6 +13372,99 @@ class TestInheritAuthGuardrails: assert body["metadata"] == {} + @pytest.mark.asyncio + @pytest.mark.parametrize( + "extra", + [ + {"guardrails": []}, + {"guardrails": [], "metadata": {"user": "u"}}, + {"guardrails": [], "metadata": {"guardrails": []}}, + {"metadata": {"guardrails": []}}, + ], + ids=["guardrails", "with-metadata", "with-metadata-guardrails", "metadata-guardrails-only"], + ) + async def test_an_empty_guardrails_list_on_a_plain_key_keeps_the_body_as_sent(self, extra: dict[str, object]): + body: Final[dict[str, object]] = self._chat_body(**extra) + + assert await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True) is None + + assert body == self._chat_body(**extra) + + @pytest.mark.asyncio + async def test_a_guardrail_the_body_turns_off_with_false_stays_out_of_the_merge(self): + body: Final[dict[str, object]] = self._chat_body( + metadata={"guardrails": {"lak-legacy": False, "body-guard": {}}} + ) + + 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 == {**self._chat_body(), "metadata": {"guardrails": ["key-guard", "body-guard"]}} + + @staticmethod + def _client_values_the_guardrails_cant_read() -> dict[str, object]: + return { + "litellm_call_id": 5, + "litellm_trace_id": 5, + "custom_llm_provider": 5, + "api_base": 5, + "api_key": 5, + "proxy_server_request": {"headers": {"user-agent": 7}}, + "metadata": { + "headers": "abc", + "requester_metadata": "abc", + "user_api_key_alias": "forged", + "user_api_key_user_id": 5, + "user": "kept", + }, + "litellm_metadata": {"headers": {"x-aim-user-email": 5}, "user_api_key_alias": 5, "tags": ["kept"]}, + } + + @pytest.mark.asyncio + async def test_client_values_the_guardrails_cant_read_are_dropped_when_guardrails_run(self): + body: Final[dict[str, object]] = self._chat_body(**self._client_values_the_guardrails_cant_read()) + + 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 == { + **self._chat_body(), + "metadata": {"user": "kept", "guardrails": ["key-guard"]}, + "litellm_metadata": {"tags": ["kept"]}, + } + + @pytest.mark.asyncio + async def test_client_values_the_guardrails_can_read_are_kept_when_guardrails_run(self): + readable: Final[dict[str, object]] = { + "litellm_call_id": "call", + "litellm_trace_id": "trace", + "custom_llm_provider": "openai", + "api_base": "https://example.com", + "api_key": "sk-client", + "proxy_server_request": {"headers": {"user-agent": "ua"}}, + } + client_metadata: Final[dict[str, object]] = {"headers": {"x-aim-user-email": "a@b.c"}, "requester_metadata": {}} + body: Final[dict[str, object]] = self._chat_body(**readable, metadata=dict(client_metadata)) + + 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 == {**self._chat_body(**readable), "metadata": {**client_metadata, "guardrails": ["key-guard"]}} + + @pytest.mark.asyncio + async def test_client_values_stay_as_sent_when_no_guardrail_runs(self): + body: Final[dict[str, object]] = self._chat_body(**self._client_values_the_guardrails_cant_read()) + + assert await _inherit_auth_guardrails(body, UserAPIKeyAuth(api_key="sk-test"), None, True) is None + + assert body == self._chat_body(**self._client_values_the_guardrails_cant_read()) + @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): diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index eeffba68969..c8e0617566f 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -5622,6 +5622,38 @@ def test_add_guardrails_from_auth_metadata_lists_each_key_and_team_guardrail_onc assert data["metadata"]["guardrails"] == [configured, "audit", "toxic"] +@pytest.mark.parametrize("level", ["key", "team", "project"]) +def test_add_guardrails_from_auth_metadata_skips_a_policy_name_that_is_not_a_string(level: str): + from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails + + metadata: Final[dict[str, object]] = {} + data: Final[dict[str, object]] = {"model": "gpt-4", "metadata": metadata} + policy_registry = get_policy_registry() + policy_registry._policies = {"baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker"]))} + policy_registry._initialized = True + policies = {"policies": [{"a": 1}, ["b"], "baseline"]} + + try: + with patch("litellm.proxy.utils.premium_user_check"): + add_guardrails_from_auth_metadata( + user_api_key_dict=UserAPIKeyAuth( + api_key="test-key", + metadata=policies if level == "key" else {}, + team_metadata=policies if level == "team" else {}, + project_metadata=policies if level == "project" else None, + ), + data=data, + metadata_variable_name="metadata", + ) + finally: + policy_registry._policies = {} + policy_registry._initialized = False + + assert metadata["guardrails"] == ["pii_blocker"] + + def test_add_guardrails_from_auth_metadata_keeps_a_selected_guardrail_s_settings_when_a_key_policy_adds_it(): from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata from litellm.proxy.policy_engine.policy_registry import get_policy_registry 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 77542f1c3f4..43f4d0fa8ea 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 @@ -949,6 +949,44 @@ async def test_pre_call_block_keeps_request_declared_guardrail_in_applied_guardr assert data["metadata"]["applied_guardrails"] == ["blocker", "declared-post-call"] +class _RecordsAppliedGuardrails(CustomGuardrail): + def __init__(self, record: object) -> None: + super().__init__(guardrail_name="recorder", event_hook=GuardrailEventHooks.pre_call, default_on=False) + self._record: Final = record + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, dict[str, object]], + call_type: CallTypesLiteral, + ) -> dict[str, dict[str, object]]: + data["metadata"]["applied_guardrails"] = self._record + return data + + +@pytest.mark.asyncio +@pytest.mark.parametrize("record", [5, "other", ("other",)], ids=["int", "str", "tuple"]) +async def test_a_non_list_applied_guardrails_record_does_not_fail_the_request( + proxy_logging, make_user_api_key_auth, monkeypatch, record: object +): + monkeypatch.setattr(litellm, "callbacks", [_RecordsAppliedGuardrails(record)]) + proxy_logging.slack_alerting_instance = MagicMock(alerting=None) + data: Final = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"guardrails": ["recorder", "declared-post-call"]}, + } + + out: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data=data, + call_type="completion", + ) + + assert out["metadata"]["applied_guardrails"] == record + + @pytest.mark.asyncio async def test_skip_guardrails_still_runs_non_guardrail_callbacks(proxy_logging, make_user_api_key_auth, monkeypatch): accountant = _Accountant()