scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(guardrails/rubrik): moderate text-completion prompts and forward rac_agent_id (#49912)

Browse source 
* fix(guardrails/rubrik): moderate text-completion prompts and forward rac_agent_id

_moderate_prompt synthesized a user message from texts for /v1/completions
requests but built the payload from structured_messages, so the webhook
received an empty messages list and the prompt went unmoderated. Pass the
resolved messages into _build_prompt_moderation_payload.

Forward metadata.rac_agent_id to the before_prompt and after_completion
webhooks so agent-scoped policies apply. Only that key is forwarded; the
rest of request_data metadata carries non-serializable objects and inbound
headers

* fix(guardrails/rubrik): also read rac_agent_id from litellm_metadata

/v1/messages, /v1/responses and other LITELLM_METADATA_ROUTES keep proxy
metadata under litellm_metadata, so a caller on those routes may set
rac_agent_id there. Read metadata first, then litellm_metadata, forwarding
only a non-empty string

* fix(guardrails/rubrik): take rac_agent_id from key or team metadata only

The Rubrik webhooks use rac_agent_id to choose which agent-scoped
policies apply. Forwarding it from the request body let a caller pick any
agent's policies, including one a policy exempts. Read it from the
caller's admin-configured key metadata, falling back to team metadata,
via CustomGuardrail._get_admin_metadata. The proxy writes those from the
authenticated key and strips client-supplied user_api_key_* fields from
both metadata dicts, so a caller cannot set or override the value

---------

Co-authored-by: yucheng-berri <yucheng@berri.ai>
This commit is contained in:
Joseph Barker 2026-11-10 16:30:41 -08:00 • committed by GitHub
parent ebd5a15b22
commit 2c9068dbab
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 209 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
litellm/integrations/rubrik.py
View file

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

179
tests/unit/integrations/test_rubrik.py
View file

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