From 9ff6c63d5c27d20aa6b55e2f8724f1994a5793bd Mon Sep 17 00:00:00 2001 From: Alexander Grattan <51346343+agrattan0820@users.noreply.github.com> Date: Tue, 14 Apr 2026 23:14:10 -0400 Subject: [PATCH] fix: ensure `litellm_metadata` is attached to `pre_call` guardrail to align with `post_call` guardrail (#25641) * fix: ensure `litellm_metadata` is attached to pre_call to align with post_call * refactor: remove unused BaseTranslation._ensure_litellm_metadata * refactor: module level imports for ensure_litellm_metadata and CodeQL * fix: update based off of Codex comment * revert: undo usage of `_guardrail_litellm_metadata` --- .../unified_guardrail/unified_guardrail.py | 18 ++++++ .../guardrail_hooks/test_grayswan.py | 62 ++++++++++++++++++- 2 files changed, 78 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index a1623121da5..367e6b2f150 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -52,6 +52,20 @@ def _get_a2a_request_id( endpoint_guardrail_translation_mappings = None +def _ensure_litellm_metadata(data: dict, user_api_key_dict: UserAPIKeyAuth) -> None: + """Populate data['litellm_metadata'] from user_api_key_dict if absent.""" + if "litellm_metadata" not in data: + from litellm.llms.base_llm.guardrail_translation.base_translation import ( + BaseTranslation, + ) + + user_metadata = BaseTranslation.transform_user_api_key_dict_to_metadata( + user_api_key_dict + ) + if user_metadata: + data["litellm_metadata"] = user_metadata + + class UnifiedLLMGuardrails(CustomLogger): def __init__( self, @@ -120,6 +134,8 @@ class UnifiedLLMGuardrails(CustomLogger): CallTypes(call_type) ]() + _ensure_litellm_metadata(data, user_api_key_dict) + data = await endpoint_translation.process_input_messages( data=data, guardrail_to_apply=guardrail_to_apply, @@ -177,6 +193,8 @@ class UnifiedLLMGuardrails(CustomLogger): CallTypes(call_type) ]() + _ensure_litellm_metadata(data, user_api_key_dict) + return await endpoint_translation.process_input_messages( data=data, guardrail_to_apply=guardrail_to_apply, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py index 109ad0bfdc8..f2e7447239f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_grayswan.py @@ -4,6 +4,7 @@ import pytest from fastapi import HTTPException from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.grayswan.grayswan import ( GraySwanGuardrail, GraySwanGuardrailAPIError, @@ -435,7 +436,10 @@ def test_format_violation_message() -> None: assert "Gray Swan Cygnal Guardrail" in message assert "the input query has a violation score of 0.85" in message assert "violating the rule(s): 1, 3, 5" in message - assert "Mutation effort to make the harmful intention disguised was DETECTED" in message + assert ( + "Mutation effort to make the harmful intention disguised was DETECTED" + in message + ) # IPI should not be in message since it's False assert "Indirect Prompt Injection was DETECTED" not in message @@ -446,4 +450,58 @@ def test_format_violation_message() -> None: assert "Gray Swan Cygnal Guardrail" in message assert "the model response has a violation score of 0.85" in message assert "violating the rule(s): 1, 3, 5" in message - assert "Mutation effort to make the harmful intention disguised was DETECTED" in message + assert ( + "Mutation effort to make the harmful intention disguised was DETECTED" + in message + ) + + +def test_prepare_payload_includes_litellm_metadata( + grayswan_guardrail: GraySwanGuardrail, +) -> None: + """Verify _prepare_payload forwards litellm_metadata from request_data.""" + messages = [{"role": "user", "content": "hello"}] + request_data = { + "litellm_metadata": { + "user_api_key_user_id": "user-123", + "user_api_key_team_id": "team-456", + "user_api_key_spend": 0, + } + } + + payload = grayswan_guardrail._prepare_payload(messages, {}, request_data) + + assert payload is not None + assert "litellm_metadata" in payload + assert payload["litellm_metadata"]["user_api_key_user_id"] == "user-123" + assert payload["litellm_metadata"]["user_api_key_team_id"] == "team-456" + + +def test_ensure_litellm_metadata_populates_from_user_api_key_dict() -> None: + """Verify _ensure_litellm_metadata populates litellm_metadata.""" + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + _ensure_litellm_metadata, + ) + + user_auth = UserAPIKeyAuth(user_id="u1", team_id="t1", api_key="sk-test-hashed") + data: dict = {} + + _ensure_litellm_metadata(data, user_auth) + + assert "litellm_metadata" in data + assert data["litellm_metadata"]["user_api_key_user_id"] == "u1" + assert data["litellm_metadata"]["user_api_key_team_id"] == "t1" + + +def test_ensure_litellm_metadata_noop_when_already_present() -> None: + """Verify _ensure_litellm_metadata does not overwrite existing litellm_metadata.""" + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + _ensure_litellm_metadata, + ) + + user_auth = UserAPIKeyAuth(user_id="should-not-appear") + data: dict = {"litellm_metadata": {"existing": "value"}} + + _ensure_litellm_metadata(data, user_auth) + + assert data["litellm_metadata"] == {"existing": "value"}