diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py index efc0b7c3aa4..95525bf9719 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -108,52 +108,30 @@ class SingulrGuardrail(CustomGuardrail): return SingulrGuardrailConfigModel @staticmethod - def _resolve_key_alias_from_request_data(request_data: Mapping[str, Any]) -> str | None: - litellm_metadata: Final = request_data.get("litellm_metadata") or _EMPTY_MAPPING - if litellm_metadata: - litellm_metadata_alias: Final = litellm_metadata.get("user_api_key_alias") - if litellm_metadata_alias: - return litellm_metadata_alias - - metadata: Final = request_data.get("metadata") or _EMPTY_MAPPING - if metadata: - metadata_alias: Final = metadata.get("user_api_key_alias") - if metadata_alias: - return metadata_alias - - return None - - - @staticmethod - def _resolve_user_id_from_request_data(request_data: Mapping[str, Any]) -> str | None: - litellm_metadata: Final = request_data.get("litellm_metadata") or _EMPTY_MAPPING - if litellm_metadata: - litellm_metadata_alias: Final = litellm_metadata.get("user_api_key_user_id") - if litellm_metadata_alias: - return litellm_metadata_alias - - metadata: Final = request_data.get("metadata") or _EMPTY_MAPPING - if metadata: - metadata_user_id: Final = metadata.get("user_api_key_user_id") - if metadata_user_id: - return metadata_user_id - + def _resolve_metadata_value(request_data: Mapping[str, Any], key: str) -> str | None: + for container_key in ("litellm_metadata", "metadata"): + container: Final = request_data.get(container_key) or _EMPTY_MAPPING + if container: + value: Final = container.get(key) + if value: + return value return None @classmethod def _build_metadata(cls, request_data: Mapping[str, Any]) -> Mapping[str, Any] | None: - user_api_key_alias: Final = cls._resolve_key_alias_from_request_data(request_data=request_data) - user_api_key_user_id: Final = cls._resolve_user_id_from_request_data(request_data=request_data) - if not user_api_key_alias and not user_api_key_user_id: - return None - return { # mutable-ok: short-lived JSON payload dict - key: value - for key, value in ( - ("user_api_key_alias", user_api_key_alias), - ("user_api_key_user_id", user_api_key_user_id), - ) - if value + fields: Final = ( + "user_api_key_alias", + "user_api_key_user_id", + "user_api_key_user_email", + "user_api_key_org_alias", + "user_api_key_team_alias", + ) + resolved: Final = { # mutable-ok: short-lived JSON payload dict + field: cls._resolve_metadata_value(request_data=request_data, key=field) for field in fields } + if not any(resolved.values()): + return None + return {key: value for key, value in resolved.items() if value} # mutable-ok: short-lived JSON payload dict @staticmethod def _build_user_message(text: str) -> Mapping[str, Any]: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py index b607c9638e8..da3773a997a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py @@ -291,12 +291,54 @@ class TestSingulrRequestPayload: assert sent_payload["metadata"] == {"user_api_key_user_id": "fallback-user-id"} @pytest.mark.asyncio - async def test_user_api_key_alias_and_user_id_both_forwarded(self, singulr_guardrail): + async def test_user_api_key_user_email_is_forwarded_in_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"litellm_metadata": {"user_api_key_user_email": "user@example.com"}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == {"user_api_key_user_email": "user@example.com"} + + @pytest.mark.asyncio + async def test_user_api_key_organization_alias_is_forwarded_in_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"litellm_metadata": {"user_api_key_org_alias": "Acme Org"}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == {"user_api_key_org_alias": "Acme Org"} + + @pytest.mark.asyncio + async def test_user_api_key_team_alias_is_forwarded_in_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"litellm_metadata": {"user_api_key_team_alias": "AI Content Security Team"}} + with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: + await singulr_guardrail.apply_guardrail( + inputs={"texts": ["hi"]}, + request_data=request_data, + input_type="request", + ) + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["metadata"] == {"user_api_key_team_alias": "AI Content Security Team"} + + @pytest.mark.asyncio + async def test_all_user_metadata_fields_forwarded_together(self, singulr_guardrail): resp = _make_response({"should_block": False}) request_data = { "litellm_metadata": { "user_api_key_alias": "my-key-alias", "user_api_key_user_id": "my-user-id", + "user_api_key_user_email": "user@example.com", + "user_api_key_org_alias": "Acme Org", + "user_api_key_team_alias": "AI Content Security Team", } } with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: @@ -309,6 +351,9 @@ class TestSingulrRequestPayload: assert sent_payload["metadata"] == { "user_api_key_alias": "my-key-alias", "user_api_key_user_id": "my-user-id", + "user_api_key_user_email": "user@example.com", + "user_api_key_org_alias": "Acme Org", + "user_api_key_team_alias": "AI Content Security Team", } @pytest.mark.asyncio