diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py index 95525bf9719..4e018310d9b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -22,6 +22,7 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.base import ( GuardrailConfigModel, @@ -117,21 +118,33 @@ class SingulrGuardrail(CustomGuardrail): return value return None + @staticmethod + def _resolve_user_role_from_request_data(request_data: Mapping[str, Any]) -> str | None: + for container_key in ("litellm_metadata", "metadata"): + container: Final = request_data.get(container_key) or _EMPTY_MAPPING + auth: Final = container.get("user_api_key_auth") if container else None + if isinstance(auth, UserAPIKeyAuth) and auth.user_role: + return auth.user_role.value + return None + @classmethod def _build_metadata(cls, request_data: Mapping[str, Any]) -> Mapping[str, Any] | None: fields: Final = ( "user_api_key_alias", "user_api_key_user_id", "user_api_key_user_email", + "user_api_key_org_id", "user_api_key_org_alias", + "user_api_key_team_id", "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()): + resolved: Final = ( + *((field, cls._resolve_metadata_value(request_data=request_data, key=field)) for field in fields), + ("user_api_key_user_role", cls._resolve_user_role_from_request_data(request_data=request_data)), + ) + if not any(value for _, value in resolved): return None - return {key: value for key, value in resolved.items() if value} # mutable-ok: short-lived JSON payload dict + return {key: value for key, value in resolved 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 da3773a997a..9b789984684 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py @@ -4,6 +4,7 @@ import httpx import pytest from litellm.exceptions import GuardrailRaisedException +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.guardrails.guardrail_hooks.singulr.singulr import SingulrGuardrail from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.singulr import ( @@ -329,16 +330,75 @@ class TestSingulrRequestPayload: 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_user_api_key_org_id_is_forwarded_in_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"litellm_metadata": {"user_api_key_org_id": "org-123"}} + 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_id": "org-123"} + + @pytest.mark.asyncio + async def test_user_api_key_team_id_is_forwarded_in_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"litellm_metadata": {"user_api_key_team_id": "team-456"}} + 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_id": "team-456"} + + @pytest.mark.asyncio + async def test_user_api_key_user_role_is_forwarded_in_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) + request_data = {"litellm_metadata": {"user_api_key_auth": auth}} + 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_role": LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value + } + + @pytest.mark.asyncio + async def test_no_user_role_available_omits_role_from_metadata(self, singulr_guardrail): + resp = _make_response({"should_block": False}) + request_data = {"litellm_metadata": {"user_api_key_alias": "my-key-alias"}} + 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 "user_api_key_user_role" not in sent_payload["metadata"] + @pytest.mark.asyncio async def test_all_user_metadata_fields_forwarded_together(self, singulr_guardrail): resp = _make_response({"should_block": False}) + auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY) 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_id": "org-123", "user_api_key_org_alias": "Acme Org", + "user_api_key_team_id": "team-456", "user_api_key_team_alias": "AI Content Security Team", + "user_api_key_auth": auth, } } with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post: @@ -352,8 +412,11 @@ class TestSingulrRequestPayload: "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_id": "org-123", "user_api_key_org_alias": "Acme Org", + "user_api_key_team_id": "team-456", "user_api_key_team_alias": "AI Content Security Team", + "user_api_key_user_role": LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value, } @pytest.mark.asyncio