scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 87

feat(guardrails): forward org id, team id, and user role to Singulr

Browse source 
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.
This commit is contained in:
aniket-kardile 2026-08-26 13:03:01 +05:30
parent 555abfdd10
commit d3367c614e
No known key found for this signature in database
2 changed files with 81 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py
View file

@ -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]:

63
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py
View file

@ -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