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

Merge b52585e276 into b71491c6c9

Browse source
This commit is contained in:
Mrunmayee Daware 2026-10-30 13:29:39 +00:00 • committed by GitHub
commit 03b8e568fb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 139 additions and 64 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
litellm/responses/utils.py
View file

@ -574,13 +574,17 @@ class ResponsesAPIRequestUtils:
response: Union["ResponsesAPIResponse", dict[str, object]],
model_id: str | None,
) -> Union["ResponsesAPIResponse", dict[str, object]]:
"""Rewrite item IDs for output items that contain ``encrypted_content``.
"""Wrap ``encrypted_content`` of output items with the ``model_id``.
Encodes ``model_id`` into the item ID so that follow-up requests can be
routed back to the originating deployment without any cache lookup.
Follow-up requests are routed back to the originating deployment from
the ``litellm_enc:`` marker in the content - no cache lookup needed.
For items without an ID (e.g., from Codex), encodes model_id directly
into the encrypted_content itself.
Item IDs are deliberately left as the provider issued them: the streamed
``output_item.added``/``.done`` events cannot rewrite IDs (at ``.added``
time the item may not carry ``encrypted_content`` yet), so rewriting
only here gave one item two IDs, and ID-keyed clients replayed both
copies (#47881). Clients holding older ``encitem_...`` IDs keep working
through ``restore_encrypted_content_item_ids_in_input``.
"""
if not model_id:
return response
@ -596,35 +600,26 @@ class ResponsesAPIRequestUtils:
for item in output:
if isinstance(item, dict):
item_id = item.get("id")
encrypted_content = item.get("encrypted_content")
if encrypted_content and isinstance(encrypted_content, str):
# Always wrap encrypted_content with model_id for redundancy
# Wrap encrypted_content with the model_id; the item ID stays
# as the provider issued it so the item keeps one ID across
# the streamed output_item events and the final response.
item["encrypted_content"] = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id(
encrypted_content, model_id
)
# Also encode the ID if present
if item_id and isinstance(item_id, str):
item["id"] = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, item_id)
else:
item_id = getattr(item, "id", None)
encrypted_content = getattr(item, "encrypted_content", None)
if encrypted_content and isinstance(encrypted_content, str):
# Always wrap encrypted_content with model_id for redundancy
# Same as the dict branch: wrap the content, keep the ID.
try:
item.encrypted_content = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id(
encrypted_content, model_id
)
except AttributeError:
pass
# Also encode the ID if present
if item_id and isinstance(item_id, str):
try:
item.id = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, item_id)
except AttributeError:
pass
return response
@ -638,10 +633,17 @@ class ResponsesAPIRequestUtils:
Handles both:
1. Items with encoded IDs (encitem_...)
2. Items with wrapped encrypted_content (litellm_enc:...)
Also drops later duplicates of the same restored item ID: clients that
collected an item from both the streamed per-item events and
``response.completed`` while those carried different IDs replay both
copies, and providers reject duplicate item IDs with a 400 (#47881).
"""
if not isinstance(request_input, list):
return request_input
seen_ids: Final = set()
restored_input: Final = []
for item in request_input:
if isinstance(item, dict):
item_id = item.get("id")
@ -649,6 +651,7 @@ class ResponsesAPIRequestUtils:
decoded = ResponsesAPIRequestUtils.decode_encrypted_item_id(item_id)
if decoded:
item["id"] = decoded["item_id"]
item_id = item["id"]
encrypted_content = item.get("encrypted_content")
if encrypted_content and isinstance(encrypted_content, str):
@ -659,7 +662,14 @@ class ResponsesAPIRequestUtils:
if unwrapped != encrypted_content:
item["encrypted_content"] = unwrapped
return request_input
if item_id and isinstance(item_id, str):
if item_id in seen_ids:
# A second copy of an item the client collected twice
continue
seen_ids.add(item_id)
restored_input.append(item)
return cast(_RequestInputT, restored_input) # cast-ok: list inputs retain their input type after deduplication
_restore_encrypted_content_item_ids_in_input = restore_encrypted_content_item_ids_in_input

5
tests/unit/responses/test_responses_websocket_all_providers.py
View file

@ -4014,9 +4014,8 @@ class TestNativeWebSocketEncryptedContentAffinity:
assert completed["response"]["id"] == ResponsesAPIRequestUtils._build_responses_api_response_id(
custom_llm_provider="openai", model_id="dep-1", response_id="resp_1"
)
assert completed["response"]["output"][0]["id"] == ResponsesAPIRequestUtils._build_encrypted_item_id(
"dep-1", "rs_1"
)
# The provider's item ID is preserved; only the content is wrapped (#47881)
assert completed["response"]["output"][0]["id"] == "rs_1"
assert completed["response"]["output"][0]["encrypted_content"] == wrapped_content
await asyncio.sleep(0)
logged = logging_obj.dispatch_success_handlers.await_args[0][0]

150
tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py
View file

@ -75,13 +75,17 @@ def _get_item_id(item) -> str:
return getattr(item, "id", "") or ""
def _extract_encoded_item_id(response) -> str:
"""Return the first ``encitem_``-prefixed item ID from the response output."""
def _extract_encrypted_item(response) -> dict:
"""Return the first output item carrying ``encrypted_content`` as a replayable dict.
The provider's item ID is preserved on responses (#47881); the ``litellm_enc:``
marker in the content carries the deployment for affinity routing.
"""
for item in response.output or []:
item_id = _get_item_id(item)
if item_id.startswith("encitem_"):
return item_id
return ""
content = item.get("encrypted_content") if isinstance(item, dict) else getattr(item, "encrypted_content", None)
if content:
return {"id": _get_item_id(item), "encrypted_content": content}
return {}
# ---------------------------------------------------------------------------
@ -128,7 +132,7 @@ class TestEncryptedItemIdCodec:
class TestUpdateEncryptedContentItemIds:
def test_rewrites_encrypted_items_in_dict_response(self):
def test_preserves_item_id_and_wraps_content_in_dict_response(self):
model_id = "deployment-1"
response = {
"id": "resp_123",
@ -140,12 +144,12 @@ class TestUpdateEncryptedContentItemIds:
result = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response(response, model_id)
# Plain message item untouched
assert result["output"][0]["id"] == "msg_abc"
# Reasoning item with encrypted_content gets encoded
encoded_id = result["output"][1]["id"]
assert encoded_id.startswith("encitem_")
decoded = ResponsesAPIRequestUtils.decode_encrypted_item_id(encoded_id)
assert decoded["model_id"] == model_id
assert decoded["item_id"] == "rs_xyz"
# Reasoning item keeps the provider's ID; only the content is wrapped (#47881)
assert result["output"][1]["id"] == "rs_xyz"
wrapped = result["output"][1]["encrypted_content"]
assert wrapped.startswith("litellm_enc:")
unwrapped_model_id, _ = ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id(wrapped)
assert unwrapped_model_id == model_id
def test_no_op_when_model_id_is_none(self):
response = {"output": [{"id": "rs_xyz", "type": "reasoning", "encrypted_content": "secret"}]}
@ -153,6 +157,62 @@ class TestUpdateEncryptedContentItemIds:
assert result["output"][0]["id"] == "rs_xyz"
class TestSingleIdAcrossStreamedAndCompletedViews:
"""#47881: one item keeps one ID across every event that carries it."""
def test_response_path_keeps_the_streamed_events_id(self):
model_id = "deployment-1"
response = {
"id": "resp_123",
"output": [{"id": "rs_streamed", "type": "reasoning", "encrypted_content": "blob"}],
}
completed = ResponsesAPIRequestUtils._update_encrypted_content_item_ids_in_response(
{"id": "resp_123", "output": [dict(response["output"][0])]}, model_id
)
# What the streaming iterator emits for output_item.added/.done: wrapped content, raw ID
streamed = {
"id": "rs_streamed",
"type": "reasoning",
"encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("blob", model_id),
}
assert completed["output"][0]["id"] == streamed["id"]
assert completed["output"][0]["encrypted_content"] == streamed["encrypted_content"]
class TestRestoreDropsDuplicateReplayedCopies:
"""#47881: clients that stored both the streamed and the completed copy replay one item."""
def test_dedupes_two_copies_of_one_item(self):
model_id = "deployment-1"
encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id(model_id, "rs_dup")
wrapped = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("blob", model_id)
request_input = [
{"role": "user", "content": "hi"},
{"type": "reasoning", "id": "rs_dup", "encrypted_content": wrapped},
{"type": "reasoning", "id": encoded_id, "encrypted_content": wrapped},
{"role": "user", "content": "again"},
]
result = ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(request_input)
assert result is not request_input
assert len(request_input) == 4
reasoning = [i for i in result if isinstance(i, dict) and i.get("type") == "reasoning"]
assert len(reasoning) == 1
assert reasoning[0]["id"] == "rs_dup"
assert reasoning[0]["encrypted_content"] == "blob"
assert len(result) == 3
def test_keeps_distinct_items(self):
model_id = "deployment-1"
wrapped = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id("blob", model_id)
request_input = [
{"type": "reasoning", "id": "rs_one", "encrypted_content": wrapped},
{"type": "reasoning", "id": "rs_two", "encrypted_content": wrapped},
{"type": "message", "id": "msg_one", "role": "assistant"},
]
result = ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(request_input)
assert [i["id"] for i in result] == ["rs_one", "rs_two", "msg_one"]
class TestEncryptedContentWrapping:
def test_wrap_and_unwrap_encrypted_content(self):
"""Test wrapping encrypted_content with model_id metadata."""
@ -249,13 +309,13 @@ class TestRestoreEncryptedContentItemIds:
@pytest.mark.asyncio
async def test_encrypted_content_affinity_tracks_and_routes():
"""
The first response rewrites encrypted-content item IDs to encoded form.
The follow-up request with those encoded IDs is pinned to the same deployment.
The first response wraps encrypted content with the deployment's model_id.
The follow-up request replaying that item is pinned to the same deployment.
Mocks ``async_response_api_handler`` (the method that makes the HTTP call)
so the test is deterministic regardless of the HTTP transport in use.
The ``@client`` decorator and ``_update_responses_api_response_id_with_model_id``
post-processing still run, so item-ID rewriting is exercised end-to-end.
post-processing still run, so content wrapping is exercised end-to-end.
"""
mock_resp = _build_mock_response(
output_items=[
@ -324,27 +384,28 @@ async def test_encrypted_content_affinity_tracks_and_routes():
first_model_id = first_response._hidden_params["model_id"]
selected_deployments.append(first_model_id)
# The response must have rewritten the encrypted item's ID to encoded form
encoded_item_id = _extract_encoded_item_id(first_response)
assert encoded_item_id.startswith("encitem_"), (
f"Expected output item ID to be rewritten to encitem_... but got {encoded_item_id!r}"
# The response keeps the provider's item ID and wraps the encrypted content
replayed_item = _extract_encrypted_item(first_response)
assert replayed_item["id"] == "rs_encrypted_item_456", (
f"Expected the provider's item ID to be preserved, but got {replayed_item['id']!r}"
)
assert replayed_item["encrypted_content"].startswith("litellm_enc:")
# Verify the encoded ID decodes back to the correct deployment + original ID
decoded = ResponsesAPIRequestUtils.decode_encrypted_item_id(encoded_item_id)
assert decoded is not None
assert decoded["model_id"] == first_model_id
assert decoded["item_id"] == "rs_encrypted_item_456"
# The wrapped content carries the correct deployment for affinity
decoded_model_id, _ = ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id(
replayed_item["encrypted_content"]
)
assert decoded_model_id == first_model_id
# Second request: use the encoded item IDs from the first response
# Second request: replay the item as an ID-keyed client stored it
second_response = await router.aresponses(
model="openai.gpt-5.1-codex",
input=[
{"type": "message", "id": "msg_abc123", "role": "assistant"},
{
"type": "reasoning",
"id": encoded_item_id,
"encrypted_content": "gAAAAABpnW_yEYmSNEyOG...",
"id": replayed_item["id"],
"encrypted_content": replayed_item["encrypted_content"],
},
],
)
@ -455,18 +516,21 @@ async def test_encrypted_content_affinity_bypasses_rpm_limits():
first_model_id = first_response._hidden_params["model_id"]
selected_deployments.append(first_model_id)
# Extract encoded item ID from the first response output
encoded_item_id = _extract_encoded_item_id(first_response)
assert encoded_item_id.startswith("encitem_"), f"Expected encitem_... but got {encoded_item_id!r}"
# Extract the replayable item from the first response output
replayed_item = _extract_encrypted_item(first_response)
assert replayed_item["id"] == "rs_encrypted_must_pin", (
f"Expected the provider's item ID, got {replayed_item['id']!r}"
)
assert replayed_item["encrypted_content"].startswith("litellm_enc:")
# Follow-up with the encoded item ID — should pin to same deployment
# Follow-up replaying the wrapped item — should pin to same deployment
second_response = await router.aresponses(
model="openai.gpt-5.1-codex",
input=[
{
"type": "reasoning",
"id": encoded_item_id,
"encrypted_content": "gAAAAABpnW_yEYmSNEyOG...",
"id": replayed_item["id"],
"encrypted_content": replayed_item["encrypted_content"],
},
],
)
@ -864,8 +928,9 @@ async def test_affinity_falls_back_to_same_encryption_boundary_on_model_group_sw
r1 = await router.aresponses(model="gpt-5.3-codex", input="hi")
assert r1._hidden_params["model_id"] == "gpt-5.3-codex-account-a"
encoded_id = _extract_encoded_item_id(r1)
assert encoded_id.startswith("encitem_")
replayed_item = _extract_encrypted_item(r1)
assert replayed_item["id"] == "rs_encrypted_xyz"
assert replayed_item["encrypted_content"].startswith("litellm_enc:")
# simple_shuffle.random.choice NOT patched: prove affinity narrows the
# candidate pool to a single deployment regardless of which one shuffle picks.
@ -879,8 +944,8 @@ async def test_affinity_falls_back_to_same_encryption_boundary_on_model_group_sw
input=[
{
"type": "reasoning",
"id": encoded_id,
"encrypted_content": "gAAAAABpnW_yEYmSNEyOG...",
"id": replayed_item["id"],
"encrypted_content": replayed_item["encrypted_content"],
},
],
)
@ -991,8 +1056,9 @@ async def test_affinity_falls_back_to_same_boundary_on_alias_switch():
):
r1 = await router.aresponses(model="gpt-5.3-codex", input="hi")
encoded_id = _extract_encoded_item_id(r1)
assert encoded_id.startswith("encitem_")
replayed_item = _extract_encrypted_item(r1)
assert replayed_item["id"] == "rs_alias_xyz"
assert replayed_item["encrypted_content"].startswith("litellm_enc:")
with patch(
"litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler",
@ -1004,8 +1070,8 @@ async def test_affinity_falls_back_to_same_boundary_on_alias_switch():
input=[
{
"type": "reasoning",
"id": encoded_id,
"encrypted_content": "gAAAAABpnW_yEYmSNEyOG...",
"id": replayed_item["id"],
"encrypted_content": replayed_item["encrypted_content"],
},
],
)