scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-13 23:11:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 100

fix(router): read the affinity pin from the messages argument and strip bridge reasoning on a cross-group route

Browse source 
The encrypted_content_affinity check only read the Anthropic history from request_kwargs["messages"], so a caller that passes it through the callback's messages argument alone skipped the pin. Read the argument first and fall back to the kwargs.

When the minting deployment is not a candidate of the routed group, the base already strips the Responses input's encrypted reasoning; do the same for the bridge-tagged thinking blocks in Anthropic messages so the routed deployment gets the readable thinking text instead of ciphertext it cannot decrypt.
This commit is contained in:
mateo-berri 2026-09-09 18:31:08 -07:00
parent 4d4906e94e
commit 4716b46c24
4 changed files with 166 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

43
litellm/litellm_core_utils/prompt_templates/common_utils.py
View file

@ -6,7 +6,7 @@ import io
import json
import mimetypes
import re
from collections.abc import Iterable, Mapping, Sequence
from collections.abc import Iterable, Iterator, Mapping, Sequence
from itertools import groupby
from os import PathLike
from pathlib import Path
@ -1889,6 +1889,47 @@ def is_encrypted_reasoning_block(block: object) -> bool:
return _carries_encrypted_reasoning(_encrypted_reasoning_field(mapping))
def strip_encrypted_reasoning_from_messages(messages: object) -> None:
"""Drop the encrypted reasoning a routed deployment cannot decrypt from Anthropic-shaped
history, keeping the readable thinking text.
Mutates the content lists in place: the router's fallback snapshot shares these
message objects, so a rebound list would replay the stripped blocks on the fallback hop.
"""
if not isinstance(messages, list):
return
for content in _anthropic_content_lists(cast(list[object], messages)): # cast-ok: untyped client json
_strip_encrypted_reasoning_from_blocks(content)
def _anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]:
return (
cast(list[object], content) # cast-ok: narrowed by isinstance
for message in messages
if isinstance(message, Mapping)
for content in (cast(Mapping[str, object], message).get("content"),) # cast-ok: narrowed by isinstance
if isinstance(content, list)
)
def _strip_encrypted_reasoning_from_blocks(content: object) -> None:
blocks: Final = cast(list[object], content) # cast-ok: narrowed by the caller's isinstance
stripped: Final = tuple(_without_encrypted_reasoning_block(block) for block in blocks)
blocks[:] = (block for block in stripped if block is not None) # rebind-ok: list shared with fallback snapshot
def _without_encrypted_reasoning_block(block: object) -> object | None:
if not is_encrypted_reasoning_block(block):
return block
mapping: Final = cast(Mapping[str, object], block) # cast-ok: narrowed by is_encrypted_reasoning_block
if mapping.get("type") != "thinking" or not mapping.get("thinking"):
return None
kept: Final[dict[str, object]] = { # mutable-ok: thinking block rebuilt without the undecryptable signature
key: value for key, value in mapping.items() if key != "signature"
}
return kept
def _reasoning_replay_group_key(indexed_block: tuple[int, Mapping[str, object]]) -> str:
index, block = indexed_block
return f"encrypted:{index}" if is_encrypted_reasoning_block(block) else "summary"

14
litellm/router_utils/pre_call_checks/encrypted_content_affinity_check.py
View file

@ -48,7 +48,10 @@ from litellm.exceptions import (
ServiceUnavailableError,
)
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_content_of_block
from litellm.litellm_core_utils.prompt_templates.common_utils import (
encrypted_content_of_block,
strip_encrypted_reasoning_from_messages,
)
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.router_utils.cooldown_cache import CooldownCacheValue
from litellm.types.llms.openai import AllMessageValues
@ -274,8 +277,9 @@ class EncryptedContentAffinityCheck(CustomLogger):
parent_otel_span: Span | None = None,
) -> list[dict]:
"""
If the request ``input`` contains litellm-encoded item IDs, decode the
embedded ``model_id`` and pin the request to that deployment. Raises
If the request ``input`` contains litellm-encoded item IDs, or its Anthropic
``messages`` replay a bridge-tagged thinking block, decode the embedded
``model_id`` and pin the request to that deployment. Raises
``RateLimitError`` / ``ServiceUnavailableError`` when the originating
deployment is a member of the routed model group but currently unavailable
and no encryption-boundary peer exists, rather than dispatching a doomed
@ -304,9 +308,10 @@ class EncryptedContentAffinityCheck(CustomLogger):
request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] = True
request_input: Final = request_kwargs.get("input")
anthropic_messages: Final = messages or request_kwargs.get("messages")
model_id: Final = self._extract_model_id_from_input(
request_input
) or self._extract_model_id_from_anthropic_messages(request_kwargs.get("messages"))
) or self._extract_model_id_from_anthropic_messages(anthropic_messages)
if not model_id:
return typed_healthy_deployments
@ -363,6 +368,7 @@ class EncryptedContentAffinityCheck(CustomLogger):
model,
)
ResponsesAPIRequestUtils.strip_encrypted_reasoning_from_input(request_input)
strip_encrypted_reasoning_from_messages(anthropic_messages)
return typed_healthy_deployments
# The origin is a member of the routed group but currently unavailable (cooled down); fail fast

60
tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py
View file

@ -1,3 +1,4 @@
import copy
import functools
import json
import os
@ -20,6 +21,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import (
is_encrypted_reasoning_block,
responses_reasoning_items_from_thinking_blocks,
split_concatenated_json_objects,
strip_encrypted_reasoning_from_messages,
update_messages_with_model_file_ids,
)
@ -1567,7 +1569,9 @@ class TestEncryptedReasoningReplay:
def test_signature_round_trips_the_encrypted_content(self):
assert encrypted_content_from_signature(encrypted_reasoning_signature("gAAAA_bytes")) == "gAAAA_bytes"
@pytest.mark.parametrize("signature", [None, "", "ErcBCkgIValidAnthropicSignature", "litellm_encrypted_reasoning:", 7])
@pytest.mark.parametrize(
"signature", [None, "", "ErcBCkgIValidAnthropicSignature", "litellm_encrypted_reasoning:", 7]
)
def test_anything_else_is_not_encrypted_content(self, signature):
assert encrypted_content_from_signature(signature) is None
@ -1576,7 +1580,11 @@ class TestEncryptedReasoningReplay:
[{"type": "thinking", "thinking": "Plan.", "signature": encrypted_reasoning_signature("gAAAA_1")}]
)
assert items == (
{"type": "reasoning", "summary": [{"type": "summary_text", "text": "Plan."}], "encrypted_content": "gAAAA_1"},
{
"type": "reasoning",
"summary": [{"type": "summary_text", "text": "Plan."}],
"encrypted_content": "gAAAA_1",
},
)
def test_encrypted_redacted_block_replays_with_an_empty_summary(self):
@ -1596,7 +1604,10 @@ class TestEncryptedReasoningReplay:
]
)
assert items == (
{"type": "reasoning", "summary": [{"type": "summary_text", "text": "A."}, {"type": "summary_text", "text": "B."}]},
{
"type": "reasoning",
"summary": [{"type": "summary_text", "text": "A."}, {"type": "summary_text", "text": "B."}],
},
{"type": "reasoning", "summary": [{"type": "summary_text", "text": "C."}], "encrypted_content": "gAAAA_c"},
{"type": "reasoning", "summary": [{"type": "summary_text", "text": "D."}]},
)
@ -1621,3 +1632,46 @@ class TestEncryptedReasoningReplay:
)
def test_is_encrypted_reasoning_block(self, block, expected):
assert is_encrypted_reasoning_block(block) is expected
def test_strip_keeps_the_readable_thinking_and_drops_the_undecryptable_bytes(self):
assistant_content = [
{"type": "thinking", "thinking": "minted by Anthropic", "signature": "ErcBCkgIValid"},
{"type": "thinking", "thinking": "packed by the bridge", "signature": encrypted_reasoning_signature("g1")},
{"type": "redacted_thinking", "data": encrypted_reasoning_signature("g2")},
{"type": "thinking", "thinking": "", "signature": encrypted_reasoning_signature("g3")},
{"type": "text", "text": "answer"},
]
messages = [
{"role": "user", "content": "question"},
{"role": "assistant", "content": assistant_content},
{"role": "user", "content": [{"type": "text", "text": "follow-up"}]},
]
strip_encrypted_reasoning_from_messages(messages)
assert messages[1]["content"] is assistant_content
assert assistant_content == [
{"type": "thinking", "thinking": "minted by Anthropic", "signature": "ErcBCkgIValid"},
{"type": "thinking", "thinking": "packed by the bridge"},
{"type": "text", "text": "answer"},
]
assert messages[0] == {"role": "user", "content": "question"}
assert messages[2] == {"role": "user", "content": [{"type": "text", "text": "follow-up"}]}
@pytest.mark.parametrize(
"messages",
[
"not a list",
None,
[{"role": "user", "content": None}],
[{"role": "user", "content": "plain string"}],
["not a message"],
[{"role": "assistant", "content": [{"type": "thinking", "thinking": "x", "signature": "ErcBCkgIValid"}]}],
],
)
def test_strip_leaves_history_without_bridge_reasoning_untouched(self, messages):
before = copy.deepcopy(messages)
strip_encrypted_reasoning_from_messages(messages)
assert messages == before

73
tests/test_litellm/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py
View file

@ -1597,26 +1597,12 @@ async def test_encrypted_content_affinity_pins_anthropic_messages_replayed_throu
{"model_info": {"id": "openai-org-a"}, "litellm_params": {"model": "openai/gpt-5.1"}},
{"model_info": {"id": "openai-org-b"}, "litellm_params": {"model": "openai/gpt-5.1"}},
]
wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA_turn_one", "openai-org-b")
request_kwargs = {
"messages": [
{"role": "user", "content": "Solve the zebra puzzle"},
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"},
{"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"},
{"type": "text", "text": "The zebra owner lives in the green house."},
],
},
{"role": "user", "content": "And who drinks water?"},
],
}
request_kwargs = {"model": "gpt-5.1"}
pinned = await check.async_filter_deployments(
model="gpt-5.1",
healthy_deployments=deployments,
messages=None,
messages=_bridge_replayed_anthropic_messages(minted_by="openai-org-b"),
request_kwargs=request_kwargs,
)
@ -1624,6 +1610,61 @@ async def test_encrypted_content_affinity_pins_anthropic_messages_replayed_throu
assert request_kwargs["_encrypted_content_affinity_pinned"] is True
def _bridge_replayed_anthropic_messages(minted_by: str) -> list:
wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA_turn_one", minted_by)
return [
{"role": "user", "content": "Solve the zebra puzzle"},
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"},
{"type": "redacted_thinking", "data": f"litellm_encrypted_reasoning:{wrapped}"},
{
"type": "thinking",
"thinking": "The bridge packed this one",
"signature": f"litellm_encrypted_reasoning:{wrapped}",
},
{"type": "text", "text": "The zebra owner lives in the green house."},
],
},
{"role": "user", "content": "And who drinks water?"},
]
@pytest.mark.asyncio
async def test_encrypted_content_affinity_strips_bridge_reasoning_from_messages_routed_to_another_group():
"""
The /v1/messages twin of the tier-change case: the routed group holds no deployment
of the org that minted the reasoning, so the bridge-tagged blocks are stripped down
to their readable thinking text and the request dispatches to the routed pool.
"""
originating = _make_originating_mock(None, "key-a", model_name="gpt-reasoning-tier")
mock_router = _make_router_mock_with_cooldown(
originating, cooldown_entries=[], routed_group_model_ids=["openai-org-b"]
)
check = EncryptedContentAffinityCheck(router=mock_router)
routed_pool = [{"model_info": {"id": "openai-org-b"}, "litellm_params": {"model": "openai/gpt-5-nano"}}]
messages = _bridge_replayed_anthropic_messages(minted_by="openai-org-a")
assistant_content = messages[1]["content"]
request_kwargs = {"model": "gpt-5.1"}
result = await check.async_filter_deployments(
model="gpt-simple-tier",
healthy_deployments=routed_pool,
messages=messages,
request_kwargs=request_kwargs,
)
assert result is routed_pool
assert "_encrypted_content_affinity_pinned" not in request_kwargs
assert messages[1]["content"] is assistant_content
assert assistant_content == [
{"type": "thinking", "thinking": "Anthropic minted this one", "signature": "ErcCCpIBCBEYAipA"},
{"type": "thinking", "thinking": "The bridge packed this one"},
{"type": "text", "text": "The zebra owner lives in the green house."},
]
class TestStripEncryptedReasoningFromInput:
def test_keeps_summary_and_drops_encrypted_content_and_id(self):
wrapped = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id("gAAAAA-blob", "deployment-a")