From ab79bc3a80f4cb2a11ba1b9fa61ea70ffbc535cc Mon Sep 17 00:00:00 2001 From: aniket-kardile Date: Wed, 26 Aug 2026 13:59:58 +0530 Subject: [PATCH] fix(guardrails): send rich user info to singulr --- .../guardrail_hooks/singulr/singulr.py | 47 +++++++--- .../guardrail_hooks/test_singulr.py | 85 ++++++++++++++++++- 2 files changed, 116 insertions(+), 16 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py index 4e018310d9b..fb92d1b8fa9 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py +++ b/litellm/proxy/guardrails/guardrail_hooks/singulr/singulr.py @@ -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: 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 9b789984684..457b7c5c7f4 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_singulr.py @@ -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 "" + 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"] == "" + 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):