scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 9

fix(guardrails): send rich user info to singulr

Browse source
This commit is contained in:
aniket-kardile 2026-08-26 13:59:58 +05:30
parent d3367c614e
commit ab79bc3a80
No known key found for this signature in database
2 changed files with 116 additions and 16 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -109,20 +109,38 @@ class SingulrGuardrail(CustomGuardrail):
return SingulrGuardrailConfigModel
@staticmethod
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
def _metadata_containers(request_data: Mapping[str, Any]) -> tuple[Mapping[str, Any], ...]:
"""Candidate metadata dicts to check, in priority order.
Most call paths put metadata at the top level of ``request_data``
(``litellm_metadata`` or ``metadata``). ``post_mcp_call`` instead hands
us ``litellm_logging_obj.model_call_details``, which nests it under
``litellm_params`` instead, so that's checked as a fallback.
"""
litellm_params: Final = request_data.get("litellm_params") or _EMPTY_MAPPING
return tuple(
container
for container in (
request_data.get("litellm_metadata"),
request_data.get("metadata"),
litellm_params.get("litellm_metadata") if litellm_params else None,
litellm_params.get("metadata") if litellm_params else None,
)
if container
)
@classmethod
def _resolve_metadata_value(cls, request_data: Mapping[str, Any], key: str) -> str | None:
for container in cls._metadata_containers(request_data=request_data):
value: Final = container.get(key)
if value:
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
@classmethod
def _resolve_user_role_from_request_data(cls, request_data: Mapping[str, Any]) -> str | None:
for container in cls._metadata_containers(request_data=request_data):
auth: Final = container.get("user_api_key_auth")
if isinstance(auth, UserAPIKeyAuth) and auth.user_role:
return auth.user_role.value
return None
@ -375,20 +393,24 @@ class SingulrGuardrail(CustomGuardrail):
try:
messages: Final = kwargs.get("messages") or ()
if messages:
request_metadata: Final = self._build_metadata(request_data=kwargs)
singulr_req_obj = SingulrGuardrailPayload(
correlation_id=kwargs.get("litellm_call_id"),
model_name=kwargs.get("model"),
guardrail_scope="request",
messages=messages,
metadata=request_metadata,
)
payload_req = singulr_req_obj.model_dump(mode="json")
await self._call_api(payload_req)
if result:
response_metadata: Final = self._build_metadata(request_data=kwargs)
singulr_res_obj = SingulrGuardrailPayload(
correlation_id=kwargs.get("litellm_call_id"),
guardrail_scope="response",
response=result,
metadata=response_metadata,
)
try:
payload = singulr_res_obj.model_dump(mode="json")
@ -398,6 +420,7 @@ class SingulrGuardrail(CustomGuardrail):
"correlation_id": kwargs.get("litellm_call_id"),
"guardrail_scope": "response",
"response": str(result),
"metadata": response_metadata,
}
await self._call_api(payload)
except GuardrailRaisedException:

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

@ -368,9 +368,7 @@ class TestSingulrRequestPayload:
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
}
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):
@ -662,6 +660,64 @@ class TestSingulrMcpResponse:
input_type="response",
)
@pytest.mark.asyncio
async def test_mcp_response_resolves_metadata_from_nested_litellm_params(self, singulr_guardrail):
"""post_mcp_call hands apply_guardrail litellm_logging_obj.model_call_details,
which nests metadata under litellm_params instead of at the top level."""
resp = _make_response({"should_block": False})
auth = UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER_VIEW_ONLY)
request_data = {
"call_type": "call_mcp_tool",
"mcp_tool_name": "search_docs",
"litellm_params": {
"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:
await singulr_guardrail.apply_guardrail(
inputs={"texts": ["Result: password reset link sent."]},
request_data=request_data,
input_type="response",
)
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",
"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
async def test_mcp_response_prefers_top_level_metadata_over_nested_litellm_params(self, singulr_guardrail):
resp = _make_response({"should_block": False})
request_data = {
"call_type": "call_mcp_tool",
"mcp_tool_name": "search_docs",
"litellm_metadata": {"user_api_key_alias": "top-level-alias"},
"litellm_params": {"metadata": {"user_api_key_alias": "nested-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="response",
)
sent_payload = mock_post.call_args.kwargs["json"]
assert sent_payload["metadata"] == {"user_api_key_alias": "top-level-alias"}
# ---------------------------------------------------------------------------
# apply_guardrail dispatch (request vs response vs unknown input_type)
@ -703,6 +759,25 @@ class TestSingulrLoggingHook:
assert response_payload["guardrail_scope"] == "response"
assert response_payload["response"] == result
@pytest.mark.asyncio
async def test_forwards_user_metadata_in_both_request_and_response_payloads(self, singulr_guardrail):
resp = _make_response({"should_block": False})
kwargs = {
"messages": [{"role": "user", "content": "hi"}],
"model": "gpt-4o",
"litellm_call_id": "call-1",
"litellm_metadata": {"user_api_key_alias": "my-key-alias", "user_api_key_org_id": "org-123"},
}
result = {"choices": [{"finish_reason": "stop", "message": {"content": "hello there"}}]}
with patch.object(singulr_guardrail.async_handler, "post", return_value=resp) as mock_post:
await singulr_guardrail.async_logging_hook(kwargs=kwargs, result=result, call_type="acompletion")
request_payload = mock_post.call_args_list[0].kwargs["json"]
response_payload = mock_post.call_args_list[1].kwargs["json"]
expected_metadata = {"user_api_key_alias": "my-key-alias", "user_api_key_org_id": "org-123"}
assert request_payload["metadata"] == expected_metadata
assert response_payload["metadata"] == expected_metadata
@pytest.mark.asyncio
async def test_forwards_a_real_model_response_without_swallowing_it(self, singulr_guardrail):
"""Regression: a normal completion callback passes a ModelResponse, not a
@ -729,11 +804,13 @@ class TestSingulrLoggingHook:
def __repr__(self) -> str:
return "<Unserializable>"
kwargs = {"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.async_logging_hook(kwargs={}, result=Unserializable(), call_type="acompletion")
await singulr_guardrail.async_logging_hook(kwargs=kwargs, result=Unserializable(), call_type="acompletion")
response_payload = mock_post.call_args.kwargs["json"]
assert response_payload["response"] == "<Unserializable>"
assert response_payload["metadata"] == {"user_api_key_alias": "my-key-alias"}
@pytest.mark.asyncio
async def test_no_messages_and_no_result_skips_both_api_calls(self, singulr_guardrail):