scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

feat(guardrails): forward user_api_key_user_id to Singulr metadata

Browse source 
Wires the existing (unused) user-id resolver into the metadata payload
alongside the key alias, and fixes its metadata fallback which read
user_api_key_alias instead of user_api_key_user_id.
This commit is contained in:
aniket-kardile 2026-08-26 12:38:36 +05:30
parent 9cfe6e7c7d
commit 2b42b4a643
No known key found for this signature in database
2 changed files with 74 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -123,12 +123,37 @@ class SingulrGuardrail(CustomGuardrail):
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
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)
if not user_api_key_alias:
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 {"user_api_key_alias": user_api_key_alias} # mutable-ok: short-lived JSON payload dict
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
}
@staticmethod
def _build_user_message(text: str) -> Mapping[str, Any]:

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

@ -264,6 +264,53 @@ class TestSingulrRequestPayload:
sent_payload = mock_post.call_args.kwargs["json"]
assert sent_payload["metadata"] == {"user_api_key_alias": "fallback-alias"}
@pytest.mark.asyncio
async def test_user_api_key_user_id_is_forwarded_in_metadata(self, singulr_guardrail):
resp = _make_response({"should_block": False})
request_data = {"litellm_metadata": {"user_api_key_user_id": "my-user-id"}}
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_id": "my-user-id"}
@pytest.mark.asyncio
async def test_falls_back_to_regular_metadata_for_user_id(self, singulr_guardrail):
resp = _make_response({"should_block": False})
request_data = {"metadata": {"user_api_key_user_id": "fallback-user-id"}}
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_id": "fallback-user-id"}
@pytest.mark.asyncio
async def test_user_api_key_alias_and_user_id_both_forwarded(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",
}
}
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_alias": "my-key-alias",
"user_api_key_user_id": "my-user-id",
}
@pytest.mark.asyncio
async def test_no_key_alias_available_sends_no_metadata(self, singulr_guardrail):
"""Regression: with no alias found, metadata must be omitted (None),