mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(router): pin bridge-replayed encrypted reasoning to the deployment that minted it
The encrypted_content_affinity check only read the pin from the Responses input, which /v1/messages builds after the router has picked a deployment, so a model group spread across OpenAI orgs sent follow-up turns to the wrong org and got invalid_encrypted_content back. The check now also decodes the pin from bridge-tagged thinking and redacted_thinking blocks in the Anthropic messages. The bridge also keeps a deployment's own include list next to reasoning.encrypted_content instead of replacing it.
This commit is contained in:
parent
f49ebc93ea
commit
e27a018a6c
5 changed files with 104 additions and 10 deletions
|
|
@ -1873,7 +1873,7 @@ def _encrypted_reasoning_field(block: Mapping[str, object]) -> object:
|
|||
return None
|
||||
|
||||
|
||||
def _encrypted_content_of_block(block: Mapping[str, object]) -> str | None:
|
||||
def encrypted_content_of_block(block: Mapping[str, object]) -> str | None:
|
||||
return encrypted_content_from_signature(_encrypted_reasoning_field(block))
|
||||
|
||||
|
||||
|
|
@ -1900,7 +1900,7 @@ def _reasoning_item_from_block_group(group: tuple[Mapping[str, object], ...]) ->
|
|||
for block in group
|
||||
if (text := _readable_thinking_text(block))
|
||||
]
|
||||
encrypted_content: Final = _encrypted_content_of_block(group[0])
|
||||
encrypted_content: Final = encrypted_content_of_block(group[0])
|
||||
if encrypted_content is not None:
|
||||
return ChatCompletionReasoningItem(type="reasoning", summary=summary, encrypted_content=encrypted_content)
|
||||
if not summary:
|
||||
|
|
|
|||
|
|
@ -126,7 +126,7 @@ def _build_responses_kwargs(
|
|||
responses_kwargs["stream"] = True
|
||||
|
||||
# Forward litellm-specific kwargs (api_key, api_base, logging obj, etc.)
|
||||
excluded: Final = {"anthropic_messages"}
|
||||
excluded: Final = frozenset(("anthropic_messages",))
|
||||
for key, value in forwarded_kwargs.items():
|
||||
if key == "litellm_logging_obj" and value is not None:
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
|
|
@ -147,6 +147,14 @@ def _build_responses_kwargs(
|
|||
if explicit_prompt_cache_key is not None:
|
||||
responses_kwargs["prompt_cache_key"] = explicit_prompt_cache_key
|
||||
|
||||
deployment_include: Final = forwarded_kwargs.get("include")
|
||||
bridge_include: Final = responses_kwargs.get("include")
|
||||
if isinstance(deployment_include, list) and isinstance(bridge_include, list):
|
||||
responses_kwargs["include"] = [
|
||||
*bridge_include,
|
||||
*(item for item in deployment_include if item not in bridge_include),
|
||||
]
|
||||
|
||||
return responses_kwargs
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -37,6 +37,7 @@ Safe to enable globally:
|
|||
"""
|
||||
|
||||
import time
|
||||
from collections.abc import Iterator, Mapping
|
||||
from typing import TYPE_CHECKING, Final, Optional, Protocol, cast
|
||||
|
||||
import httpx
|
||||
|
|
@ -48,6 +49,7 @@ 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.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.router_utils.cooldown_cache import CooldownCacheValue
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
|
@ -138,15 +140,48 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
# If no encoded ID, check if encrypted_content itself is wrapped
|
||||
encrypted_content = item.get("encrypted_content")
|
||||
if encrypted_content and isinstance(encrypted_content, str):
|
||||
(
|
||||
model_id,
|
||||
_,
|
||||
) = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content)
|
||||
model_id = EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(encrypted_content)
|
||||
if model_id:
|
||||
return model_id
|
||||
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _anthropic_content_blocks(messages: object) -> Iterator[Mapping[str, object]]:
|
||||
if not isinstance(messages, list):
|
||||
return iter(())
|
||||
return (
|
||||
cast(Mapping[str, object], block) # cast-ok: narrowed by isinstance
|
||||
for message in cast(list[object], messages) # cast-ok: narrowed by isinstance
|
||||
if isinstance(message, Mapping)
|
||||
for content in (cast(Mapping[str, object], message).get("content"),) # cast-ok: narrowed by isinstance
|
||||
if isinstance(content, list)
|
||||
for block in cast(list[object], content) # cast-ok: narrowed by isinstance
|
||||
if isinstance(block, Mapping)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _model_id_from_wrapped_encrypted_content(encrypted_content: str) -> str | None:
|
||||
model_id, _ = ResponsesAPIRequestUtils._unwrap_encrypted_content_with_model_id(encrypted_content)
|
||||
return model_id or None
|
||||
|
||||
@staticmethod
|
||||
def _extract_model_id_from_anthropic_messages(messages: object) -> str | None:
|
||||
return next(
|
||||
(
|
||||
model_id
|
||||
for block in EncryptedContentAffinityCheck._anthropic_content_blocks(messages)
|
||||
if (encrypted_content := encrypted_content_of_block(block)) is not None
|
||||
if (
|
||||
model_id := EncryptedContentAffinityCheck._model_id_from_wrapped_encrypted_content(
|
||||
encrypted_content
|
||||
)
|
||||
)
|
||||
is not None
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None:
|
||||
for deployment in healthy_deployments:
|
||||
|
|
@ -248,13 +283,14 @@ class EncryptedContentAffinityCheck(CustomLogger):
|
|||
if "litellm_metadata" in request_kwargs:
|
||||
request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] = True
|
||||
|
||||
request_input: Final = request_kwargs.get("input")
|
||||
model_id: Final = self._extract_model_id_from_input(request_input)
|
||||
model_id: Final = self._extract_model_id_from_input(
|
||||
request_kwargs.get("input")
|
||||
) or self._extract_model_id_from_anthropic_messages(request_kwargs.get("messages"))
|
||||
if not model_id:
|
||||
return typed_healthy_deployments
|
||||
|
||||
verbose_router_logger.debug(
|
||||
"EncryptedContentAffinityCheck: decoded model_id=%s from input item IDs",
|
||||
"EncryptedContentAffinityCheck: decoded model_id=%s from the request's encrypted content markers",
|
||||
model_id,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -77,6 +77,16 @@ def test_build_responses_kwargs_skips_include_for_a_responses_provider_that_reje
|
|||
assert "reasoning" in responses_kwargs
|
||||
|
||||
|
||||
def test_build_responses_kwargs_keeps_the_deployment_include_next_to_encrypted_reasoning():
|
||||
responses_kwargs = _build_responses_kwargs(
|
||||
max_tokens=1024,
|
||||
messages=MESSAGES,
|
||||
model="openai/gpt-5.6-luna",
|
||||
extra_kwargs={"custom_llm_provider": "openai", "include": ["file_search_call.results"]},
|
||||
)
|
||||
assert responses_kwargs["include"] == ["reasoning.encrypted_content", "file_search_call.results"]
|
||||
|
||||
|
||||
def test_build_responses_kwargs_without_metadata_sets_no_prompt_cache_key():
|
||||
responses_kwargs = _build_responses_kwargs(
|
||||
max_tokens=1024,
|
||||
|
|
|
|||
|
|
@ -1656,3 +1656,43 @@ async def test_model_group_encrypted_content_affinity_overrides_global_deploymen
|
|||
assert request_kwargs.get("_encrypted_content_affinity_pinned") is True
|
||||
finally:
|
||||
router.discard()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_encrypted_content_affinity_pins_anthropic_messages_replayed_through_the_bridge():
|
||||
"""
|
||||
Claude Code behind /v1/messages replays the encrypted reasoning the bridge packed
|
||||
into a thinking block's signature (or a redacted block's data). The pin has to be
|
||||
read from those blocks because the bridge builds the Responses `input` only after
|
||||
the router has picked a deployment.
|
||||
"""
|
||||
check = EncryptedContentAffinityCheck()
|
||||
deployments = [
|
||||
{"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?"},
|
||||
],
|
||||
}
|
||||
|
||||
pinned = await check.async_filter_deployments(
|
||||
model="gpt-5.1",
|
||||
healthy_deployments=deployments,
|
||||
messages=None,
|
||||
request_kwargs=request_kwargs,
|
||||
)
|
||||
|
||||
assert [d["model_info"]["id"] for d in pinned] == ["openai-org-b"]
|
||||
assert request_kwargs["_encrypted_content_affinity_pinned"] is True
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue