mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-07 07:58:09 +00:00
fix(guardrails): align the max_messages window and stop identity being taken from the request body
Three findings from review of the previous commit. max_messages applied the same number to two lists that count different things: texts is fragment-based (a multimodal turn contributes several entries, a tool-call-only turn contributes none) while structured_messages is message-based, so the two views of the payload could cover different parts of the conversation, sending text from an omitted turn or dropping text from a retained one. The message window is now the source of truth and the text window is the number of fragments those retained messages contain, counted with a role-agnostic helper so an unenumerated role cannot silently shrink it. /guardrails/apply_guardrail built request_data straight from the caller's metadata, so a caller could claim a user_api_key_alias or user_api_key_team_id that a guardrail exempts and get its text back unscanned. The authenticated identity is now layered on top of the caller's metadata via the same sanitizer the in-line path uses, so what authentication established wins. Caller metadata that is not identity still passes through untouched, and every guardrail on that route now sees the caller information it already sees in-line. The per_session dedup cache was keyed only on the caller-supplied session id, so whoever claimed an id first suppressed another tenant's entry for the same id. The key is now namespaced by the authenticated key hash, falling back to the team id.
This commit is contained in:
parent
c3ad8ac34e
commit
155e4080c4
7 changed files with 248 additions and 38 deletions
|
|
@ -112,6 +112,20 @@ def iter_message_text(data: dict[str, Any]) -> Iterator[str]:
|
|||
yield from _iter_text_parts_in_content(message.get("content"))
|
||||
|
||||
|
||||
def iter_messages_text(messages: Sequence[Mapping[str, object]]) -> Iterator[str]:
|
||||
"""Yield every text fragment carried by a message list, whatever the role.
|
||||
|
||||
Role-agnostic on purpose: a caller counting what a set of messages
|
||||
contributes (payload windowing, size accounting) must not miss a turn whose
|
||||
role it did not think to enumerate.
|
||||
"""
|
||||
for message in messages:
|
||||
# The annotation is the contract; the values arrive from provider
|
||||
# translation handlers, so a stray non-message entry must not raise.
|
||||
if isinstance(message, Mapping): # pyright: ignore[reportUnnecessaryIsInstance] # untyped upstream data
|
||||
yield from _iter_text_parts_in_content(message.get("content"))
|
||||
|
||||
|
||||
def iter_role_text(messages: Sequence[Mapping[str, object]], roles: Collection[str]) -> Iterator[str]:
|
||||
"""Yield every text fragment carried by messages whose role is in ``roles``.
|
||||
|
||||
|
|
|
|||
|
|
@ -2288,6 +2288,7 @@ async def apply_guardrail(
|
|||
executor as thread_pool_executor,
|
||||
)
|
||||
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.proxy.proxy_server import (
|
||||
general_settings,
|
||||
proxy_config,
|
||||
|
|
@ -2332,9 +2333,16 @@ async def apply_guardrail(
|
|||
if litellm_logging_obj is not None:
|
||||
_patch_logging_obj_for_guardrail(litellm_logging_obj, request)
|
||||
|
||||
# The caller supplies this metadata, so the authenticated identity is
|
||||
# layered on top of it: a guardrail keyed on user_api_key_alias or
|
||||
# user_api_key_team_id (exemptions, per-tenant behavior) must not be
|
||||
# steerable by what the request body claims to be.
|
||||
authenticated_metadata: Final = LiteLLMProxyRequestSetup.get_sanitized_user_information_from_key(
|
||||
user_api_key_dict
|
||||
)
|
||||
request_data: Final[dict] = {
|
||||
**({"messages": request.messages} if request.messages is not None else {}),
|
||||
**({"metadata": request.metadata} if request.metadata is not None else {}),
|
||||
"metadata": {**(request.metadata or {}), **(authenticated_metadata or {})},
|
||||
}
|
||||
_input_type: Final = _resolve_guardrail_input_type(active_guardrail, request.input_type)
|
||||
guardrailed_inputs: Final = await active_guardrail.apply_guardrail(
|
||||
|
|
|
|||
|
|
@ -605,7 +605,17 @@ class GenericGuardrailAPI(CustomGuardrail):
|
|||
input_type=input_type,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
if self._record_scope.should_suppress(get_session_id_from_request_data(request_data or {})):
|
||||
# The session id is caller-supplied, so the dedup key is namespaced by the
|
||||
# authenticated caller: the key hash is unique per virtual key, with the
|
||||
# team id as the fallback. Both come from auth, not from the body.
|
||||
identity: Final = (
|
||||
self._extract_user_api_key_metadata(request_data) if request_data else GenericGuardrailAPIMetadata()
|
||||
)
|
||||
session_id: Final = get_session_id_from_request_data(request_data) if request_data else None
|
||||
if self._record_scope.should_suppress(
|
||||
session_id,
|
||||
tenant=identity.get("user_api_key_hash") or identity.get("user_api_key_team_id"),
|
||||
):
|
||||
suppress_guardrail_information_record()
|
||||
return result
|
||||
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from pydantic import JsonValue
|
|||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy.guardrails._content_utils import (
|
||||
iter_messages_text,
|
||||
map_messages_image_urls,
|
||||
map_messages_text,
|
||||
)
|
||||
|
|
@ -121,7 +122,39 @@ def _shape_text(text: str, policy: PayloadPolicy) -> str:
|
|||
return stripped[: policy.max_text_chars] if policy.max_text_chars is not None else stripped
|
||||
|
||||
|
||||
def _shape_texts(texts: Sequence[str], policy: PayloadPolicy) -> tuple[JsonValue, frozenset[int]]:
|
||||
def _retained_messages(messages: JsonValue, policy: PayloadPolicy) -> list[JsonValue] | None:
|
||||
"""The message window ``max_messages`` keeps, or None when there is nothing to window."""
|
||||
if not isinstance(messages, list) or not messages:
|
||||
return None
|
||||
if policy.max_messages is None:
|
||||
return messages
|
||||
return messages[-policy.max_messages :]
|
||||
|
||||
|
||||
def _text_window_size(retained: list[JsonValue] | None, policy: PayloadPolicy, total_texts: int) -> int:
|
||||
"""How many trailing ``texts`` entries belong to the retained messages.
|
||||
|
||||
``texts`` is fragment-based (a multimodal turn contributes several entries, a
|
||||
tool-call-only turn contributes none) while ``max_messages`` counts messages,
|
||||
so applying the same number to both would leave the two views of the payload
|
||||
covering different parts of the conversation. Counting the fragments inside
|
||||
the retained messages keeps them aligned. Without structured messages to
|
||||
count (a response payload, an embedding call) the message count is the only
|
||||
bound available, so it is applied to the fragments directly.
|
||||
"""
|
||||
if policy.max_messages is None:
|
||||
return total_texts
|
||||
if retained is None:
|
||||
return min(policy.max_messages, total_texts)
|
||||
fragments: Final = sum(1 for _ in iter_messages_text(retained))
|
||||
return min(fragments, total_texts)
|
||||
|
||||
|
||||
def _shape_texts(
|
||||
texts: Sequence[str],
|
||||
policy: PayloadPolicy,
|
||||
window_size: int,
|
||||
) -> tuple[JsonValue, frozenset[int]]:
|
||||
"""Window and rewrite the flat text list, and report which indices changed.
|
||||
|
||||
``max_messages`` has to bound ``texts`` as well as ``structured_messages``:
|
||||
|
|
@ -131,19 +164,18 @@ def _shape_texts(texts: Sequence[str], policy: PayloadPolicy) -> tuple[JsonValue
|
|||
shifts every position, so the caller treats the whole list as unusable for
|
||||
write-back (see ``merge_guardrailed_texts``).
|
||||
"""
|
||||
windowed: Final = texts[-policy.max_messages :] if policy.max_messages is not None else texts
|
||||
windowed: Final = texts[len(texts) - window_size :] if window_size < len(texts) else texts
|
||||
shaped: Final = [_shape_text(text, policy) for text in windowed] # mutable-ok: JSON texts is an array
|
||||
if len(windowed) != len(texts):
|
||||
return shaped, frozenset(range(len(texts)))
|
||||
return shaped, frozenset(index for index, text in enumerate(texts) if text != shaped[index])
|
||||
|
||||
|
||||
def _shape_messages(messages: JsonValue, policy: PayloadPolicy) -> JsonValue:
|
||||
if not isinstance(messages, list) or not messages:
|
||||
def _shape_messages(retained: list[JsonValue] | None, policy: PayloadPolicy) -> JsonValue:
|
||||
if retained is None:
|
||||
return None
|
||||
windowed: Final = messages[-policy.max_messages :] if policy.max_messages is not None else messages
|
||||
texted: Final = (
|
||||
map_messages_text(windowed, lambda text: _shape_text(text, policy)) if policy.shapes_text else windowed
|
||||
map_messages_text(retained, lambda text: _shape_text(text, policy)) if policy.shapes_text else retained
|
||||
)
|
||||
# send_images=False has to cover the inline image parts too, otherwise the
|
||||
# base64 payload still leaves LiteLLM inside structured_messages. The part is
|
||||
|
|
@ -170,8 +202,12 @@ def shape_payload(
|
|||
|
||||
texts: Final = request.texts
|
||||
texts_sent: Final = bool(texts) and "texts" not in exclude
|
||||
shaped_texts, altered_indices = _shape_texts(texts or (), policy) if texts_sent else (None, frozenset[int]())
|
||||
shaped_messages: Final = _shape_messages(dumped.get("structured_messages"), policy)
|
||||
retained: Final = _retained_messages(dumped.get("structured_messages"), policy)
|
||||
window_size: Final = _text_window_size(retained, policy, len(texts or ()))
|
||||
shaped_texts, altered_indices = (
|
||||
_shape_texts(texts or (), policy, window_size) if texts_sent else (None, frozenset[int]())
|
||||
)
|
||||
shaped_messages: Final = _shape_messages(retained, policy)
|
||||
|
||||
payload: Final = { # mutable-ok: the POST body is a JSON object
|
||||
**dumped,
|
||||
|
|
|
|||
|
|
@ -44,13 +44,17 @@ class RecordScope:
|
|||
def scope(self) -> GuardrailInformationScope:
|
||||
return self._scope
|
||||
|
||||
def should_suppress(self, session_id: str | None) -> bool:
|
||||
def should_suppress(self, session_id: str | None, *, tenant: str | None = None) -> bool:
|
||||
"""Whether this successful call should skip its logging entry.
|
||||
|
||||
``per_session`` claims the session on its first call, so the calls that
|
||||
follow it suppress. A session id is required: without one there is
|
||||
nothing to dedup against, so the call records as it would under
|
||||
``per_call`` rather than silently dropping every entry.
|
||||
|
||||
The session id is caller-supplied, so the key is namespaced by the
|
||||
authenticated caller: otherwise whoever claims an id first would
|
||||
suppress another tenant's telemetry for the same id.
|
||||
"""
|
||||
if self._scope == "per_call":
|
||||
return False
|
||||
|
|
@ -58,7 +62,8 @@ class RecordScope:
|
|||
return True
|
||||
if session_id is None:
|
||||
return False
|
||||
if self._recorded_sessions.get_cache(session_id) is True:
|
||||
key: Final = f"{tenant or ''}::{session_id}"
|
||||
if self._recorded_sessions.get_cache(key) is True:
|
||||
return True
|
||||
self._recorded_sessions.set_cache(session_id, True, ttl=_SESSION_CACHE_TTL_SECONDS)
|
||||
self._recorded_sessions.set_cache(key, True, ttl=_SESSION_CACHE_TTL_SECONDS)
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -3190,3 +3190,116 @@ class TestConfigValidationWarnings:
|
|||
|
||||
assert "tools" not in handler.payloads[0]
|
||||
assert result["tools"] == original_tools
|
||||
|
||||
|
||||
class TestMaxMessagesWindowAlignment:
|
||||
"""The texts window must cover the same turns as the message window.
|
||||
|
||||
`texts` is fragment-based and `max_messages` counts messages, so applying the
|
||||
same number to both would let the two views of the payload describe different
|
||||
parts of the conversation.
|
||||
"""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_multipart_turn_does_not_leak_earlier_messages(self):
|
||||
"""A retained turn with two text parts must not push an omitted turn's text in."""
|
||||
handler = _RecordingHandler()
|
||||
guardrail = _make_guardrail(handler, max_messages=1)
|
||||
messages = [
|
||||
{"role": "user", "content": "old turn that must not be sent"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": [
|
||||
{"type": "text", "text": "kept part one"},
|
||||
{"type": "text", "text": "kept part two"},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={
|
||||
"texts": ["old turn that must not be sent", "kept part one", "kept part two"],
|
||||
"structured_messages": messages,
|
||||
},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
payload = handler.payloads[0]
|
||||
assert payload["texts"] == ["kept part one", "kept part two"]
|
||||
assert [m["role"] for m in payload["structured_messages"]] == ["assistant"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_textless_turn_does_not_drop_retained_text(self):
|
||||
"""A retained tool-call-only turn contributes no fragment, so the window shrinks."""
|
||||
handler = _RecordingHandler()
|
||||
guardrail = _make_guardrail(handler, max_messages=2)
|
||||
messages = [
|
||||
{"role": "user", "content": "first"},
|
||||
{"role": "user", "content": "second"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": [{"id": "c1", "type": "function", "function": {"name": "run", "arguments": "{}"}}],
|
||||
},
|
||||
]
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["first", "second"], "structured_messages": messages},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
|
||||
payload = handler.payloads[0]
|
||||
# The window keeps the last two messages, which carry exactly one fragment.
|
||||
assert payload["texts"] == ["second"]
|
||||
assert [m["role"] for m in payload["structured_messages"]] == ["user", "assistant"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_response_payload_without_messages_still_bounded(self):
|
||||
"""Nothing to align against, so the message count bounds the fragments."""
|
||||
handler = _RecordingHandler()
|
||||
guardrail = _make_guardrail(handler, max_messages=2)
|
||||
|
||||
await guardrail.apply_guardrail(
|
||||
inputs={"texts": ["a", "b", "c", "d"]},
|
||||
request_data={},
|
||||
input_type="response",
|
||||
)
|
||||
|
||||
assert handler.payloads[0]["texts"] == ["c", "d"]
|
||||
|
||||
|
||||
class TestSessionScopeIsolation:
|
||||
"""per_session dedup must not let one caller suppress another's telemetry."""
|
||||
|
||||
@staticmethod
|
||||
def _request_data(session_id: str, key_hash: str) -> dict:
|
||||
return {
|
||||
"litellm_session_id": session_id,
|
||||
"metadata": {"user_api_key_hash": key_hash},
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_session_id_from_another_key_still_records(self):
|
||||
handler = _RecordingHandler()
|
||||
guardrail = _make_guardrail(handler, guardrail_information_scope="per_session")
|
||||
first = self._request_data("shared-id", "hash-tenant-a")
|
||||
second = self._request_data("shared-id", "hash-tenant-b")
|
||||
|
||||
await guardrail.apply_guardrail(inputs={"texts": ["a"]}, request_data=first, input_type="request")
|
||||
await guardrail.apply_guardrail(inputs={"texts": ["b"]}, request_data=second, input_type="request")
|
||||
|
||||
assert len(_recorded_entries(first)) == 1
|
||||
assert len(_recorded_entries(second)) == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_key_and_session_still_dedups(self):
|
||||
handler = _RecordingHandler()
|
||||
guardrail = _make_guardrail(handler, guardrail_information_scope="per_session")
|
||||
data = self._request_data("shared-id", "hash-tenant-a")
|
||||
|
||||
for _ in range(3):
|
||||
await guardrail.apply_guardrail(inputs={"texts": ["a"]}, request_data=data, input_type="request")
|
||||
|
||||
assert len(_recorded_entries(data)) == 1
|
||||
|
|
|
|||
|
|
@ -1552,11 +1552,41 @@ async def test_apply_guardrail_forwards_metadata_to_guardrail(mocker):
|
|||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
mock_guardrail.apply_guardrail.assert_awaited_once_with(
|
||||
inputs={"texts": ["What are tax loopholes?"]},
|
||||
request_data={"metadata": {"forbidden_topics": ["tax"]}},
|
||||
input_type="request",
|
||||
call = mock_guardrail.apply_guardrail.await_args.kwargs
|
||||
assert call["inputs"] == {"texts": ["What are tax loopholes?"]}
|
||||
assert call["input_type"] == "request"
|
||||
assert call["request_data"]["metadata"]["forbidden_topics"] == ["tax"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_authenticated_identity_overrides_client_metadata(mocker):
|
||||
"""A caller must not be able to claim another identity in the request body.
|
||||
|
||||
Guardrails key exemptions and per-tenant behavior off user_api_key_* metadata,
|
||||
so the authenticated values have to win over whatever the body says.
|
||||
"""
|
||||
mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]})
|
||||
|
||||
request = ApplyGuardrailRequest(
|
||||
guardrail_name="test-guardrail",
|
||||
text="hello",
|
||||
metadata={
|
||||
"user_api_key_alias": "exempt-batch-worker",
|
||||
"user_api_key_team_id": "exempt-team",
|
||||
"forbidden_topics": ["tax"],
|
||||
},
|
||||
)
|
||||
await apply_guardrail(
|
||||
fastapi_request=mocker.Mock(),
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(key_alias="real-caller", team_id="real-team"),
|
||||
)
|
||||
|
||||
metadata = mock_guardrail.apply_guardrail.await_args.kwargs["request_data"]["metadata"]
|
||||
assert metadata["user_api_key_alias"] == "real-caller"
|
||||
assert metadata["user_api_key_team_id"] == "real-team"
|
||||
# Non-identity metadata the caller sent is still forwarded untouched.
|
||||
assert metadata["forbidden_topics"] == ["tax"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1578,33 +1608,28 @@ async def test_apply_guardrail_forwards_metadata_and_messages_together(mocker):
|
|||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
mock_guardrail.apply_guardrail.assert_awaited_once_with(
|
||||
inputs={"texts": ["What are tax loopholes?"]},
|
||||
request_data={
|
||||
"messages": messages,
|
||||
"metadata": {"forbidden_topics": ["tax"]},
|
||||
},
|
||||
input_type="request",
|
||||
)
|
||||
call = mock_guardrail.apply_guardrail.await_args.kwargs
|
||||
assert call["request_data"]["messages"] == messages
|
||||
assert call["request_data"]["metadata"]["forbidden_topics"] == ["tax"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apply_guardrail_omits_metadata_when_not_sent(mocker):
|
||||
"""Without metadata, request_data stays empty (backward-compatible)."""
|
||||
async def test_apply_guardrail_carries_authenticated_metadata_when_none_sent(mocker):
|
||||
"""request_data always carries the authenticated identity, even with no client
|
||||
metadata, so a guardrail sees the same caller information it sees in-line."""
|
||||
mock_guardrail = _patch_apply_guardrail_env(mocker, {"texts": ["ok"]})
|
||||
|
||||
request = ApplyGuardrailRequest(guardrail_name="test-guardrail", text="hello")
|
||||
await apply_guardrail(
|
||||
fastapi_request=mocker.Mock(),
|
||||
request=request,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
user_api_key_dict=UserAPIKeyAuth(key_alias="known-caller"),
|
||||
)
|
||||
|
||||
mock_guardrail.apply_guardrail.assert_awaited_once_with(
|
||||
inputs={"texts": ["hello"]},
|
||||
request_data={},
|
||||
input_type="request",
|
||||
)
|
||||
call = mock_guardrail.apply_guardrail.await_args.kwargs
|
||||
assert call["inputs"] == {"texts": ["hello"]}
|
||||
assert "messages" not in call["request_data"]
|
||||
assert call["request_data"]["metadata"]["user_api_key_alias"] == "known-caller"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1625,11 +1650,10 @@ async def test_apply_guardrail_forwards_explicit_empty_messages_and_metadata(moc
|
|||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
mock_guardrail.apply_guardrail.assert_awaited_once_with(
|
||||
inputs={"texts": ["hello"]},
|
||||
request_data={"messages": [], "metadata": {}},
|
||||
input_type="request",
|
||||
)
|
||||
call = mock_guardrail.apply_guardrail.await_args.kwargs
|
||||
assert call["request_data"]["messages"] == []
|
||||
# The explicitly-empty metadata is still forwarded, now carrying identity.
|
||||
assert "metadata" in call["request_data"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue