diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index 57e8323cdc3..7d04b1df95e 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -31,7 +31,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionUserMessage from litellm.types.utils import ( ChatCompletionMessageToolCall, Function, @@ -103,6 +103,8 @@ class _RequestContextSource(Protocol): @overload def get(self, key: Literal["optional_params"], /) -> Mapping[str, object] | None: ... @overload + def get(self, key: Literal["metadata", "litellm_metadata"], /) -> Mapping[str, object] | None: ... + @overload def get(self, key: str, /) -> object | None: ... def __contains__(self, key: object, /) -> bool: ... def __getitem__(self, key: str, /) -> object: ... @@ -441,7 +443,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): ) -> GenericGuardrailAPIInputs: """Send the (normalized) prompt to the before_prompt webhook and raise if the prompt is blocked.""" - messages = inputs.get("structured_messages") + messages: Sequence[AllMessageValues] | None = inputs.get("structured_messages") if not messages: # For non-chat request types (e.g. /v1/completions), litellm # supplies the prompt as ``texts`` with no structured_messages. @@ -450,11 +452,11 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): if texts: joined: Final = "\n".join(t for t in texts if t) if joined: - messages = [{"role": "user", "content": joined}] + messages = (ChatCompletionUserMessage(role="user", content=joined),) if not messages: return inputs - payload: Final = self._build_prompt_moderation_payload(inputs, request_data) + payload: Final = self._build_prompt_moderation_payload(inputs, messages, request_data) service_response: Final = await self._post_to_prompt_moderation_endpoint(payload) refusal: Final = self._extract_prompt_refusal(service_response) if refusal is None: @@ -616,7 +618,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod def _build_prompt_moderation_payload( inputs: GenericGuardrailAPIInputs, - request_data: Mapping[str, object], + messages: Sequence[AllMessageValues | None], + request_data: _RequestContextSource, ) -> Mapping[str, object]: """Build the bare OpenAI request the before_prompt webhook consumes. @@ -637,10 +640,11 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): correlation_key: Final = request_data.get("correlation_key") or request_data.get("litellm_call_id") return { "model": inputs.get("model") or request_data.get("model") or "", - "messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")), + "messages": RubrikLogger._flatten_messages_for_moderation(messages), **({"tools": tools} if tools is not None else _EMPTY_MAPPING), **({"user": user} if user else _EMPTY_MAPPING), **({"correlation_key": correlation_key} if correlation_key else _EMPTY_MAPPING), + **RubrikLogger._rac_agent_metadata(request_data), } @staticmethod @@ -671,7 +675,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): else: tools = optional_params.get("tools") - # The response moderation service consumes only messages/model/tools. + # The response moderation service consumes only messages/model/tools + # (plus metadata.rac_agent_id when set). # Don't forward proxy_server_request -- in litellm >=1.83 its ``body`` # snapshot carries a UserAPIKeyAuth instance that breaks json.dumps, # silently fail-opening the guardrail. @@ -679,8 +684,31 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): "messages": call_details.get("messages"), "model": call_details.get("model"), "tools": tools, + **RubrikLogger._rac_agent_metadata(request_data), } + @staticmethod + def _rac_agent_metadata(request_data: _RequestContextSource) -> Mapping[str, object]: + """Forward ``rac_agent_id`` from the caller's key or team metadata. + + The webhooks use it to choose which agent-scoped policies apply, so it + must come from configuration an admin controls, never from the request + body: a caller who could set it would select any agent's policies, + including one a policy exempts. ``_get_admin_metadata`` reads only the + ``user_api_key_metadata`` / ``user_api_key_team_metadata`` the proxy + writes from the authenticated key (key overrides team); the proxy strips + client-supplied ``user_api_key_*`` fields from both metadata dicts. + """ + rac_agent_id: Final[object] = RubrikLogger._admin_rac_agent_id(request_data) + if not isinstance(rac_agent_id, str) or not rac_agent_id: + return _EMPTY_MAPPING + return {"metadata": {"rac_agent_id": rac_agent_id}} + + @staticmethod + def _admin_rac_agent_id(request_data: _RequestContextSource) -> object: + """``rac_agent_id`` from the merged key/team admin metadata, untyped.""" + return RubrikLogger._get_admin_metadata(request_data).get("rac_agent_id") # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportArgumentType] # upstream helper takes and returns a bare dict; caller narrows to str + @staticmethod def _sanitize_proxy_server_request(proxy_server_request: Mapping[str, object] | str | None) -> object: """Allowlist only routing fields (``url``, ``method``) when forwarding diff --git a/tests/unit/integrations/test_rubrik.py b/tests/unit/integrations/test_rubrik.py index f8aec70a2f7..cc9e03ba4b9 100644 --- a/tests/unit/integrations/test_rubrik.py +++ b/tests/unit/integrations/test_rubrik.py @@ -5,6 +5,7 @@ Covers initialization, apply_guardrail (prompt moderation + response/tool blocking), batch logging, and Anthropic format handling. """ +import json import os from typing import Any, Dict from unittest.mock import AsyncMock, Mock, patch @@ -534,6 +535,22 @@ def _mock_service_response(response_json): return mock_client +def _capturing_service(response_json): + """Like ``_mock_service_response`` but also records each posted JSON body.""" + captured: list = [] + + async def mock_post(*_args, **kwargs): + captured.append(kwargs.get("json")) + mock_resp = Mock() + mock_resp.json.return_value = response_json + mock_resp.raise_for_status = Mock() + return mock_resp + + mock_client = AsyncMock() + mock_client.post = mock_post + return mock_client, captured + + def _echo_service(): """Create a mock tool blocking client that echoes the payload back.""" @@ -1230,6 +1247,35 @@ class TestPromptModeration: ) assert result is inputs + async def test_text_completion_prompt_is_sent_as_user_message(self, handler): + """/v1/completions requests carry the prompt as ``texts`` with no + structured_messages; the webhook must still receive it.""" + handler.moderation_client, captured = _capturing_service({}) + + await handler.apply_guardrail( + inputs={"texts": ["write a poem", "", "about cats"]}, + request_data={}, + input_type="request", + ) + + assert list(captured[0]["messages"]) == [ + {"role": "user", "content": "write a poem\nabout cats"} + ] + + async def test_structured_messages_take_precedence_over_texts(self, handler): + handler.moderation_client, captured = _capturing_service({}) + + await handler.apply_guardrail( + inputs={ + "texts": ["ignored"], + "structured_messages": [{"role": "user", "content": "hello"}], + }, + request_data={}, + input_type="request", + ) + + assert list(captured[0]["messages"]) == [{"role": "user", "content": "hello"}] + async def test_prompt_moderation_stashes_logging_obj_on_block(self, handler): """On a prompt block, _stash_block_context must set the blocked flag.""" inputs = { @@ -1378,7 +1424,7 @@ class TestBuildPromptModerationPayload: "structured_messages": [{"role": "user", "content": "hi"}], "tools": [{"type": "function", "function": {"name": "fn"}}], } - payload = RubrikLogger._build_prompt_moderation_payload(inputs, {}) + payload = RubrikLogger._build_prompt_moderation_payload(inputs, inputs["structured_messages"], {}) assert payload["tools"] == [{"type": "function", "function": {"name": "fn"}}] def test_payload_includes_user_when_present(self): @@ -1386,7 +1432,7 @@ class TestBuildPromptModerationPayload: "structured_messages": [{"role": "user", "content": "hi"}], } request_data = {"user": "alice"} - payload = RubrikLogger._build_prompt_moderation_payload(inputs, request_data) + payload = RubrikLogger._build_prompt_moderation_payload(inputs, inputs["structured_messages"], request_data) assert payload["user"] == "alice" def test_payload_uses_explicit_correlation_key(self): @@ -1395,22 +1441,145 @@ class TestBuildPromptModerationPayload: "correlation_key": "corr-123", "litellm_call_id": "litellm-456", } - payload = RubrikLogger._build_prompt_moderation_payload(inputs, request_data) + payload = RubrikLogger._build_prompt_moderation_payload(inputs, inputs["structured_messages"], request_data) assert payload["correlation_key"] == "corr-123" def test_payload_falls_back_to_litellm_call_id(self): """When correlation_key is absent, litellm_call_id is used.""" inputs = {"structured_messages": [{"role": "user", "content": "hi"}]} request_data = {"litellm_call_id": "litellm-789"} - payload = RubrikLogger._build_prompt_moderation_payload(inputs, request_data) + payload = RubrikLogger._build_prompt_moderation_payload(inputs, inputs["structured_messages"], request_data) assert payload["correlation_key"] == "litellm-789" def test_payload_omits_optional_fields_when_absent(self): inputs = {"structured_messages": [{"role": "user", "content": "hi"}]} - payload = RubrikLogger._build_prompt_moderation_payload(inputs, {}) + payload = RubrikLogger._build_prompt_moderation_payload(inputs, inputs["structured_messages"], {}) assert "tools" not in payload assert "user" not in payload assert "correlation_key" not in payload + assert "metadata" not in payload + + +# -- rac_agent_id forwarding -------------------------------------------------- + + +def _proxy_metadata( + key_meta: Dict[str, Any] | None = None, + team_meta: Dict[str, Any] | None = None, + **caller_fields: Any, +) -> Dict[str, Any]: + """Request metadata as the proxy hands it to guardrails: the authenticated + key/team metadata under ``user_api_key_*``, next to fields the caller sent""" + return { + **caller_fields, + "headers": {"authorization": "Bearer sk-secret"}, + "litellm_parent_otel_span": object(), + "user_api_key_metadata": key_meta or {}, + "user_api_key_team_metadata": team_meta or {}, + } + + +def _prompt_and_response_metadata(request_data: Dict[str, Any]): + inputs = {"structured_messages": [{"role": "user", "content": "hi"}]} + prompt_payload = RubrikLogger._build_prompt_moderation_payload( + inputs, inputs["structured_messages"], request_data + ) + response_request = RubrikLogger._extract_request_data( + {"model": "gpt-4"}, request_data + ) + return prompt_payload.get("metadata"), response_request.get("metadata") + + +class TestRacAgentIdForwarding: + def test_forwards_rac_agent_id_from_key_metadata_only(self): + request_data = {"metadata": _proxy_metadata(key_meta={"rac_agent_id": "agent-42", "other": 1})} + assert _prompt_and_response_metadata(request_data) == ( + {"rac_agent_id": "agent-42"}, + {"rac_agent_id": "agent-42"}, + ) + + def test_forwards_rac_agent_id_from_team_metadata(self): + request_data = {"metadata": _proxy_metadata(team_meta={"rac_agent_id": "team-agent"})} + assert _prompt_and_response_metadata(request_data) == ( + {"rac_agent_id": "team-agent"}, + {"rac_agent_id": "team-agent"}, + ) + + def test_key_metadata_overrides_team_metadata(self): + request_data = { + "metadata": _proxy_metadata( + key_meta={"rac_agent_id": "key-agent"}, + team_meta={"rac_agent_id": "team-agent"}, + ) + } + assert _prompt_and_response_metadata(request_data)[0] == {"rac_agent_id": "key-agent"} + + def test_admin_metadata_under_litellm_metadata(self): + """/v1/messages and /v1/responses put proxy metadata under litellm_metadata""" + request_data = {"litellm_metadata": _proxy_metadata(key_meta={"rac_agent_id": "agent-42"})} + assert _prompt_and_response_metadata(request_data)[0] == {"rac_agent_id": "agent-42"} + + @pytest.mark.parametrize("container", ["metadata", "litellm_metadata"]) + def test_caller_supplied_rac_agent_id_is_ignored(self, container): + """A caller must not choose the policy scope: the request-body value is + never forwarded, even when the key has no agent configured""" + request_data = {container: _proxy_metadata(rac_agent_id="exempt-agent")} + assert _prompt_and_response_metadata(request_data) == (None, None) + + def test_caller_cannot_override_the_keys_agent(self): + request_data = { + "metadata": _proxy_metadata( + key_meta={"rac_agent_id": "assigned-agent"}, + rac_agent_id="exempt-agent", + ), + "litellm_metadata": {"rac_agent_id": "exempt-agent"}, + } + assert _prompt_and_response_metadata(request_data) == ( + {"rac_agent_id": "assigned-agent"}, + {"rac_agent_id": "assigned-agent"}, + ) + + @pytest.mark.parametrize( + "request_data", + [ + {}, + {"metadata": None}, + {"metadata": "not-a-dict"}, + {"metadata": _proxy_metadata()}, + {"metadata": _proxy_metadata(key_meta={"rac_agent_id": ""})}, + {"metadata": _proxy_metadata(key_meta={"rac_agent_id": {"nested": "object"}})}, + {"metadata": _proxy_metadata(key_meta={"rac_agent_id": 42})}, + ], + ) + def test_metadata_omitted_without_a_valid_admin_rac_agent_id(self, request_data): + assert _prompt_and_response_metadata(request_data) == (None, None) + + @pytest.mark.asyncio + async def test_both_webhooks_receive_rac_agent_id_and_payloads_serialize( + self, handler + ): + request_data = {"metadata": _proxy_metadata(key_meta={"rac_agent_id": "agent-42"})} + handler.moderation_client, prompt_captured = _capturing_service({}) + await handler.apply_guardrail( + inputs={"structured_messages": [{"role": "user", "content": "hi"}]}, + request_data=request_data, + input_type="request", + ) + handler.moderation_client, response_captured = _capturing_service( + {"choices": [{"message": {"role": "assistant", "content": "hello"}}]} + ) + await handler.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data=request_data, + input_type="response", + ) + + assert prompt_captured[0]["metadata"] == {"rac_agent_id": "agent-42"} + assert response_captured[0]["request"]["metadata"] == { + "rac_agent_id": "agent-42" + } + json.dumps(prompt_captured[0]) + json.dumps(response_captured[0]) # -- _extract_request_data tools preference -----------------------------------