From d3367c614e3a2a9397e84ee85cfe3e7d53d541d3 Mon Sep 17 00:00:00 2001 From: aniket-kardile Date: Wed, 26 Aug 2026 13:03:01 +0530 Subject: [PATCH] feat(guardrails): forward org id, team id, and user role to Singulr Adds user_api_key_org_id, user_api_key_team_id, and user_api_key_user_role to the Singulr guardrail metadata payload. Role is read off the UserAPIKeyAuth object the proxy attaches to request metadata, since it isn't flattened into the metadata dict the way the other fields are. --- .../guardrail_hooks/singulr/singulr.py | 23 +++++-- .../guardrail_hooks/test_singulr.py | 63 +++++++++++++++++++ 2 files changed, 81 insertions(+), 5 deletions(-) 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