From 93115527953a0b61f7a54ccd55711e2565cfcec9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:11:44 -0700 Subject: [PATCH 01/21] fix(responses): round-trip Anthropic compaction blocks as compaction items /v1/responses on Anthropic routes now returns Claude's compaction block as a compaction output item (and streams it as output_item.added/done), and a compaction input item is replayed to Claude as its native compaction block, so the next turn reuses the summary instead of compacting again. Other routes receive the summary as assistant text. The Anthropic stream handler now keeps the summary text that arrives in compaction_delta. Fixes #41456 --- litellm/llms/anthropic/chat/handler.py | 14 ++- litellm/proxy/_lazy_openapi_snapshot.json | 3 + .../compaction_items.py | 60 ++++++++++ .../streaming_iterator.py | 49 +++++++- .../transformation.py | 49 ++++++++ litellm/types/llms/openai.py | 2 + .../test_anthropic_responses_bridge.py | 108 ++++++++++++++++++ .../test_litellm_completion_responses.py | 35 ++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 9 files changed, 316 insertions(+), 6 deletions(-) create mode 100644 litellm/responses/litellm_completion_transformation/compaction_items.py diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index a68ca1f5534..cf84f11b4c2 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -851,11 +851,21 @@ class ModelResponseIterator: if reasoning_content is None: reasoning_content = "" elif "content" in content_block["delta"] and content_block["delta"].get("type") == "compaction_delta": - # Handle compaction delta + compaction_delta: Final = content_block["delta"]["content"] provider_specific_fields["compaction_delta"] = { "type": "compaction_delta", - "content": content_block["delta"]["content"], + "content": compaction_delta, } + if self.compaction_blocks: + streamed_summary: Final = self.compaction_blocks[-1].get("content") + self.compaction_blocks = [ + *self.compaction_blocks[:-1], + { + **self.compaction_blocks[-1], + "content": (streamed_summary if isinstance(streamed_summary, str) else "") + compaction_delta, + }, + ] + provider_specific_fields["compaction_blocks"] = self.compaction_blocks return ( text, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 9cf0b813836..da788a04643 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -27149,6 +27149,9 @@ }, { "$ref": "#/components/schemas/OutputToolSearchCall" + }, + { + "$ref": "#/components/schemas/ResponseCompactionItem" } ] }, diff --git a/litellm/responses/litellm_completion_transformation/compaction_items.py b/litellm/responses/litellm_completion_transformation/compaction_items.py new file mode 100644 index 00000000000..d9f4c95a060 --- /dev/null +++ b/litellm/responses/litellm_completion_transformation/compaction_items.py @@ -0,0 +1,60 @@ +import json +import uuid +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Final + +from openai.types.responses.response_compaction_item import ResponseCompactionItem +from pydantic import TypeAdapter, ValidationError + +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) +_OBJECT_LIST: Final = TypeAdapter(list[object]) + + +@dataclass(frozen=True, slots=True) +class CompactionBlock: + summary: str + wire_block: Mapping[str, object] + + +def _compaction_block(value: object) -> CompactionBlock | None: + try: + block: Final = _JSON_OBJECT.validate_python(value) + except ValidationError: + return None + summary: Final = block.get("content") + if block.get("type") != "compaction" or not isinstance(summary, str) or not summary: + return None + return CompactionBlock(summary=summary, wire_block=block) + + +def _newest_compaction_block(compaction_blocks: object) -> CompactionBlock | None: + if not isinstance(compaction_blocks, list): + return None + blocks: Final = _OBJECT_LIST.validate_python(compaction_blocks) + return next((block for value in reversed(blocks) if (block := _compaction_block(value)) is not None), None) + + +def compaction_output_item(provider_specific_fields: object) -> ResponseCompactionItem | None: + try: + fields: Final = _JSON_OBJECT.validate_python(provider_specific_fields) + except ValidationError: + return None + block: Final = _newest_compaction_block(fields.get("compaction_blocks")) + if block is None: + return None + return ResponseCompactionItem( + id=f"cmp_{uuid.uuid4()}", + encrypted_content=json.dumps(dict(block.wire_block), separators=(",", ":")), + type="compaction", + ) + + +def decode_compaction_item(encrypted_content: object) -> CompactionBlock | None: + if not isinstance(encrypted_content, str): + return None + try: + decoded: Final = _JSON_OBJECT.validate_json(encrypted_content) + except ValidationError: + return None + return _compaction_block(decoded) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 418bdb3d6db..0b3fb8ef25f 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -7,6 +7,7 @@ from typing import Any, Final, cast import litellm from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params from litellm.main import stream_chunk_builder +from litellm.responses.litellm_completion_transformation.compaction_items import compaction_output_item from litellm.responses.litellm_completion_transformation.custom_tools import ( build_tool_call_item_kwargs, extract_custom_tool_names, @@ -82,6 +83,11 @@ def _output_items_with_id(items: tuple[Any, ...], item_type: str, item_id: str | ) +def _streams_compaction(chunk: ModelResponseStream) -> bool: + delta_fields: Final = chunk.choices[0].delta.provider_specific_fields if chunk.choices else None + return delta_fields is not None and "compaction_blocks" in delta_fields + + def _delta_has_signed_thinking_block(delta: object) -> bool: blocks: Final = getattr(delta, "thinking_blocks", None) or () return any(isinstance(b, dict) and (b.get("signature") or b.get("data")) for b in blocks) @@ -140,6 +146,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._cached_item_id: str | None = None self._message_output_index: int | None = None self._reasoning_output_index: int | None = None + self._compaction_output_index: int | None = None + self._compaction_item_id: str | None = None self._cached_response_id: str | None = None self._buffered_chunk: ModelResponseStream | None = None self._upstream_exhausted: bool = False @@ -180,7 +188,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._tool_search_call_ids = self._tool_search_call_ids | {call_id} return True - def _queue_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> None: + def _next_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> OutputItemAddedEvent: self._sequence_number += 1 added: Final = OutputItemAddedEvent( type=ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, @@ -188,7 +196,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): item=item, ) added.__dict__["sequence_number"] = self._sequence_number - self._pending_tool_events.append(added) + return added + + def _queue_output_item_added(self, output_index: int, item: BaseLiteLLMOpenAIResponseObject) -> None: + self._pending_tool_events.append(self._next_output_item_added(output_index, item)) def _queue_tool_search_call_added(self, call_id: str, output_index: int, args_delta: str) -> None: if call_id not in self._tool_args_by_call_id: @@ -242,6 +253,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): return self._message_output_index case "reasoning": return self._reasoning_output_index + case "compaction": + return self._compaction_output_index case _: call_id: Final = getattr(item, "call_id", None) or str(getattr(item, "id", "")).removeprefix("ws_") return self._tool_output_index_by_call_id.get(str(call_id)) @@ -250,6 +263,28 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): index: Final = self._streamed_output_index(item) return (index is None, index or 0) + def _queue_compaction_item_events(self) -> None: + if ( + self._compaction_output_index is not None + or "compaction_blocks" not in self._accumulated_provider_specific_fields + ): + return + item: Final = compaction_output_item(self._accumulated_provider_specific_fields) + if item is None: + return + self._compaction_item_id = item.id + self._compaction_output_index = self._allocate_output_index() + streamed_item: Final = BaseLiteLLMOpenAIResponseObject(**item.model_dump(exclude_none=True)) + added: Final = self._next_output_item_added(self._compaction_output_index, streamed_item) + self._sequence_number += 1 + done: Final = OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=self._compaction_output_index, + sequence_number=self._sequence_number, + item=streamed_item, + ) + self._pending_response_events.extend((added, done)) + def _normalize_tool_call_index(self, tool_call: object) -> int | None: idx_raw: Final = tool_call.get("index") if isinstance(tool_call, dict) else getattr(tool_call, "index", None) if idx_raw is None: @@ -978,6 +1013,9 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if not self.litellm_model_response or isinstance(self.litellm_model_response, TextCompletionResponse): self.litellm_model_response = self.create_litellm_model_response() if self.litellm_model_response: + self._queue_compaction_item_events() + if self._pending_response_events: + return self._pending_response_events.pop(0) # If tool calls exist, emit tool events before finishing/response.completed. if isinstance(self.litellm_model_response, ModelResponse): self._queue_final_tool_call_done_events(self.litellm_model_response) @@ -1079,6 +1117,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): chunk = await self.litellm_custom_stream_wrapper.__anext__() if chunk is not None: chunk = cast(ModelResponseStream, chunk) + if not _streams_compaction(chunk): + self._queue_compaction_item_events() for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -1187,6 +1227,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): raise StopIteration else: chunk = self.litellm_custom_stream_wrapper.__next__() + if not _streams_compaction(chunk): + self._queue_compaction_item_events() for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -1342,7 +1384,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): "reasoning", self._cached_reasoning_item_id, ) - streamed_items: Final = filterfalse(self._is_unstreamed_empty_message, reasoning_aligned) + compaction_aligned: Final = _output_items_with_id(reasoning_aligned, "compaction", self._compaction_item_id) + streamed_items: Final = filterfalse(self._is_unstreamed_empty_message, compaction_aligned) return tuple(sorted(streamed_items, key=self._streamed_output_position)) def _is_unstreamed_empty_message(self, item: object) -> bool: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 7f5d143e98e..d922f67fc77 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -29,6 +29,7 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( Function as NamedToolChoiceFunction, ) from openai.types.responses import ResponseFunctionToolCall, ResponseFunctionWebSearch +from openai.types.responses.response_compaction_item import ResponseCompactionItem from openai.types.responses.response_create_params import ResponseInputParam from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam @@ -115,6 +116,7 @@ from litellm.types.utils import ( Usage, ) +from .compaction_items import compaction_output_item, decode_compaction_item from .custom_tools import ( convert_custom_tool_to_function_tool, extract_custom_tool_names, @@ -339,6 +341,10 @@ def without_item_cache_control(input: str | ResponseInputParam) -> str | Respons ) +def _carries_compaction_item(input: str | ResponseInputParam) -> bool: + return isinstance(input, list) and any(item.get("type") == "compaction" for item in input) + + _PROVIDERS_KEEPING_MESSAGE_CACHE_CONTROL: Final = frozenset(("anthropic", "dashscope", "minimax", "sap", "zai")) @@ -866,6 +872,8 @@ class LiteLLMCompletionResponsesConfig: responses_api_request=responses_api_request, replay_reasoning=True, item_cache_control=keeps_item_marks, + native_compaction=_carries_compaction_item(input) + and provider_builds_anthropic_request(model, custom_llm_provider), ) litellm_completion_request: dict = { @@ -927,6 +935,7 @@ class LiteLLMCompletionResponsesConfig: responses_api_request: ResponsesAPIOptionalRequestParams | dict, replay_reasoning: bool = False, item_cache_control: bool = False, + native_compaction: bool = False, ) -> list[ AllMessageValues | GenericChatCompletionMessage @@ -948,6 +957,11 @@ class LiteLLMCompletionResponsesConfig: guardrail scanning) leave both off, because they need every piece of text in the request to stay readable as message ``content``, with no item mark a strict message schema could refuse. + + ``native_compaction`` belongs to callers whose target builds an + Anthropic request: a ``compaction`` item then goes back as the + provider's own compaction block. Everyone else reads its summary as + assistant text. """ messages: list[ AllMessageValues @@ -967,6 +981,7 @@ class LiteLLMCompletionResponsesConfig: LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message( input=input if item_cache_control else without_item_cache_control(input), replay_reasoning=replay_reasoning, + native_compaction=native_compaction, ) ) @@ -1097,6 +1112,7 @@ class LiteLLMCompletionResponsesConfig: def _transform_response_input_param_to_chat_completion_message( input: str | ResponseInputParam, replay_reasoning: bool = False, + native_compaction: bool = False, ) -> list[ AllMessageValues | GenericChatCompletionMessage | ChatCompletionMessageToolCall | ChatCompletionResponseMessage ]: @@ -1122,6 +1138,7 @@ class LiteLLMCompletionResponsesConfig: LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( input_item=_input, replay_reasoning=replay_reasoning, + native_compaction=native_compaction, ) ) @@ -1821,6 +1838,7 @@ class LiteLLMCompletionResponsesConfig: def _transform_responses_api_input_item_to_chat_completion_message( input_item: Mapping[str, object], replay_reasoning: bool = False, + native_compaction: bool = False, ) -> list[AllMessageValues | GenericChatCompletionMessage | ChatCompletionResponseMessage]: """ Transform a Responses API input item into a Chat Completion message @@ -1859,6 +1877,12 @@ class LiteLLMCompletionResponsesConfig: Mapping[str, str], input_item ) ) + elif input_item.get("type") == "compaction": + return list( + LiteLLMCompletionResponsesConfig._transform_compaction_item_to_chat_completion_messages( + input_item, native_compaction + ) + ) elif input_item.get("type") == "reasoning": # A ResponseReasoningItemParam carries the prior-turn chain-of-thought. # Chat-completions providers (DeepSeek V4, Kimi K2.6, ...) expect this @@ -1910,6 +1934,23 @@ class LiteLLMCompletionResponsesConfig: return [] return [LiteLLMCompletionResponsesConfig._transform_responses_api_message_item(input_item, content)] + @staticmethod + def _transform_compaction_item_to_chat_completion_messages( + input_item: Mapping[str, object], native_compaction: bool + ) -> tuple[GenericChatCompletionMessage | ChatCompletionResponseMessage, ...]: + compaction: Final = decode_compaction_item(input_item.get("encrypted_content")) + if compaction is None: + return () + if native_compaction: + return ( + ChatCompletionResponseMessage( + role="assistant", + content=None, + provider_specific_fields={"compaction_blocks": [dict(compaction.wire_block)]}, + ), + ) + return (GenericChatCompletionMessage(role="assistant", content=compaction.summary),) + @staticmethod def _transform_responses_api_message_item( input_item: Mapping[str, object], content: object @@ -3020,6 +3061,7 @@ class LiteLLMCompletionResponsesConfig: | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | ResponseCompactionItem ]: responses_output: list[ GenericResponseOutputItem @@ -3030,8 +3072,10 @@ class LiteLLMCompletionResponsesConfig: | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | ResponseCompactionItem ] = [] + responses_output.extend(LiteLLMCompletionResponsesConfig._extract_compaction_output_items(choices)) responses_output.extend( LiteLLMCompletionResponsesConfig._extract_reasoning_output_items(chat_completion_response, choices) ) @@ -3097,6 +3141,11 @@ class LiteLLMCompletionResponsesConfig: output_items.append(item) return output_items + @staticmethod + def _extract_compaction_output_items(choices: Sequence[Choices]) -> tuple[ResponseCompactionItem, ...]: + item: Final = compaction_output_item(choices[0].message.provider_specific_fields) if choices else None + return () if item is None else (item,) + @staticmethod def _encode_thinking_blocks(message: Message) -> str | None: thinking_blocks: Final[Sequence[Mapping[str, object]]] = getattr(message, "thinking_blocks", None) or () diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 1e4f78b2204..abcc2feae5e 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -55,6 +55,7 @@ except (ImportError, AttributeError): from typing import Annotated +from openai.types.responses.response_compaction_item import ResponseCompactionItem from openai.types.responses.response_create_params import ( Reasoning, ResponseIncludable, @@ -1465,6 +1466,7 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | ResponseCompactionItem ] ) parallel_tool_calls: bool | None = None diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 61b23428d62..efdbe84babb 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -540,3 +540,111 @@ def test_streamed_responses_over_anthropic_emit_the_responses_event_sequence(res assert deltas == "Argentina won." assert completed.response.status == "completed" assert completed.response.output[0].content[0].text == "Argentina won." + + +_ANTHROPIC_MESSAGES_URL: Final = "https://api.anthropic.com/v1/messages" +_COMPACTION_BLOCK: Final = {"type": "compaction", "content": "The user pasted a long document and asked about it."} +_FIRST_INPUT: Final = ({"role": "user", "content": "What is the first word of the document?"},) +_NEXT_INPUT: Final = ({"role": "user", "content": "And the last word?"},) + + +def _anthropic_message(*content: Mapping[str, object]) -> httpx.Response: + return httpx.Response( + 200, + json={ + "id": "msg_compacted", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": list(content), + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + }, + ) + + +def _replay(output: Sequence[BaseModel]) -> list[object]: + return [*_FIRST_INPUT, *(item.model_dump(exclude_none=True) for item in output), *_NEXT_INPUT] + + +def _sent_assistant_content(route: respx.Route) -> list[object]: + sent: Final = json.loads(route.calls.last.request.content) + return next(message["content"] for message in sent["messages"] if message["role"] == "assistant") + + +def test_responses_over_anthropic_return_the_compaction_block_and_send_it_back( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + side_effect=[ + _anthropic_message(_COMPACTION_BLOCK, {"type": "text", "text": "The."}), + _anthropic_message({"type": "text", "text": "dog."}), + ] + ) + + first: Final = litellm.responses(model=_MODEL, input=list(_FIRST_INPUT), api_key="mock_api_key") + litellm.responses(model=_MODEL, input=_replay(first.output), api_key="mock_api_key") + + assert [item.type for item in first.output] == ["compaction", "message"] + assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] + + +def test_streamed_responses_over_anthropic_emit_the_compaction_item_and_send_it_back( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + side_effect=[ + httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_compacted_stream", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "compaction", "content": None}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "compaction_delta", "content": _COMPACTION_BLOCK["content"]}, + }, + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "text_delta", "text": "The."}}, + {"type": "content_block_stop", "index": 1}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ), + _anthropic_message({"type": "text", "text": "dog."}), + ] + ) + + events: Final = tuple(litellm.responses(model=_MODEL, input=list(_FIRST_INPUT), stream=True, api_key="mock_api_key")) + item_events: Final = tuple( + event + for event in events + if event.type in (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE) + ) + completed: Final = events[-1].response + litellm.responses(model=_MODEL, input=_replay(completed.output), api_key="mock_api_key") + + assert [(event.type, event.output_index, event.item.type) for event in item_events] == [ + (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, 0, "compaction"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, 0, "compaction"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, 1, "message"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, 1, "message"), + ] + assert [(item.type, getattr(item, "encrypted_content", None)) for item in completed.output[:1]] == [ + ("compaction", item_events[1].item.encrypted_content) + ] + assert completed.output[0].id == item_events[1].item.id + assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 0d55fc35e1d..98c045bcd6e 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -6464,3 +6464,38 @@ def test_transform_request_reads_service_tier_from_responses_api_request() -> No ) assert result["service_tier"] == "priority" + + +_COMPACTION_BLOCK: Final = {"type": "compaction", "content": "Earlier turns covered the Tokyo trip."} + + +@pytest.mark.parametrize( + ("model", "replayed_compaction"), + [ + ( + "anthropic/claude-sonnet-5-5", + {"role": "assistant", "content": None, "provider_specific_fields": {"compaction_blocks": [_COMPACTION_BLOCK]}}, + ), + ("gemini/gemini-3.8-flash", {"role": "assistant", "content": _COMPACTION_BLOCK["content"]}), + ], +) +def test_compaction_input_item_replays_as_a_compaction_block_only_on_anthropic_routes( + model: str, replayed_compaction: Mapping[str, object] +) -> None: + result: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model=model, + input=cast( + ResponseInputParam, + [ + {"type": "compaction", "id": "cmp_openai", "encrypted_content": "gAAAAABopaque"}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"role": "user", "content": "And Osaka?"}, + ], + ), + responses_api_request={}, + ) + + assert [dict(message) for message in result["messages"]] == [ + replayed_compaction, + {"role": "user", "content": "And Osaka?"}, + ] diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 95461192a8c..c51e3d1c50b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -43958,7 +43958,7 @@ export interface components { /** Output */ output: (components["schemas"]["ResponseOutputMessage"] | components["schemas"]["ResponseFileSearchToolCall"] | components["schemas"]["ResponseFunctionToolCall"] | components["schemas"]["ResponseFunctionToolCallOutputItem"] | components["schemas"]["ResponseFunctionWebSearch"] | components["schemas"]["ResponseComputerToolCall"] | components["schemas"]["ResponseComputerToolCallOutputItem"] | components["schemas"]["ResponseReasoningItem"] | components["schemas"]["ResponseToolSearchCall"] | components["schemas"]["ResponseToolSearchOutputItem"] | components["schemas"]["ResponseCompactionItem"] | components["schemas"]["ImageGenerationCall"] | components["schemas"]["ResponseCodeInterpreterToolCall"] | components["schemas"]["LocalShellCall"] | components["schemas"]["LocalShellCallOutput"] | components["schemas"]["ResponseFunctionShellToolCall"] | components["schemas"]["ResponseFunctionShellToolCallOutput"] | components["schemas"]["ResponseApplyPatchToolCall"] | components["schemas"]["ResponseApplyPatchToolCallOutput"] | components["schemas"]["McpCall"] | components["schemas"]["McpListTools"] | components["schemas"]["McpApprovalRequest"] | components["schemas"]["McpApprovalResponse"] | components["schemas"]["ResponseCustomToolCall"] | components["schemas"]["ResponseCustomToolCallOutputItem"] | { [key: string]: unknown; - })[] | (components["schemas"]["GenericResponseOutputItem"] | components["schemas"]["OutputCodeInterpreterCall"] | components["schemas"]["OutputFunctionToolCall"] | components["schemas"]["OutputImageGenerationCall"] | components["schemas"]["ResponseFunctionToolCall"] | components["schemas"]["ResponseFunctionWebSearch"] | components["schemas"]["CustomToolCallOutputItem"] | components["schemas"]["OutputToolSearchCall"])[]; + })[] | (components["schemas"]["GenericResponseOutputItem"] | components["schemas"]["OutputCodeInterpreterCall"] | components["schemas"]["OutputFunctionToolCall"] | components["schemas"]["OutputImageGenerationCall"] | components["schemas"]["ResponseFunctionToolCall"] | components["schemas"]["ResponseFunctionWebSearch"] | components["schemas"]["CustomToolCallOutputItem"] | components["schemas"]["OutputToolSearchCall"] | components["schemas"]["ResponseCompactionItem"])[]; /** Parallel Tool Calls */ parallel_tool_calls?: boolean | null; /** Previous Response Id */ From 5588fe3454da18e32379e51dbbdd8e6bfab03434 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:26:18 -0700 Subject: [PATCH 02/21] fix(responses): replay compaction blocks natively only when the request compacts --- .../transformation.py | 23 +++++++---- .../test_anthropic_responses_bridge.py | 40 +++++++++++++++++-- .../test_litellm_completion_responses.py | 18 ++++++--- 3 files changed, 65 insertions(+), 16 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index d922f67fc77..eab8fe92b93 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -79,6 +79,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolParam, ChatCompletionToolParamFunctionChunk, ChatCompletionUserMessage, + ContextManagementEntry, GenericChatCompletionMessage, IncompleteDetails, InputTokensDetails, @@ -341,8 +342,14 @@ def without_item_cache_control(input: str | ResponseInputParam) -> str | Respons ) -def _carries_compaction_item(input: str | ResponseInputParam) -> bool: - return isinstance(input, list) and any(item.get("type") == "compaction" for item in input) +def _replays_compaction_block( + input: str | ResponseInputParam, context_management: Sequence[ContextManagementEntry] | None +) -> bool: + return ( + isinstance(input, list) + and any(item.get("type") == "compaction" for item in input) + and any(entry.get("type") == "compaction" for entry in context_management or ()) + ) _PROVIDERS_KEEPING_MESSAGE_CACHE_CONTROL: Final = frozenset(("anthropic", "dashscope", "minimax", "sap", "zai")) @@ -864,6 +871,7 @@ class LiteLLMCompletionResponsesConfig: ) request_kwargs: Final = _STR_KEY_DICT_ADAPTER.validate_python({**kwargs, "drop_params": drop_params}) + context_management: Final = responses_api_request.get("context_management") keeps_item_marks: Final = _target_keeps_item_marks( model, custom_llm_provider, drops_unsupported_params(request_kwargs.get("drop_params")) ) @@ -872,7 +880,7 @@ class LiteLLMCompletionResponsesConfig: responses_api_request=responses_api_request, replay_reasoning=True, item_cache_control=keeps_item_marks, - native_compaction=_carries_compaction_item(input) + native_compaction=_replays_compaction_block(input, context_management) and provider_builds_anthropic_request(model, custom_llm_provider), ) @@ -905,7 +913,7 @@ class LiteLLMCompletionResponsesConfig: "response_format": response_format, "reasoning_effort": reasoning.effort, "reasoning_summary": reasoning.summary, - "context_management": responses_api_request.get("context_management"), + "context_management": context_management, # litellm specific params "custom_llm_provider": custom_llm_provider, "extra_headers": extra_headers, @@ -959,9 +967,10 @@ class LiteLLMCompletionResponsesConfig: a strict message schema could refuse. ``native_compaction`` belongs to callers whose target builds an - Anthropic request: a ``compaction`` item then goes back as the - provider's own compaction block. Everyone else reads its summary as - assistant text. + Anthropic request that also asks for compaction in + ``context_management``: a ``compaction`` item then goes back as the + provider's own compaction block, which Anthropic rejects without that + compaction strategy. Everyone else reads its summary as assistant text. """ messages: list[ AllMessageValues diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index efdbe84babb..c135ae57e26 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -546,6 +546,7 @@ _ANTHROPIC_MESSAGES_URL: Final = "https://api.anthropic.com/v1/messages" _COMPACTION_BLOCK: Final = {"type": "compaction", "content": "The user pasted a long document and asked about it."} _FIRST_INPUT: Final = ({"role": "user", "content": "What is the first word of the document?"},) _NEXT_INPUT: Final = ({"role": "user", "content": "And the last word?"},) +_COMPACT_AT_50K: Final = ({"type": "compaction", "compact_threshold": 50000},) def _anthropic_message(*content: Mapping[str, object]) -> httpx.Response: @@ -583,8 +584,12 @@ def test_responses_over_anthropic_return_the_compaction_block_and_send_it_back( ] ) - first: Final = litellm.responses(model=_MODEL, input=list(_FIRST_INPUT), api_key="mock_api_key") - litellm.responses(model=_MODEL, input=_replay(first.output), api_key="mock_api_key") + first: Final = litellm.responses( + model=_MODEL, input=list(_FIRST_INPUT), context_management=list(_COMPACT_AT_50K), api_key="mock_api_key" + ) + litellm.responses( + model=_MODEL, input=_replay(first.output), context_management=list(_COMPACT_AT_50K), api_key="mock_api_key" + ) assert [item.type for item in first.output] == ["compaction", "message"] assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] @@ -628,14 +633,24 @@ def test_streamed_responses_over_anthropic_emit_the_compaction_item_and_send_it_ ] ) - events: Final = tuple(litellm.responses(model=_MODEL, input=list(_FIRST_INPUT), stream=True, api_key="mock_api_key")) + events: Final = tuple( + litellm.responses( + model=_MODEL, + input=list(_FIRST_INPUT), + context_management=list(_COMPACT_AT_50K), + stream=True, + api_key="mock_api_key", + ) + ) item_events: Final = tuple( event for event in events if event.type in (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE) ) completed: Final = events[-1].response - litellm.responses(model=_MODEL, input=_replay(completed.output), api_key="mock_api_key") + litellm.responses( + model=_MODEL, input=_replay(completed.output), context_management=list(_COMPACT_AT_50K), api_key="mock_api_key" + ) assert [(event.type, event.output_index, event.item.type) for event in item_events] == [ (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, 0, "compaction"), @@ -648,3 +663,20 @@ def test_streamed_responses_over_anthropic_emit_the_compaction_item_and_send_it_ ] assert completed.output[0].id == item_events[1].item.id assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] + + +def test_responses_over_anthropic_send_a_compaction_item_as_summary_text_when_the_request_does_not_compact( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=_anthropic_message({"type": "text", "text": "dog."}) + ) + compaction_item: Final = { + "type": "compaction", + "id": "cmp_earlier", + "encrypted_content": json.dumps(_COMPACTION_BLOCK), + } + + litellm.responses(model=_MODEL, input=[compaction_item, *_NEXT_INPUT], api_key="mock_api_key") + + assert _sent_assistant_content(route) == [{"type": "text", "text": _COMPACTION_BLOCK["content"]}] diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 98c045bcd6e..e539f47e451 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -6469,18 +6469,26 @@ def test_transform_request_reads_service_tier_from_responses_api_request() -> No _COMPACTION_BLOCK: Final = {"type": "compaction", "content": "Earlier turns covered the Tokyo trip."} +_COMPACTION_REQUESTED: Final = ResponsesAPIOptionalRequestParams( + context_management=[{"type": "compaction", "compact_threshold": 50000}] +) +_SUMMARY_AS_TEXT: Final = {"role": "assistant", "content": _COMPACTION_BLOCK["content"]} + + @pytest.mark.parametrize( - ("model", "replayed_compaction"), + ("model", "responses_api_request", "replayed_compaction"), [ ( "anthropic/claude-sonnet-5-5", + _COMPACTION_REQUESTED, {"role": "assistant", "content": None, "provider_specific_fields": {"compaction_blocks": [_COMPACTION_BLOCK]}}, ), - ("gemini/gemini-3.8-flash", {"role": "assistant", "content": _COMPACTION_BLOCK["content"]}), + ("anthropic/claude-sonnet-5-5", ResponsesAPIOptionalRequestParams(), _SUMMARY_AS_TEXT), + ("gemini/gemini-3.8-flash", _COMPACTION_REQUESTED, _SUMMARY_AS_TEXT), ], ) -def test_compaction_input_item_replays_as_a_compaction_block_only_on_anthropic_routes( - model: str, replayed_compaction: Mapping[str, object] +def test_compaction_input_item_replays_as_a_compaction_block_only_when_anthropic_compacts( + model: str, responses_api_request: ResponsesAPIOptionalRequestParams, replayed_compaction: Mapping[str, object] ) -> None: result: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( model=model, @@ -6492,7 +6500,7 @@ def test_compaction_input_item_replays_as_a_compaction_block_only_on_anthropic_r {"role": "user", "content": "And Osaka?"}, ], ), - responses_api_request={}, + responses_api_request=responses_api_request, ) assert [dict(message) for message in result["messages"]] == [ From 7a91a525afd8ebc6814cf98d85c654361ca01b85 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 14:45:20 -0700 Subject: [PATCH 03/21] fix(responses): carry compaction through dict-form edits, SDK replay, affinity, and guardrails Read Anthropic's dict-form context_management, accept SDK ResponseCompactionItem objects on replay, strip the encrypted_content affinity wrap before decoding, scan and rewrite compaction summaries in Responses guardrails, drop the history a summary replay replaces, send the summary as a user message so thinking ordering holds, and send the streamed compaction .done with the newest block. --- .../guardrail_translation/handler.py | 15 ++ .../compaction_items.py | 74 ++++++++- .../streaming_iterator.py | 46 +++-- .../transformation.py | 44 +++-- ...test_openai_responses_guardrail_handler.py | 26 +++ .../test_anthropic_responses_bridge.py | 157 +++++++++++++++++- .../test_litellm_completion_responses.py | 51 +++++- 7 files changed, 373 insertions(+), 40 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 435c11c0602..1f5d04b5d60 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -64,6 +64,10 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( unappliable_request_rewrite, ) from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools +from litellm.responses.litellm_completion_transformation.compaction_items import ( + compaction_summary, + with_compaction_summary, +) from litellm.responses.litellm_completion_transformation.custom_tools import without_malformed_tool_flags from litellm.responses.litellm_completion_transformation.tool_search import is_server_tool_search_tool from litellm.responses.litellm_completion_transformation.transformation import ( @@ -803,6 +807,11 @@ class OpenAIResponsesHandler(BaseTranslation): Override this method to customize text/image extraction logic. """ + summary: Final = compaction_summary(message) + if summary is not None: + texts_to_check.append(summary) + task_mappings.append((msg_idx, None)) + return content: Final = message.get("content", None) if content is None: return @@ -848,6 +857,12 @@ class OpenAIResponsesHandler(BaseTranslation): msg_idx = cast(int, mapping[0]) content_idx_optional = cast(int | None, mapping[1]) + if compaction_summary(messages[msg_idx]) is not None: + messages[msg_idx]["encrypted_content"] = with_compaction_summary( # rebind-ok: in-place input rewrite + messages[msg_idx]["encrypted_content"], guardrail_response + ) + continue + content = messages[msg_idx].get("content", None) if content is None: continue diff --git a/litellm/responses/litellm_completion_transformation/compaction_items.py b/litellm/responses/litellm_completion_transformation/compaction_items.py index d9f4c95a060..81f080e0d7a 100644 --- a/litellm/responses/litellm_completion_transformation/compaction_items.py +++ b/litellm/responses/litellm_completion_transformation/compaction_items.py @@ -5,10 +5,21 @@ from dataclasses import dataclass from typing import Final from openai.types.responses.response_compaction_item import ResponseCompactionItem +from openai.types.responses.response_create_params import ResponseInputParam from pydantic import TypeAdapter, ValidationError +from litellm.responses.utils import ResponsesAPIRequestUtils + _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) _OBJECT_LIST: Final = TypeAdapter(list[object]) +_INSTRUCTION_ROLES: Final = frozenset(("system", "developer")) + + +def _json_object(value: object) -> Mapping[str, object] | None: + try: + return _JSON_OBJECT.validate_python(value) + except ValidationError: + return None @dataclass(frozen=True, slots=True) @@ -18,9 +29,8 @@ class CompactionBlock: def _compaction_block(value: object) -> CompactionBlock | None: - try: - block: Final = _JSON_OBJECT.validate_python(value) - except ValidationError: + block: Final = _json_object(value) + if block is None: return None summary: Final = block.get("content") if block.get("type") != "compaction" or not isinstance(summary, str) or not summary: @@ -50,11 +60,67 @@ def compaction_output_item(provider_specific_fields: object) -> ResponseCompacti ) +def _entry_type(entry: object) -> object: + fields: Final = _json_object(entry) + return None if fields is None else fields.get("type") + + +def _has_entry_of_type(entries: object, entry_type: str) -> bool: + try: + rows: Final = _OBJECT_LIST.validate_python(entries) + except ValidationError: + return False + return any(_entry_type(row) == entry_type for row in rows) + + +def requests_compaction(context_management: object) -> bool: + if _has_entry_of_type(context_management, "compaction"): + return True + anthropic_form: Final = _json_object(context_management) + return anthropic_form is not None and _has_entry_of_type(anthropic_form.get("edits"), "compact_20260112") + + +def _without_affinity_wrap(encrypted_content: str) -> str: + _, original = ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id(encrypted_content) + return original + + def decode_compaction_item(encrypted_content: object) -> CompactionBlock | None: if not isinstance(encrypted_content, str): return None try: - decoded: Final = _JSON_OBJECT.validate_json(encrypted_content) + decoded: Final = _JSON_OBJECT.validate_json(_without_affinity_wrap(encrypted_content)) except ValidationError: return None return _compaction_block(decoded) + + +def compaction_summary(item: object) -> str | None: + fields: Final = _json_object(item) + if fields is None or fields.get("type") != "compaction": + return None + block: Final = decode_compaction_item(fields.get("encrypted_content")) + return None if block is None else block.summary + + +def with_compaction_summary(encrypted_content: str, summary: str) -> str: + original: Final = _without_affinity_wrap(encrypted_content) + block: Final = decode_compaction_item(original) + if block is None: + return encrypted_content + rewritten: Final = json.dumps({**block.wire_block, "content": summary}, separators=(",", ":")) + return encrypted_content.removesuffix(original) + rewritten + + +def _is_instruction_item(item: object) -> bool: + fields: Final = _json_object(item) + return fields is not None and fields.get("role") in _INSTRUCTION_ROLES + + +def without_items_before_compaction(input: str | ResponseInputParam) -> str | ResponseInputParam: + if not isinstance(input, list): + return input + newest: Final = next( + (index for index in range(len(input) - 1, 0, -1) if compaction_summary(input[index]) is not None), 0 + ) + return [*(item for item in input[:newest] if _is_instruction_item(item)), *input[newest:]] diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 0b3fb8ef25f..68544d466af 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -148,6 +148,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._reasoning_output_index: int | None = None self._compaction_output_index: int | None = None self._compaction_item_id: str | None = None + self._sent_compaction_item_done: bool = False self._cached_response_id: str | None = None self._buffered_chunk: ModelResponseStream | None = None self._upstream_exhausted: bool = False @@ -263,27 +264,41 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): index: Final = self._streamed_output_index(item) return (index is None, index or 0) - def _queue_compaction_item_events(self) -> None: + def _streamed_compaction_item(self) -> BaseLiteLLMOpenAIResponseObject | None: + item: Final = compaction_output_item(self._accumulated_provider_specific_fields) + if item is None: + return None + self._compaction_item_id = self._compaction_item_id or item.id + return BaseLiteLLMOpenAIResponseObject(**{**item.model_dump(exclude_none=True), "id": self._compaction_item_id}) + + def _queue_compaction_item_added(self) -> None: if ( self._compaction_output_index is not None or "compaction_blocks" not in self._accumulated_provider_specific_fields ): return - item: Final = compaction_output_item(self._accumulated_provider_specific_fields) - if item is None: + streamed_item: Final = self._streamed_compaction_item() + if streamed_item is None: return - self._compaction_item_id = item.id self._compaction_output_index = self._allocate_output_index() - streamed_item: Final = BaseLiteLLMOpenAIResponseObject(**item.model_dump(exclude_none=True)) - added: Final = self._next_output_item_added(self._compaction_output_index, streamed_item) + self._pending_response_events.append(self._next_output_item_added(self._compaction_output_index, streamed_item)) + + def _queue_compaction_item_done(self) -> None: + if self._compaction_output_index is None or self._sent_compaction_item_done: + return + streamed_item: Final = self._streamed_compaction_item() + if streamed_item is None: + return + self._sent_compaction_item_done = True self._sequence_number += 1 - done: Final = OutputItemDoneEvent( - type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, - output_index=self._compaction_output_index, - sequence_number=self._sequence_number, - item=streamed_item, + self._pending_response_events.append( + OutputItemDoneEvent( + type=ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, + output_index=self._compaction_output_index, + sequence_number=self._sequence_number, + item=streamed_item, + ) ) - self._pending_response_events.extend((added, done)) def _normalize_tool_call_index(self, tool_call: object) -> int | None: idx_raw: Final = tool_call.get("index") if isinstance(tool_call, dict) else getattr(tool_call, "index", None) @@ -1013,7 +1028,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if not self.litellm_model_response or isinstance(self.litellm_model_response, TextCompletionResponse): self.litellm_model_response = self.create_litellm_model_response() if self.litellm_model_response: - self._queue_compaction_item_events() + self._queue_compaction_item_added() + self._queue_compaction_item_done() if self._pending_response_events: return self._pending_response_events.pop(0) # If tool calls exist, emit tool events before finishing/response.completed. @@ -1118,7 +1134,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if chunk is not None: chunk = cast(ModelResponseStream, chunk) if not _streams_compaction(chunk): - self._queue_compaction_item_events() + self._queue_compaction_item_added() for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -1228,7 +1244,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): else: chunk = self.litellm_custom_stream_wrapper.__next__() if not _streams_compaction(chunk): - self._queue_compaction_item_events() + self._queue_compaction_item_added() for src in ( getattr(chunk, "provider_specific_fields", None), getattr( diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index eab8fe92b93..c13d8cbe65c 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -79,7 +79,6 @@ from litellm.types.llms.openai import ( ChatCompletionToolParam, ChatCompletionToolParamFunctionChunk, ChatCompletionUserMessage, - ContextManagementEntry, GenericChatCompletionMessage, IncompleteDetails, InputTokensDetails, @@ -117,7 +116,13 @@ from litellm.types.utils import ( Usage, ) -from .compaction_items import compaction_output_item, decode_compaction_item +from .compaction_items import ( + compaction_output_item, + compaction_summary, + decode_compaction_item, + requests_compaction, + without_items_before_compaction, +) from .custom_tools import ( convert_custom_tool_to_function_tool, extract_custom_tool_names, @@ -342,13 +347,23 @@ def without_item_cache_control(input: str | ResponseInputParam) -> str | Respons ) -def _replays_compaction_block( - input: str | ResponseInputParam, context_management: Sequence[ContextManagementEntry] | None -) -> bool: +def _as_input_param(item: object) -> object: + return item.model_dump(exclude_none=True) if isinstance(item, ResponseCompactionItem) else item + + +def with_compaction_items_as_params(input: str | ResponseInputParam) -> str | ResponseInputParam: + if not isinstance(input, list): + return input + return cast( # cast-ok: the same input items, a compaction item replayed from SDK output dumped to its param dict + ResponseInputParam, [_as_input_param(item) for item in input] + ) + + +def _replays_compaction_block(input: str | ResponseInputParam, context_management: object) -> bool: return ( - isinstance(input, list) - and any(item.get("type") == "compaction" for item in input) - and any(entry.get("type") == "compaction" for entry in context_management or ()) + requests_compaction(context_management) + and isinstance(input, list) + and any(compaction_summary(item) is not None for item in input) ) @@ -875,13 +890,16 @@ class LiteLLMCompletionResponsesConfig: keeps_item_marks: Final = _target_keeps_item_marks( model, custom_llm_provider, drops_unsupported_params(request_kwargs.get("drop_params")) ) + request_input: Final = with_compaction_items_as_params(input) + native_compaction: Final = _replays_compaction_block( + request_input, context_management + ) and provider_builds_anthropic_request(model, custom_llm_provider) bridged_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=input, + input=request_input if native_compaction else without_items_before_compaction(request_input), responses_api_request=responses_api_request, replay_reasoning=True, item_cache_control=keeps_item_marks, - native_compaction=_replays_compaction_block(input, context_management) - and provider_builds_anthropic_request(model, custom_llm_provider), + native_compaction=native_compaction, ) litellm_completion_request: dict = { @@ -1946,7 +1964,7 @@ class LiteLLMCompletionResponsesConfig: @staticmethod def _transform_compaction_item_to_chat_completion_messages( input_item: Mapping[str, object], native_compaction: bool - ) -> tuple[GenericChatCompletionMessage | ChatCompletionResponseMessage, ...]: + ) -> tuple[ChatCompletionUserMessage | ChatCompletionResponseMessage, ...]: compaction: Final = decode_compaction_item(input_item.get("encrypted_content")) if compaction is None: return () @@ -1958,7 +1976,7 @@ class LiteLLMCompletionResponsesConfig: provider_specific_fields={"compaction_blocks": [dict(compaction.wire_block)]}, ), ) - return (GenericChatCompletionMessage(role="assistant", content=compaction.summary),) + return (ChatCompletionUserMessage(role="user", content=compaction.summary),) @staticmethod def _transform_responses_api_message_item( diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index bf871210074..6dbd19e991d 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -298,6 +298,32 @@ class TestOpenAIResponsesHandlerInputProcessing: assert result["input"][1]["content"] == "World [GUARDRAILED]" assert result["model"] == "gpt-4" + @pytest.mark.asyncio + async def test_compaction_summary_is_scanned_and_rewritten_inside_its_item(self) -> None: + affinity_prefix: Final = "litellm_enc:bW9kZWxfaWQ6ZGVwLWE=;" + data: Final = { + "input": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": affinity_prefix + + json.dumps({"type": "compaction", "content": "Cat is Biscuit."}), + }, + {"role": "user", "content": "Cat name?", "type": "message"}, + ], + "model": "claude-sonnet-5-5", + } + + result: Final = await OpenAIResponsesHandler().process_input_messages(data, MockGuardrail(guardrail_name="t")) + + rewritten: Final = result["input"][0]["encrypted_content"] + assert rewritten.startswith(affinity_prefix) + assert json.loads(rewritten.removeprefix(affinity_prefix)) == { + "type": "compaction", + "content": "Cat is Biscuit. [GUARDRAILED]", + } + assert result["input"][1]["content"] == "Cat name? [GUARDRAILED]" + @pytest.mark.asyncio async def test_process_input_list_with_multimodal_content(self): """Test processing list input with multimodal content""" diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index c135ae57e26..bf105da7793 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -615,7 +615,11 @@ def test_streamed_responses_over_anthropic_emit_the_compaction_item_and_send_it_ "usage": {"input_tokens": 9, "output_tokens": 1}, }, }, - {"type": "content_block_start", "index": 0, "content_block": {"type": "compaction", "content": None}}, + { + "type": "content_block_start", + "index": 0, + "content_block": {"type": "compaction", "content": None}, + }, { "type": "content_block_delta", "index": 0, @@ -654,14 +658,142 @@ def test_streamed_responses_over_anthropic_emit_the_compaction_item_and_send_it_ assert [(event.type, event.output_index, event.item.type) for event in item_events] == [ (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, 0, "compaction"), - (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, 0, "compaction"), (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, 1, "message"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, 0, "compaction"), (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, 1, "message"), ] assert [(item.type, getattr(item, "encrypted_content", None)) for item in completed.output[:1]] == [ - ("compaction", item_events[1].item.encrypted_content) + ("compaction", item_events[2].item.encrypted_content) ] - assert completed.output[0].id == item_events[1].item.id + assert completed.output[0].id == item_events[2].item.id + assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] + + +def _compaction_stream_events(*summaries: str) -> tuple[Mapping[str, object], ...]: + return tuple( + chain.from_iterable( + ( + { + "type": "content_block_start", + "index": 2 * index, + "content_block": {"type": "compaction", "content": None}, + }, + { + "type": "content_block_delta", + "index": 2 * index, + "delta": {"type": "compaction_delta", "content": summary}, + }, + {"type": "content_block_stop", "index": 2 * index}, + {"type": "content_block_start", "index": 2 * index + 1, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 2 * index + 1, "delta": {"type": "text_delta", "text": "Ok."}}, + {"type": "content_block_stop", "index": 2 * index + 1}, + ) + for index, summary in enumerate(summaries) + ) + ) + + +def test_streamed_compaction_item_done_carries_the_same_newest_block_as_the_completed_response( + respx_mock: respx.MockRouter, +) -> None: + respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + return_value=httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_compacted_twice", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + *_compaction_stream_events("First summary.", "Second summary."), + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ) + ) + + events: Final = tuple( + litellm.responses( + model=_MODEL, + input=list(_FIRST_INPUT), + context_management=list(_COMPACT_AT_50K), + stream=True, + api_key="mock_api_key", + ) + ) + streamed: Final = tuple( + (event.item.id, json.loads(event.item.encrypted_content)["content"]) + for event in events + if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and event.item.type == "compaction" + ) + completed: Final = tuple( + (item.id, json.loads(item.encrypted_content)["content"]) + for item in events[-1].response.output + if item.type == "compaction" + ) + + assert streamed == completed + assert [summary for _, summary in completed] == ["Second summary."] + + +def test_responses_over_anthropic_send_back_a_compaction_item_replayed_as_sdk_output( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + side_effect=[ + _anthropic_message(_COMPACTION_BLOCK, {"type": "text", "text": "The."}), + _anthropic_message({"type": "text", "text": "dog."}), + ] + ) + + first: Final = litellm.responses( + model=_MODEL, input=list(_FIRST_INPUT), context_management=list(_COMPACT_AT_50K), api_key="mock_api_key" + ) + litellm.responses( + model=_MODEL, + input=[*_FIRST_INPUT, *first.output, *_NEXT_INPUT], + context_management=list(_COMPACT_AT_50K), + api_key="mock_api_key", + ) + + assert _sent_assistant_content(route)[0] == _COMPACTION_BLOCK + + +@pytest.mark.asyncio +async def test_router_with_encrypted_content_affinity_sends_the_compaction_block_back( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + route: Final = respx_mock.post(_ANTHROPIC_MESSAGES_URL).mock( + side_effect=[ + _anthropic_message(_COMPACTION_BLOCK, {"type": "text", "text": "The."}), + _anthropic_message({"type": "text", "text": "dog."}), + ] + ) + router: Final = litellm.Router( + model_list=[ + { + "model_name": "claude", + "litellm_params": {"model": _MODEL, "api_key": "mock_api_key"}, + "model_info": {"id": "deployment-a"}, + } + ], + optional_pre_call_checks=["encrypted_content_affinity"], + ) + + first: Final = await router.aresponses( + model="claude", input=list(_FIRST_INPUT), context_management=list(_COMPACT_AT_50K) + ) + await router.aresponses(model="claude", input=_replay(first.output), context_management=list(_COMPACT_AT_50K)) + + assert first.output[0].encrypted_content.startswith("litellm_enc:") assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] @@ -677,6 +809,19 @@ def test_responses_over_anthropic_send_a_compaction_item_as_summary_text_when_th "encrypted_content": json.dumps(_COMPACTION_BLOCK), } - litellm.responses(model=_MODEL, input=[compaction_item, *_NEXT_INPUT], api_key="mock_api_key") + litellm.responses( + model=_MODEL, + input=[*_FIRST_INPUT, {"role": "assistant", "content": "The."}, compaction_item, *_NEXT_INPUT], + api_key="mock_api_key", + ) - assert _sent_assistant_content(route) == [{"type": "text", "text": _COMPACTION_BLOCK["content"]}] + sent: Final = json.loads(route.calls.last.request.content) + assert sent["messages"] == [ + { + "role": "user", + "content": [ + {"type": "text", "text": _COMPACTION_BLOCK["content"]}, + {"type": "text", "text": _NEXT_INPUT[0]["content"]}, + ], + } + ] diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index e539f47e451..2deeeaeac57 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -6472,7 +6472,13 @@ _COMPACTION_BLOCK: Final = {"type": "compaction", "content": "Earlier turns cove _COMPACTION_REQUESTED: Final = ResponsesAPIOptionalRequestParams( context_management=[{"type": "compaction", "compact_threshold": 50000}] ) -_SUMMARY_AS_TEXT: Final = {"role": "assistant", "content": _COMPACTION_BLOCK["content"]} +_SUMMARY_AS_TEXT: Final = {"role": "user", "content": _COMPACTION_BLOCK["content"]} + + +def _anthropic_form_context_management(edit: Mapping[str, object]) -> ResponsesAPIOptionalRequestParams: + return cast( # cast-ok: Anthropic's native {"edits": [...]} form, which the param type does not declare + ResponsesAPIOptionalRequestParams, {"context_management": {"edits": [edit]}} + ) @pytest.mark.parametrize( @@ -6481,9 +6487,27 @@ _SUMMARY_AS_TEXT: Final = {"role": "assistant", "content": _COMPACTION_BLOCK["co ( "anthropic/claude-sonnet-5-5", _COMPACTION_REQUESTED, - {"role": "assistant", "content": None, "provider_specific_fields": {"compaction_blocks": [_COMPACTION_BLOCK]}}, + { + "role": "assistant", + "content": None, + "provider_specific_fields": {"compaction_blocks": [_COMPACTION_BLOCK]}, + }, + ), + ( + "anthropic/claude-sonnet-5-5", + _anthropic_form_context_management({"type": "compact_20260112"}), + { + "role": "assistant", + "content": None, + "provider_specific_fields": {"compaction_blocks": [_COMPACTION_BLOCK]}, + }, ), ("anthropic/claude-sonnet-5-5", ResponsesAPIOptionalRequestParams(), _SUMMARY_AS_TEXT), + ( + "anthropic/claude-sonnet-5-5", + _anthropic_form_context_management({"type": "clear_tool_uses_20250919"}), + _SUMMARY_AS_TEXT, + ), ("gemini/gemini-3.8-flash", _COMPACTION_REQUESTED, _SUMMARY_AS_TEXT), ], ) @@ -6507,3 +6531,26 @@ def test_compaction_input_item_replays_as_a_compaction_block_only_when_anthropic replayed_compaction, {"role": "user", "content": "And Osaka?"}, ] + + +def test_summary_replay_drops_the_history_its_compaction_item_replaces() -> None: + result: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="gemini/gemini-3.8-flash", + input=cast( + ResponseInputParam, + [ + {"role": "system", "content": "Answer briefly."}, + {"role": "user", "content": "Plan a Tokyo trip."}, + {"role": "assistant", "content": "Day one, Asakusa."}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"role": "user", "content": "And Osaka?"}, + ], + ), + responses_api_request=_COMPACTION_REQUESTED, + ) + + assert [dict(message) for message in result["messages"]] == [ + {"role": "system", "content": "Answer briefly."}, + _SUMMARY_AS_TEXT, + {"role": "user", "content": "And Osaka?"}, + ] From 33f90b261663553d821df30f2a5c7142d0f1370d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 17:30:09 -0700 Subject: [PATCH 04/21] fix(responses): count, route, guard, and redact compaction items consistently A compaction item now reads as its summary everywhere the proxy counts or routes a Responses request: pre-call context checks, the TPM limiter, the input token endpoint, the stream usage estimate, and the previous_response_id session rebuild all drop the history the item replaces. A Claude compaction item sent to a non-Anthropic model (including a fallback) becomes a user summary message instead of an encrypted blob the provider rejects, and Claude on Bedrock Claude Platform takes the summary path. Output guardrails see the summary of a returned compaction item, and logging redaction hides it. --- litellm/litellm_core_utils/redact_messages.py | 46 ++++---- .../guardrail_translation/handler.py | 59 ++++++++-- .../hooks/parallel_request_limiter_v3.py | 7 +- .../proxy/response_api_endpoints/endpoints.py | 1 + .../compaction_items.py | 55 ++++++++-- .../session_handler.py | 20 +++- .../streaming_iterator.py | 5 +- .../tool_search.py | 18 ++-- .../transformation.py | 28 +++-- litellm/responses/main.py | 14 ++- litellm/responses/streaming_iterator.py | 1 + litellm/router.py | 1 + .../test_redact_messages.py | 21 ++++ ...test_openai_responses_guardrail_handler.py | 101 ++++++++++++++++++ tests/unit/proxy/hooks/test_tpm_concurrent.py | 34 +++++- .../response_api_endpoints/test_endpoints.py | 23 ++++ .../test_anthropic_responses_bridge.py | 2 + .../test_litellm_completion_responses.py | 1 + .../test_session_handler.py | 55 ++++++++++ .../test_responses_api_request_body.py | 48 +++++++++ .../unit/responses/test_streaming_iterator.py | 21 ++++ tests/unit/test_router/test_router.py | 23 ++++ 22 files changed, 526 insertions(+), 58 deletions(-) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index c2cca5739cf..7247ba2e4be 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items from litellm.llms.vertex_ai.common_utils import ( redact_vertex_ai_metadata_from_litellm_params, redact_vertex_ai_metadata_from_logged_object, @@ -147,33 +148,39 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[ return { **fields, **({"reasoning_content": REDACTED_BY_LITELLM} if fields.get("reasoning_content") is not None else {}), - **{key: None for key in ("thinking_blocks", "reasoningContentBlocks", "citations") if key in fields}, + **{ + key: None + for key in ("thinking_blocks", "reasoningContentBlocks", "citations", "compaction_blocks") + if key in fields + }, } +def _redact_field_when_present(item: object, field: str) -> None: + if getattr(item, field, None) is not None: + setattr(item, field, REDACTED_BY_LITELLM) + + def _redact_responses_api_output(output_items): """Helper to redact ResponsesAPIResponse output items.""" for output_item in output_items: + output_type = stream_item_field(output_item, "type") if getattr(output_item, "text", None) is not None: output_item.text = REDACTED_BY_LITELLM - if hasattr(output_item, "content") and isinstance(output_item.content, list): - for content_part in output_item.content: - if getattr(content_part, "text", None) is not None: - content_part.text = REDACTED_BY_LITELLM - if getattr(content_part, "refusal", None) is not None: - content_part.refusal = REDACTED_BY_LITELLM + for content_part in stream_item_items(output_item, "content"): + _redact_field_when_present(content_part, "text") + _redact_field_when_present(content_part, "refusal") - # Redact reasoning items in output array - if hasattr(output_item, "type") and output_item.type == "reasoning": - if hasattr(output_item, "summary") and isinstance(output_item.summary, list): - for summary_item in output_item.summary: - if getattr(summary_item, "text", None) is not None: - summary_item.text = REDACTED_BY_LITELLM + if output_type == "reasoning": + for summary_item in stream_item_items(output_item, "summary"): + _redact_field_when_present(summary_item, "text") - if hasattr(output_item, "type") and output_item.type == "function_call" and hasattr(output_item, "arguments"): + if output_type == "compaction" and hasattr(output_item, "encrypted_content"): + output_item.encrypted_content = REDACTED_BY_LITELLM + if output_type == "function_call" and hasattr(output_item, "arguments"): output_item.arguments = REDACTED_BY_LITELLM - if hasattr(output_item, "type") and output_item.type == "custom_tool_call" and hasattr(output_item, "input"): + if output_type == "custom_tool_call" and hasattr(output_item, "input"): output_item.input = REDACTED_BY_LITELLM @@ -182,6 +189,7 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str): for output_item in output_items: if not isinstance(output_item, dict): continue + output_type = stream_item_field(output_item, "type") if output_item.get("text") is not None: output_item["text"] = redacted_str @@ -193,14 +201,16 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str): if isinstance(content_item, dict) and content_item.get("refusal") is not None: content_item["refusal"] = redacted_str - if output_item.get("type") == "reasoning" and isinstance(output_item.get("summary"), list): + if output_type == "reasoning" and isinstance(output_item.get("summary"), list): for summary_item in output_item["summary"]: if isinstance(summary_item, dict) and summary_item.get("text") is not None: summary_item["text"] = redacted_str - if output_item.get("type") == "function_call" and "arguments" in output_item: + if output_type == "compaction" and "encrypted_content" in output_item: + output_item["encrypted_content"] = redacted_str + if output_type == "function_call" and "arguments" in output_item: output_item["arguments"] = redacted_str - if output_item.get("type") == "custom_tool_call" and "input" in output_item: + if output_type == "custom_tool_call" and "input" in output_item: output_item["input"] = redacted_str diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 1f5d04b5d60..bda2eab26b4 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -66,6 +66,7 @@ from litellm.llms.base_llm.guardrail_translation.utils import ( from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.responses.litellm_completion_transformation.compaction_items import ( compaction_summary, + decode_compaction_item, with_compaction_summary, ) from litellm.responses.litellm_completion_transformation.custom_tools import without_malformed_tool_flags @@ -113,6 +114,16 @@ if TYPE_CHECKING: from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +_COMPACTION_SUMMARY_POSITION: Final = -1 + + +def _output_compaction_summary(item: object) -> str | None: + if stream_item_field(item, "type") != "compaction": + return None + block: Final = decode_compaction_item(stream_item_field(item, "encrypted_content")) + return None if block is None else block.summary + + class _ToolCallShape(NamedTuple): name: str | None arguments: str @@ -1130,13 +1141,17 @@ class OpenAIResponsesHandler(BaseTranslation): # delivered, so a truncated buffer still reports text undeliverable. # # ------------------------------------------------------------------ # if final_chunk.get("type") == "response.output_item.done": + finished_summary: Final = _output_compaction_summary(stream_item_field(final_chunk, "item")) model_response_stream: Final = ( OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream(final_chunk) ) tool_calls: Final = model_response_stream.choices[0].delta.tool_calls - if tool_calls: + if tool_calls or finished_summary is not None: inputs = GenericGuardrailAPIInputs() - inputs["tool_calls"] = cast(list[ChatCompletionToolCallChunk], tool_calls) + if finished_summary is not None: + inputs["texts"] = [finished_summary] + if tool_calls: + inputs["tool_calls"] = cast(list[ChatCompletionToolCallChunk], tool_calls) if hasattr(model_response_stream, "model") and model_response_stream.model: inputs["model"] = model_response_stream.model await guardrail_to_apply.apply_guardrail( @@ -1228,6 +1243,14 @@ class OpenAIResponsesHandler(BaseTranslation): else: setattr(event, field, value) + @staticmethod + def _write_compaction_summary(item: object, summary: str) -> None: + encrypted_content: Final = stream_item_field(item, "encrypted_content") + if isinstance(encrypted_content, str): + OpenAIResponsesHandler._write_event_field( + item, "encrypted_content", with_compaction_summary(encrypted_content, summary) + ) + def _sync_stream_events_with_rewrites( self, stream_events: Sequence[Any], @@ -1270,11 +1293,13 @@ class OpenAIResponsesHandler(BaseTranslation): output_index: int, rewrites_by_position: Mapping[tuple[int, int], str], ) -> None: - content: Final = item.get("content") if isinstance(item, dict) else getattr(item, "content", None) - if not isinstance(content, list): + summary_rewrite: Final = rewrites_by_position.get((output_index, _COMPACTION_SUMMARY_POSITION)) + if summary_rewrite is not None: + OpenAIResponsesHandler._write_compaction_summary(item, summary_rewrite) return + content: Final = stream_item_items(item, "content") for (item_idx, content_idx), rewritten in rewrites_by_position.items(): - if item_idx != output_index or content_idx >= len(content): + if item_idx != output_index or not 0 <= content_idx < len(content): continue OpenAIResponsesHandler._write_event_field(content[content_idx], "text", rewritten) @@ -1495,13 +1520,15 @@ class OpenAIResponsesHandler(BaseTranslation): message_items: Final = tuple( item for item in output_items if stream_item_field(item, "type") not in _OUTPUT_TOOL_CALL_ITEM_TYPES ) + part_texts: Final = tuple( + text + for item in message_items + for part in stream_item_items(item, "content") + if isinstance(text := stream_item_field(part, "text"), str) and text + ) + summaries: Final = tuple(summary for item in output_items if (summary := _output_compaction_summary(item))) return StreamingScanKey( - texts=tuple( - text - for item in message_items - for part in stream_item_items(item, "content") - if isinstance(text := stream_item_field(part, "text"), str) and text - ), + texts=(*part_texts, *summaries), tool_calls=tuple( stream_item_fingerprint(item) for item in output_items @@ -1624,6 +1651,12 @@ class OpenAIResponsesHandler(BaseTranslation): ) return + summary: Final = _output_compaction_summary(output_item) + if summary is not None: + texts_to_check.append(summary) + task_mappings.append((output_idx, _COMPACTION_SUMMARY_POSITION)) + return + # Handle both GenericResponseOutputItem and dict content: list[OutputText] | list[dict] | None = None if isinstance(output_item, BaseModel): @@ -1691,6 +1724,10 @@ class OpenAIResponsesHandler(BaseTranslation): output_item = response_output[output_idx] + if content_idx == _COMPACTION_SUMMARY_POSITION: + self._write_compaction_summary(output_item, guardrail_response) + continue + # Handle both GenericResponseOutputItem, BaseModel, and dict if isinstance(output_item, GenericResponseOutputItem): if output_item.content and content_idx < len(output_item.content): diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 81e77283ed8..0a8b0c9ec05 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -71,6 +71,7 @@ from litellm.proxy.hooks.batch_enqueued_tokens import ( canonical_provider_batch_id, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.responses.litellm_completion_transformation.compaction_items import input_since_newest_compaction from litellm.router_utils.add_retry_fallback_headers import ( ensure_response_additional_headers, response_has_hidden_params, @@ -1225,7 +1226,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): estimable_data.get("input"), ) ) - messages, prompt, input_text = selected_fields + messages, prompt, raw_input = selected_fields + input_text: Final = ( + input_since_newest_compaction(raw_input) if call_type in RESPONSES_API_CALL_TYPES else raw_input + ) total_chars: Final = ( len(get_str_from_messages(messages)) @@ -3799,6 +3803,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=data.get("input") or "", responses_api_request=data, + only_since_compaction=True, ) @staticmethod diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 7c52499bd6f..6b825baa001 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1406,6 +1406,7 @@ def _responses_input_as_token_count_messages( transformed: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=input_value, responses_api_request=request_params, + only_since_compaction=True, ) return tuple( message if isinstance(message, dict) else message.model_dump(exclude_none=True) for message in transformed diff --git a/litellm/responses/litellm_completion_transformation/compaction_items.py b/litellm/responses/litellm_completion_transformation/compaction_items.py index 81f080e0d7a..e79b064d690 100644 --- a/litellm/responses/litellm_completion_transformation/compaction_items.py +++ b/litellm/responses/litellm_completion_transformation/compaction_items.py @@ -1,17 +1,22 @@ import json import uuid -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from dataclasses import dataclass -from typing import Final +from typing import Final, TypeVar +from openai.types.responses.easy_input_message_param import EasyInputMessageParam from openai.types.responses.response_compaction_item import ResponseCompactionItem from openai.types.responses.response_create_params import ResponseInputParam +from openai.types.responses.response_input_param import ResponseInputItemParam from pydantic import TypeAdapter, ValidationError from litellm.responses.utils import ResponsesAPIRequestUtils +from .tool_search import chat_config_builds_anthropic_request, provider_chat_config + _JSON_OBJECT: Final = TypeAdapter(dict[str, object]) _OBJECT_LIST: Final = TypeAdapter(list[object]) +_Item: Final = TypeVar("_Item") _INSTRUCTION_ROLES: Final = frozenset(("system", "developer")) @@ -80,6 +85,15 @@ def requests_compaction(context_management: object) -> bool: return anthropic_form is not None and _has_entry_of_type(anthropic_form.get("edits"), "compact_20260112") +def provider_takes_compaction_blocks(model: str, custom_llm_provider: str | None) -> bool: + from litellm.llms.bedrock.claude_platform.transformation import BedrockClaudePlatformConfig + + chat_config: Final = provider_chat_config(model, custom_llm_provider) + return chat_config_builds_anthropic_request(chat_config) and not isinstance( + chat_config, BedrockClaudePlatformConfig + ) + + def _without_affinity_wrap(encrypted_content: str) -> str: _, original = ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id(encrypted_content) return original @@ -112,15 +126,44 @@ def with_compaction_summary(encrypted_content: str, summary: str) -> str: return encrypted_content.removesuffix(original) + rewritten +def holds_compaction_summary(input: object) -> bool: + try: + items: Final = _OBJECT_LIST.validate_python(input) + except ValidationError: + return False + return any(compaction_summary(item) is not None for item in items) + + def _is_instruction_item(item: object) -> bool: fields: Final = _json_object(item) return fields is not None and fields.get("role") in _INSTRUCTION_ROLES -def without_items_before_compaction(input: str | ResponseInputParam) -> str | ResponseInputParam: - if not isinstance(input, list): - return input +def _items_from_newest_compaction(input: Sequence[_Item]) -> tuple[_Item, ...]: newest: Final = next( (index for index in range(len(input) - 1, 0, -1) if compaction_summary(input[index]) is not None), 0 ) - return [*(item for item in input[:newest] if _is_instruction_item(item)), *input[newest:]] + return (*(item for item in input[:newest] if _is_instruction_item(item)), *input[newest:]) + + +def without_items_before_compaction(input: str | ResponseInputParam) -> str | ResponseInputParam: + return list(_items_from_newest_compaction(input)) if isinstance(input, list) else input + + +def input_since_newest_compaction(input: object) -> object: + try: + items: Final = _OBJECT_LIST.validate_python(input) + except ValidationError: + return input + return list(_items_from_newest_compaction(items)) + + +def _as_summary_message(item: ResponseInputItemParam) -> ResponseInputItemParam: + summary: Final = compaction_summary(item) + return item if summary is None else EasyInputMessageParam(role="user", content=summary, type="message") + + +def with_compaction_items_as_summaries(input: str | ResponseInputParam) -> str | ResponseInputParam: + if not isinstance(input, list) or not holds_compaction_summary(input): + return input + return [_as_summary_message(item) for item in _items_from_newest_compaction(input)] diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 732c32ff2cb..24e848ceb8d 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -1,12 +1,13 @@ import asyncio import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler +from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( AllMessageValues, @@ -42,6 +43,17 @@ def _normalize_redacted_tool_call_arguments(message: Message) -> None: function_call.arguments = REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER +def _first_replayed_request(proxy_server_requests: Sequence[Mapping[str, object] | None]) -> int: + return next( + ( + index + for index in range(len(proxy_server_requests) - 1, -1, -1) + if (request := proxy_server_requests[index]) is not None and holds_compaction_summary(request.get("input")) + ), + 0, + ) + + def _stored_instructions(proxy_server_request: Mapping[str, object] | None) -> str | None: instructions: Final = None if proxy_server_request is None else proxy_server_request.get("instructions") return instructions if isinstance(instructions, str) and instructions else None @@ -81,7 +93,10 @@ class ResponsesSessionHandler: await ResponsesSessionHandler.get_proxy_server_request_from_spend_log(spend_log=spend_log) for spend_log in all_spend_logs ] - for spend_log, proxy_server_request_dict in zip(all_spend_logs, proxy_server_requests): + first_replayed: Final = _first_replayed_request(proxy_server_requests) + for spend_log, proxy_server_request_dict in zip( + all_spend_logs[first_replayed:], proxy_server_requests[first_replayed:] + ): chat_completion_message_history = ( ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload( spend_log=spend_log, @@ -152,6 +167,7 @@ class ResponsesSessionHandler: responses_api_request={}, replay_reasoning=True, item_cache_control=item_cache_control, + only_since_compaction=True, ) ) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 68544d466af..6918c2634ce 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -281,7 +281,10 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): if streamed_item is None: return self._compaction_output_index = self._allocate_output_index() - self._pending_response_events.append(self._next_output_item_added(self._compaction_output_index, streamed_item)) + unfinished_item: Final = streamed_item.model_copy(update={"encrypted_content": ""}) + self._pending_response_events.append( + self._next_output_item_added(self._compaction_output_index, unfinished_item) + ) def _queue_compaction_item_done(self) -> None: if self._compaction_output_index is None or self._sent_compaction_item_done: diff --git a/litellm/responses/litellm_completion_transformation/tool_search.py b/litellm/responses/litellm_completion_transformation/tool_search.py index 0a94acdd461..780a23e02b2 100644 --- a/litellm/responses/litellm_completion_transformation/tool_search.py +++ b/litellm/responses/litellm_completion_transformation/tool_search.py @@ -1,5 +1,5 @@ from collections.abc import Iterable, Iterator, Mapping, Sequence, Set -from typing import Final, Literal, TypeAlias +from typing import TYPE_CHECKING, Final, Literal, TypeAlias from pydantic import ConfigDict, TypeAdapter, ValidationError @@ -12,6 +12,9 @@ from litellm.types.utils import LlmProviders, Message from .custom_tools import openai_shaped_tool_call_item_id +if TYPE_CHECKING: + from litellm.llms.base_llm.chat.transformation import BaseConfig + __all__ = [ "ANTHROPIC_TOOL_SEARCH_TOOL_NAME", "ANTHROPIC_TOOL_SEARCH_TOOL_TYPE", @@ -31,6 +34,7 @@ __all__ = [ "is_openai_tool_search_tool", "is_server_tool_search_tool", "provider_builds_anthropic_request", + "provider_chat_config", "runs_tool_search_on_the_server", "tool_search_call_ids", "tools_with_anthropic_tool_search_mapped", @@ -159,7 +163,7 @@ def input_without_unanswered_anthropic_tool_search_calls( return input if len(kept) == len(input) else kept -def provider_builds_anthropic_request(model: str, custom_llm_provider: str | None) -> bool: +def provider_chat_config(model: str, custom_llm_provider: str | None) -> "BaseConfig | None": from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.utils import ProviderConfigManager @@ -167,10 +171,12 @@ def provider_builds_anthropic_request(model: str, custom_llm_provider: str | Non try: llm_provider: Final = LlmProviders(provider) except ValueError: - return False - return chat_config_builds_anthropic_request( - ProviderConfigManager.get_provider_chat_config(model=model, provider=llm_provider) - ) + return None + return ProviderConfigManager.get_provider_chat_config(model=model, provider=llm_provider) + + +def provider_builds_anthropic_request(model: str, custom_llm_provider: str | None) -> bool: + return chat_config_builds_anthropic_request(provider_chat_config(model, custom_llm_provider)) def chat_config_builds_anthropic_request(chat_config: object) -> bool: diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index c13d8cbe65c..b8075d7c078 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -118,8 +118,9 @@ from litellm.types.utils import ( from .compaction_items import ( compaction_output_item, - compaction_summary, decode_compaction_item, + holds_compaction_summary, + provider_takes_compaction_blocks, requests_compaction, without_items_before_compaction, ) @@ -360,11 +361,7 @@ def with_compaction_items_as_params(input: str | ResponseInputParam) -> str | Re def _replays_compaction_block(input: str | ResponseInputParam, context_management: object) -> bool: - return ( - requests_compaction(context_management) - and isinstance(input, list) - and any(compaction_summary(item) is not None for item in input) - ) + return requests_compaction(context_management) and holds_compaction_summary(input) _PROVIDERS_KEEPING_MESSAGE_CACHE_CONTROL: Final = frozenset(("anthropic", "dashscope", "minimax", "sap", "zai")) @@ -893,13 +890,14 @@ class LiteLLMCompletionResponsesConfig: request_input: Final = with_compaction_items_as_params(input) native_compaction: Final = _replays_compaction_block( request_input, context_management - ) and provider_builds_anthropic_request(model, custom_llm_provider) + ) and provider_takes_compaction_blocks(model, custom_llm_provider) bridged_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=request_input if native_compaction else without_items_before_compaction(request_input), + input=request_input, responses_api_request=responses_api_request, replay_reasoning=True, item_cache_control=keeps_item_marks, native_compaction=native_compaction, + only_since_compaction=not native_compaction, ) litellm_completion_request: dict = { @@ -962,6 +960,7 @@ class LiteLLMCompletionResponsesConfig: replay_reasoning: bool = False, item_cache_control: bool = False, native_compaction: bool = False, + only_since_compaction: bool = False, ) -> list[ AllMessageValues | GenericChatCompletionMessage @@ -988,7 +987,14 @@ class LiteLLMCompletionResponsesConfig: Anthropic request that also asks for compaction in ``context_management``: a ``compaction`` item then goes back as the provider's own compaction block, which Anthropic rejects without that - compaction strategy. Everyone else reads its summary as assistant text. + compaction strategy. Everyone else reads its summary as a user message. + + ``only_since_compaction`` belongs to callers that send, count, or + rebuild what the model reads: the items before the newest + ``compaction`` item are dropped, keeping system and developer + messages, the way Anthropic ignores content before a compaction block. + Guardrail scans leave it off because they write rewrites back by input + item position. """ messages: list[ AllMessageValues @@ -1004,9 +1010,11 @@ class LiteLLMCompletionResponsesConfig: ) ) + dumped_input: Final = with_compaction_items_as_params(input) + read_input: Final = without_items_before_compaction(dumped_input) if only_since_compaction else dumped_input messages.extend( LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message( - input=input if item_cache_control else without_item_cache_control(input), + input=read_input if item_cache_control else without_item_cache_control(read_input), replay_reasoning=replay_reasoning, native_compaction=native_compaction, ) diff --git a/litellm/responses/main.py b/litellm/responses/main.py index c90455e8e75..906b56a0c42 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -37,9 +37,15 @@ from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_h from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig +from litellm.responses.litellm_completion_transformation.compaction_items import ( + with_compaction_items_as_summaries, +) from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) +from litellm.responses.litellm_completion_transformation.transformation import ( + with_compaction_items_as_params, +) from litellm.responses.mcp.request_context import MCPRequestContext from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( @@ -1526,7 +1532,9 @@ def responses( ) # Decode any litellm-encoded encrypted-content item IDs back to their original IDs - input = ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(input) + input = with_compaction_items_as_summaries( + with_compaction_items_as_params(ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(input)) + ) # Call the handler with _is_async flag instead of directly calling the async handler if custom_llm_provider is None: @@ -2364,7 +2372,9 @@ def compact_responses( # Decode any litellm-encoded encrypted-content item IDs back to their original IDs # before forwarding to the upstream provider. - input = ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(input) + input = with_compaction_items_as_summaries( + with_compaction_items_as_params(ResponsesAPIRequestUtils.restore_encrypted_content_item_ids_in_input(input)) + ) # Call the handler with _is_async flag instead of directly calling the async handler response = base_llm_http_handler.compact_response_api_handler( diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 7532a528248..0f99fe43a09 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1566,6 +1566,7 @@ def _estimate_usage_from_text( messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped input=request_input, # pyright: ignore[reportArgumentType] # the raw Responses API input is a str or ResponseInputParam list, matching the helper's declared union responses_api_request=dict(responses_api_request), + only_since_compaction=True, ) input_tokens: Final = litellm.token_counter( # pyright: ignore[reportUnknownMemberType] # token_counter's public signature is untyped model=model, messages=messages diff --git a/litellm/router.py b/litellm/router.py index 8fea9126c85..3d059652094 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13104,6 +13104,7 @@ class Router: input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=typed_input, responses_api_request={"instructions": instructions} if instructions is not None else {}, + only_since_compaction=True, ) return litellm.token_counter( messages=cast(list, input_messages), # cast-ok: transformed chat messages diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index d14f236f765..7e1bce8e46b 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -55,6 +55,7 @@ def _provider_fields() -> dict[str, object]: "reasoningContentBlocks": [{"reasoningText": {"text": SECRET}}], "citations": [{"cited_text": SECRET}], "thought_signatures": [SIGNATURE], + "compaction_blocks": [{"type": "compaction", "content": SECRET}], } @@ -163,6 +164,26 @@ def test_streaming_redaction_redacts_delta_provider_fields_without_touching_the_ assert fields == _provider_fields() +@pytest.mark.parametrize("shape", ["object", "dict"]) +def test_redaction_hides_a_compaction_summary_in_responses_api_output(shape: str) -> None: + response: Final = ResponsesAPIResponse( + id="resp_1", + created_at=1, + model="claude-sonnet-5-5", + object="response", + status="completed", + output=[ + {"type": "compaction", "id": "cmp_1", "encrypted_content": json.dumps({"type": "compaction", "content": SECRET})} + ], + ) + result: Final = response if shape == "object" else response.model_dump() + + redacted: Final = perform_redaction({}, result) + + assert SECRET not in str(redacted), redacted + assert SECRET in str(result) + + @pytest.fixture(autouse=True) def _reset_global_redaction(): """Ensure the global setting is off for every test.""" diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index 6dbd19e991d..f5efaa4106b 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -4200,3 +4200,104 @@ class TestOpenAIResponsesHandlerStreamingScanKey: assert [inputs.get("texts") for inputs in guardrail.seen_inputs] == [["hi"]], guardrail.seen_inputs tool_calls = guardrail.seen_inputs[0].get("tool_calls") or [] assert [call["function"]["arguments"] for call in tool_calls] == ['{"city": "Paris"}'], tool_calls + + +class TestCompactionOutputSummaryReachesOutputGuardrails: + SUMMARY: Final = "The user lives at 1 Secret Lane." + + @classmethod + def _compaction_item(cls, summary: str | None = None) -> dict[str, str]: + return { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": summary or cls.SUMMARY}), + } + + @staticmethod + def _summary(item: object) -> str: + encrypted_content: Final = item["encrypted_content"] if isinstance(item, dict) else item.encrypted_content + return json.loads(encrypted_content)["content"] + + @staticmethod + def _lane_masking_guardrail() -> CustomGuardrail: + class MaskLane(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + return { + **inputs, + "texts": [text.replace("Secret Lane", "[MASKED]") for text in inputs.get("texts", [])], + } + + return MaskLane(guardrail_name="mask-lane") + + @pytest.mark.asyncio + async def test_non_streaming_summary_is_masked_inside_its_item(self) -> None: + response: Final = ResponsesAPIResponse( + id="resp_1", + created_at=1, + model="claude-sonnet-5-5", + object="response", + status="completed", + output=[self._compaction_item()], + ) + + result: Final = await OpenAIResponsesHandler().process_output_response(response, self._lane_masking_guardrail()) + + assert self._summary(result.output[0]) == "The user lives at 1 [MASKED]." + + @pytest.mark.asyncio + async def test_ended_stream_summary_is_masked_in_the_done_item_and_the_envelope(self) -> None: + events: Final = [ + { + "type": "response.output_item.added", + "output_index": 0, + "item": {**self._compaction_item(), "encrypted_content": ""}, + }, + {"type": "response.output_item.done", "output_index": 0, "item": self._compaction_item()}, + { + "type": "response.completed", + "response": { + "id": "resp_1", + "created_at": 1, + "model": "claude-sonnet-5-5", + "output": [self._compaction_item()], + "status": "completed", + }, + }, + ] + + await OpenAIResponsesHandler().process_output_streaming_response( + responses_so_far=events, + guardrail_to_apply=self._lane_masking_guardrail(), + deliver_ended_stream_rewrites=True, + ) + + assert self._summary(events[1]["item"]) == "The user lives at 1 [MASKED]." + assert self._summary(events[2]["response"]["output"][0]) == "The user lives at 1 [MASKED]." + + @pytest.mark.asyncio + async def test_a_blocking_guardrail_rejects_a_finished_compaction_item_mid_stream(self) -> None: + with pytest.raises(HTTPException) as exc_info: + await OpenAIResponsesHandler().process_output_streaming_response( + responses_so_far=[ + {"type": "response.output_item.done", "output_index": 0, "item": self._compaction_item()} + ], + guardrail_to_apply=MockGuardrail(guardrail_name="block"), + ) + + assert exc_info.value.detail["texts"] == [self.SUMMARY] + + def test_a_changed_summary_changes_the_ended_stream_scan_key(self) -> None: + handler: Final = OpenAIResponsesHandler() + + def ended_key(summary: str) -> object: + return handler.get_streaming_scan_key( + [{"type": "response.completed", "response": {"output": [self._compaction_item(summary)]}}] + ) + + assert ended_key(self.SUMMARY) != ended_key("The user lives elsewhere.") diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index bb2a89b1d98..9419ebbca4e 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -15,8 +15,9 @@ Redis. """ import asyncio +import json from datetime import datetime, timedelta -from typing import Any, Dict +from typing import Any, Dict, Final import pytest @@ -3727,3 +3728,34 @@ async def test_the_project_itpm_reservation_counts_the_request_off_the_event_loo if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) + + +_COMPACTION_SUMMARY: Final = "The user planned a Tokyo trip." +_NEXT_TURN: Final = {"type": "message", "role": "user", "content": "And Osaka?"} +_COMPACTED_INPUT: Final = [ + {"type": "message", "role": "user", "content": "Plan a Tokyo trip. " * 200}, + {"type": "message", "role": "assistant", "content": "Day one, Asakusa. " * 200}, + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": _COMPACTION_SUMMARY}), + }, + _NEXT_TURN, +] +_SUMMARY_INPUT: Final = [{"type": "message", "role": "user", "content": _COMPACTION_SUMMARY}, _NEXT_TURN] + + +def test_reservation_skips_the_history_a_compaction_item_replaces(rate_limiter): + handler, _cache = rate_limiter + compacted_request: Final = {"model": "gpt-4o", "input": _COMPACTED_INPUT, "max_output_tokens": 1} + + assert handler._estimate_input_and_output_tokens( + data=compacted_request, call_type="aresponses" + ) == handler._estimate_input_and_output_tokens( + data={**compacted_request, "input": _COMPACTED_INPUT[2:]}, call_type="aresponses" + ) + assert handler._estimate_precise_input_tokens( + compacted_request, model="gpt-4o", call_type="aresponses" + ) == handler._estimate_precise_input_tokens( + {**compacted_request, "input": _SUMMARY_INPUT}, model="gpt-4o", call_type="aresponses" + ) diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 1af42947da7..12cb57f5e02 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -3704,3 +3704,26 @@ def test_responses_routes_document_response_models_in_openapi_schema(): assert "output" in ok_200_properties("/v1/responses/{response_id}", "get") assert "deleted" in ok_200_properties("/v1/responses/{response_id}", "delete") assert "data" in ok_200_properties("/v1/responses/{response_id}/input_items", "get") + + +_COMPACTION_SUMMARY: Final = "The user planned a Tokyo trip." +_NEXT_TURN: Final = {"type": "message", "role": "user", "content": "And Osaka?"} +_COMPACTED_INPUT: Final = [ + {"type": "message", "role": "user", "content": "Plan a Tokyo trip. " * 200}, + {"type": "message", "role": "assistant", "content": "Day one, Asakusa. " * 200}, + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": _COMPACTION_SUMMARY}), + }, + _NEXT_TURN, +] +_SUMMARY_INPUT: Final = [{"type": "message", "role": "user", "content": _COMPACTION_SUMMARY}, _NEXT_TURN] + + +def test_input_token_count_reads_a_compaction_item_as_its_summary(): + from litellm.proxy.response_api_endpoints.endpoints import _responses_input_as_token_count_messages + + assert _responses_input_as_token_count_messages(_COMPACTED_INPUT, None) == ( + _responses_input_as_token_count_messages(_SUMMARY_INPUT, None) + ) diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index bf105da7793..72e539a1987 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -662,6 +662,8 @@ def test_streamed_responses_over_anthropic_emit_the_compaction_item_and_send_it_ (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, 0, "compaction"), (ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE, 1, "message"), ] + assert item_events[0].item.encrypted_content == "" + assert item_events[0].item.id == item_events[2].item.id assert [(item.type, getattr(item, "encrypted_content", None)) for item in completed.output[:1]] == [ ("compaction", item_events[2].item.encrypted_content) ] diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 2deeeaeac57..bd1dbaa86d9 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -6509,6 +6509,7 @@ def _anthropic_form_context_management(edit: Mapping[str, object]) -> ResponsesA _SUMMARY_AS_TEXT, ), ("gemini/gemini-3.8-flash", _COMPACTION_REQUESTED, _SUMMARY_AS_TEXT), + ("bedrock/claude_platform/claude-sonnet-5-5", _COMPACTION_REQUESTED, _SUMMARY_AS_TEXT), ], ) def test_compaction_input_item_replays_as_a_compaction_block_only_when_anthropic_compacts( diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index 002a595a5e7..406a33b5672 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -787,3 +787,58 @@ async def test_message_history_replays_real_key_named_tool_payloads() -> None: tool_message: Final = result["messages"][2] assert json.loads(assistant_message["tool_calls"][0]["function"]["arguments"]) == function_arguments assert json.loads(tool_message["content"]) == function_output + + +def _turn_spend_log(turn_input: list[dict[str, str]], answer: str) -> dict[str, object]: + return { + "request_id": f"resp_{len(turn_input)}", + "session_id": "session-compacted", + "proxy_server_request": json.dumps({"input": turn_input}), + "response": { + "id": f"chatcmpl-{len(turn_input)}", + "object": "chat.completion", + "choices": [{"index": 0, "message": {"role": "assistant", "content": answer}}], + }, + } + + +def _role_and_content(message: object) -> tuple[object, object]: + if isinstance(message, dict): + return message.get("role"), message.get("content") + return getattr(message, "role", None), getattr(message, "content", None) + + +@pytest.mark.asyncio +async def test_previous_response_history_starts_at_the_newest_turn_that_sent_a_compaction_item(): + first_turn: Final = [{"role": "user", "content": "Plan a Tokyo trip."}] + compaction_item: Final = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Turn one planned Tokyo."}), + } + second_turn: Final = [ + *first_turn, + {"role": "assistant", "content": "Day one, Asakusa."}, + compaction_item, + {"role": "user", "content": "And Osaka?"}, + ] + spend_logs: Final = [ + _turn_spend_log(first_turn, "Day one, Asakusa."), + _turn_spend_log(second_turn, "Dotonbori."), + ] + + with patch.object( + ResponsesSessionHandler, + "get_all_spend_logs_for_previous_response_id", + new_callable=AsyncMock, + return_value=spend_logs, + ): + session: Final = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + "resp_4" + ) + + assert [_role_and_content(message) for message in session["messages"]] == [ + ("user", "Turn one planned Tokyo."), + ("user", "And Osaka?"), + ("assistant", "Dotonbori."), + ] diff --git a/tests/unit/responses/test_responses_api_request_body.py b/tests/unit/responses/test_responses_api_request_body.py index 8c8b347be82..e06c88ba24e 100644 --- a/tests/unit/responses/test_responses_api_request_body.py +++ b/tests/unit/responses/test_responses_api_request_body.py @@ -1035,3 +1035,51 @@ async def test_a_native_responses_provider_places_every_point_itself(): assert user_only == [0, 2, 4, 6] assert with_unmatchable_system == user_only + + +_CLAUDE_COMPACTION_ITEM: Final = { + "type": "compaction", + "id": "cmp_claude", + "encrypted_content": json.dumps({"type": "compaction", "content": "The user's cat is named Biscuit."}), +} +_OPENAI_COMPACTION_ITEM: Final = {"type": "compaction", "id": "cmp_openai", "encrypted_content": "gAAAAABopaque"} +_HISTORY: Final = ( + {"role": "system", "content": "Answer briefly."}, + {"role": "user", "content": "My cat is named Biscuit."}, + {"role": "assistant", "content": "Noted."}, +) +_NEXT_TURN: Final = {"role": "user", "content": "What is my cat's name?"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("compaction_item", "sent_input"), + [ + ( + _CLAUDE_COMPACTION_ITEM, + [ + _HISTORY[0], + {"role": "user", "content": "The user's cat is named Biscuit.", "type": "message"}, + _NEXT_TURN, + ], + ), + (_OPENAI_COMPACTION_ITEM, [*_HISTORY, _OPENAI_COMPACTION_ITEM, _NEXT_TURN]), + ], +) +async def test_aresponses_sends_a_claude_compaction_item_to_openai_as_its_summary( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + compaction_item: dict[str, str], + sent_input: list[dict[str, str]], +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=_minimal_responses_api_payload("resp_compacted", "gpt-5.6")) + ) + + await litellm.aresponses( + model="openai/gpt-5.6", api_key="fake-api-key", input=[*_HISTORY, compaction_item, _NEXT_TURN] + ) + + assert json.loads(upstream.calls[0].request.read())["input"] == sent_input diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index 7cdac6c7581..204221148bb 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -2557,3 +2557,24 @@ def test_faked_stream_builds_each_refusal_summary_and_custom_tool_input_once_for (1, "content", 1): "No forecasts for Mars.", (2, "input", 0): "curl wttr.in/Berlin", } + + +_COMPACTION_SUMMARY: Final = "The user planned a Tokyo trip." +_NEXT_TURN: Final = {"type": "message", "role": "user", "content": "And Osaka?"} +_COMPACTED_INPUT: Final = [ + {"type": "message", "role": "user", "content": "Plan a Tokyo trip. " * 200}, + {"type": "message", "role": "assistant", "content": "Day one, Asakusa. " * 200}, + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": _COMPACTION_SUMMARY}), + }, + _NEXT_TURN, +] +_SUMMARY_INPUT: Final = [{"type": "message", "role": "user", "content": _COMPACTION_SUMMARY}, _NEXT_TURN] + + +def test_usage_estimate_counts_a_compaction_item_as_its_summary(): + assert _estimate_usage_from_text("gpt-4o-mini", _COMPACTED_INPUT, {}, "Dotonbori.") == _estimate_usage_from_text( + "gpt-4o-mini", _SUMMARY_INPUT, {}, "Dotonbori." + ) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 5a6da70d756..4c8e2ccdae8 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -6539,6 +6539,29 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): router._count_pre_call_check_tokens(messages=None, input=None) +def test_count_pre_call_check_tokens_skips_the_history_a_compaction_item_replaces(): + router = litellm.Router(model_list=[{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}]) + summary: Final = "The user planned a Tokyo trip." + compaction_item: Final = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": summary}), + } + history: Final = [ + {"role": "user", "content": "Plan a Tokyo trip. " * 200}, + {"role": "assistant", "content": "Day one, Asakusa. " * 200}, + ] + next_turn: Final = {"role": "user", "content": "And Osaka?"} + + compacted_tokens: Final = router._count_pre_call_check_tokens( + messages=None, input=[*history, compaction_item, next_turn] + ) + + assert compacted_tokens == router._count_pre_call_check_tokens( + messages=None, input=[{"role": "user", "content": summary}, next_turn] + ) + + def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): """ When neither messages nor input is provided (e.g. endpoints without prompt text), From 49f054757cfed481c2c5224fbdb1f70e3c2af330 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 18:37:32 -0700 Subject: [PATCH 05/21] fix(responses): start previous_response_id replay at a stored compaction and count reservations from it --- litellm/proxy/spend_tracking/input_tokens.py | 9 ++- .../compaction_items.py | 39 +++++++++-- .../session_handler.py | 57 ++++++++++++---- .../transformation.py | 42 ++++++++++-- .../proxy/spend_tracking/test_input_tokens.py | 29 +++++++++ .../test_litellm_completion_responses.py | 65 +++++++++++++++++++ .../test_session_handler.py | 41 +++++++++++- 7 files changed, 258 insertions(+), 24 deletions(-) diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 6c7083fb6db..520b5661438 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -19,6 +19,10 @@ from typing import Final import litellm from litellm._logging import verbose_proxy_logger +from litellm.responses.litellm_completion_transformation.compaction_items import ( + holds_compaction_summary, + input_since_newest_compaction, +) from litellm.rust_bridge import runtime from litellm.rust_bridge.catalog import Route, RouteContext from litellm.rust_bridge.token_counter import ( @@ -54,10 +58,11 @@ async def count_input_tokens( (model, rust_tokenizer(model)) for model in models ) groups: Final[tuple[RustTokenizer | None, ...]] = tuple(dict.fromkeys(tokenizer for _, tokenizer in tokenizers)) + counted_raw_body: Final = None if holds_compaction_summary(request_body.get("input")) else raw_body group_counts: Final = [ await _count_group( request_body=request_body, - raw_body=raw_body, + raw_body=counted_raw_body, tokenizer=tokenizer, models=tuple(model for model, selected in tokenizers if selected == tokenizer), ) @@ -140,7 +145,7 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: - return _count_text_tokens(model=model, text=request_body.get("input")) + return _count_text_tokens(model=model, text=input_since_newest_compaction(request_body.get("input"))) if "query" in request_body or "documents" in request_body: query_tokens: Final = _count_text_tokens(model=model, text=request_body.get("query")) document_tokens: Final = _count_text_tokens( diff --git a/litellm/responses/litellm_completion_transformation/compaction_items.py b/litellm/responses/litellm_completion_transformation/compaction_items.py index e79b064d690..746b6261a58 100644 --- a/litellm/responses/litellm_completion_transformation/compaction_items.py +++ b/litellm/responses/litellm_completion_transformation/compaction_items.py @@ -21,6 +21,8 @@ _INSTRUCTION_ROLES: Final = frozenset(("system", "developer")) def _json_object(value: object) -> Mapping[str, object] | None: + if not isinstance(value, Mapping): + return None try: return _JSON_OBJECT.validate_python(value) except ValidationError: @@ -50,12 +52,33 @@ def _newest_compaction_block(compaction_blocks: object) -> CompactionBlock | Non return next((block for value in reversed(blocks) if (block := _compaction_block(value)) is not None), None) -def compaction_output_item(provider_specific_fields: object) -> ResponseCompactionItem | None: +def _provider_compaction_block(provider_specific_fields: object) -> CompactionBlock | None: + fields: Final = _json_object(provider_specific_fields) + return None if fields is None else _newest_compaction_block(fields.get("compaction_blocks")) + + +def provider_compaction_summary(provider_specific_fields: object) -> str | None: + block: Final = _provider_compaction_block(provider_specific_fields) + return None if block is None else block.summary + + +def _choice_provider_specific_fields(choice: object) -> object: + choice_fields: Final = _json_object(choice) + message: Final = None if choice_fields is None else _json_object(choice_fields.get("message")) + return None if message is None else message.get("provider_specific_fields") + + +def response_holds_compaction_block(response: object) -> bool: + fields: Final = _json_object(response) try: - fields: Final = _JSON_OBJECT.validate_python(provider_specific_fields) + choices: Final = _OBJECT_LIST.validate_python(None if fields is None else fields.get("choices")) except ValidationError: - return None - block: Final = _newest_compaction_block(fields.get("compaction_blocks")) + return False + return any(provider_compaction_summary(_choice_provider_specific_fields(choice)) is not None for choice in choices) + + +def compaction_output_item(provider_specific_fields: object) -> ResponseCompactionItem | None: + block: Final = _provider_compaction_block(provider_specific_fields) if block is None: return None return ResponseCompactionItem( @@ -139,6 +162,14 @@ def _is_instruction_item(item: object) -> bool: return fields is not None and fields.get("role") in _INSTRUCTION_ROLES +def instruction_items(input: object) -> tuple[object, ...]: + try: + items: Final = _OBJECT_LIST.validate_python(input) + except ValidationError: + return () + return tuple(item for item in items if _is_instruction_item(item)) + + def _items_from_newest_compaction(input: Sequence[_Item]) -> tuple[_Item, ...]: newest: Final = next( (index for index in range(len(input) - 1, 0, -1) if compaction_summary(input[index]) is not None), 0 diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 24e848ceb8d..b0cb3d41cce 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -1,13 +1,18 @@ import asyncio import json from collections.abc import Mapping, Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler -from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary +from litellm.responses.litellm_completion_transformation.compaction_items import ( + holds_compaction_summary, + instruction_items, + response_holds_compaction_block, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( AllMessageValues, @@ -43,17 +48,47 @@ def _normalize_redacted_tool_call_arguments(message: Message) -> None: function_call.arguments = REDACTED_TOOL_CALL_ARGUMENTS_PLACEHOLDER -def _first_replayed_request(proxy_server_requests: Sequence[Mapping[str, object] | None]) -> int: +@dataclass(frozen=True, slots=True) +class _ReplayStart: + index: int + at_output: bool + + +def _compaction_point( + index: int, spend_log: "SpendLogsPayload", proxy_server_request: Mapping[str, object] | None +) -> _ReplayStart | None: + if response_holds_compaction_block(spend_log.get("response")): + return _ReplayStart(index=index, at_output=True) + if proxy_server_request is not None and holds_compaction_summary(proxy_server_request.get("input")): + return _ReplayStart(index=index, at_output=False) + return None + + +def _replay_start( + spend_logs: Sequence["SpendLogsPayload"], proxy_server_requests: Sequence[Mapping[str, object] | None] +) -> _ReplayStart: return next( ( - index - for index in range(len(proxy_server_requests) - 1, -1, -1) - if (request := proxy_server_requests[index]) is not None and holds_compaction_summary(request.get("input")) + start + for index in range(len(spend_logs) - 1, -1, -1) + if (start := _compaction_point(index, spend_logs[index], proxy_server_requests[index])) is not None ), - 0, + _ReplayStart(index=0, at_output=False), ) +def _replayed_request( + index: int, proxy_server_request: Mapping[str, object] | None, start: _ReplayStart +) -> Mapping[str, object] | None: + if proxy_server_request is None or not start.at_output or index != start.index: + return proxy_server_request + return { + **proxy_server_request, + "input": list(instruction_items(proxy_server_request.get("input"))), + "messages": None, + } + + def _stored_instructions(proxy_server_request: Mapping[str, object] | None) -> str | None: instructions: Final = None if proxy_server_request is None else proxy_server_request.get("instructions") return instructions if isinstance(instructions, str) and instructions else None @@ -93,14 +128,12 @@ class ResponsesSessionHandler: await ResponsesSessionHandler.get_proxy_server_request_from_spend_log(spend_log=spend_log) for spend_log in all_spend_logs ] - first_replayed: Final = _first_replayed_request(proxy_server_requests) - for spend_log, proxy_server_request_dict in zip( - all_spend_logs[first_replayed:], proxy_server_requests[first_replayed:] - ): + start: Final = _replay_start(all_spend_logs, proxy_server_requests) + for index in range(start.index, len(all_spend_logs)): chat_completion_message_history = ( ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload( - spend_log=spend_log, - proxy_server_request_dict=proxy_server_request_dict, + spend_log=all_spend_logs[index], + proxy_server_request_dict=_replayed_request(index, proxy_server_requests[index], start), chat_completion_message_history=chat_completion_message_history, item_cache_control=item_cache_control, ) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index b8075d7c078..8e483b6922d 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -120,6 +120,7 @@ from .compaction_items import ( compaction_output_item, decode_compaction_item, holds_compaction_summary, + provider_compaction_summary, provider_takes_compaction_blocks, requests_compaction, without_items_before_compaction, @@ -310,6 +311,26 @@ def _without_stored_tool_search_uses(message: SessionMessage, search_call_ids: f return message +def _stored_compaction_summary(message: SessionMessage) -> str | None: + return provider_compaction_summary(message.provider_specific_fields) if isinstance(message, Message) else None + + +def _with_stored_compaction_as_summary(message: SessionMessage) -> tuple[SessionMessage, ...]: + if not isinstance(message, Message): + return (message,) + summary: Final = provider_compaction_summary(message.provider_specific_fields) + if summary is None: + return (message,) + other_fields: Final = { + key: value + for key, value in _STR_KEY_DICT_ADAPTER.validate_python(message.provider_specific_fields).items() + if key != "compaction_blocks" + } + answer: Final = message.model_copy(update={"provider_specific_fields": other_fields or None}) + summary_message: Final = ChatCompletionUserMessage(role="user", content=summary) + return (summary_message, answer) if answer.content or answer.tool_calls else (summary_message,) + + def _input_item_cache_control(input_item: Mapping[str, object]) -> ChatCompletionCachedContent | None: cache_control: Final = input_item.get("cache_control") if not cache_control: @@ -1035,6 +1056,18 @@ class LiteLLMCompletionResponsesConfig: if (kept := _without_stored_tool_search_uses(message, search_call_ids)) is not None ) + @staticmethod + def _with_compaction_blocks_the_route_takes( + messages: Sequence[SessionMessage], route: CompletionRoute, context_management: object + ) -> tuple[SessionMessage, ...]: + if not any(_stored_compaction_summary(message) is not None for message in messages): + return tuple(messages) + if requests_compaction(context_management) and provider_takes_compaction_blocks( + route.model, route.custom_llm_provider + ): + return tuple(messages) + return tuple(chain.from_iterable(map(_with_stored_compaction_as_summary, messages))) + @staticmethod async def async_responses_api_session_handler( previous_response_id: str, @@ -1063,11 +1096,12 @@ class LiteLLMCompletionResponsesConfig: else ChatCompletionSession(messages=[], litellm_session_id=None) ) _messages: Final = litellm_completion_request.get("messages") or [] - stored_messages: Final = chat_completion_session.get("messages") or [] + route: Final = CompletionRoute(model=model, custom_llm_provider=provider) + stored_messages: Final = LiteLLMCompletionResponsesConfig._with_compaction_blocks_the_route_takes( + chat_completion_session.get("messages") or [], route, request_fields.get("context_management") + ) replayed_messages: Final = ( - LiteLLMCompletionResponsesConfig._without_tool_search_uses_off_anthropic( - stored_messages, CompletionRoute.model_validate(litellm_completion_request) - ) + LiteLLMCompletionResponsesConfig._without_tool_search_uses_off_anthropic(stored_messages, route) if any(holds_tool_search_use(message) for message in stored_messages) else stored_messages ) diff --git a/tests/unit/proxy/spend_tracking/test_input_tokens.py b/tests/unit/proxy/spend_tracking/test_input_tokens.py index a1da6c79cd7..bfec1c10212 100644 --- a/tests/unit/proxy/spend_tracking/test_input_tokens.py +++ b/tests/unit/proxy/spend_tracking/test_input_tokens.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json from types import MappingProxyType from typing import Final @@ -27,3 +28,31 @@ async def test_large_input_is_still_counted() -> None: assert counts[CL100K_MODEL] == count_input_tokens_for_model(request_body=request_body, model=CL100K_MODEL) assert isinstance(counts, MappingProxyType) + + +_COMPACTION_ITEM: Final = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Turn one planned Tokyo."}), +} +_TURN_AFTER_COMPACTION: Final = {"role": "user", "content": "And Osaka?"} + + +@pytest.mark.parametrize("model", [CL100K_MODEL, "claude-sonnet-5-5"]) +@pytest.mark.asyncio +async def test_history_before_a_compaction_item_is_not_counted(model: str) -> None: + compacted_body: Final = {"model": model, "input": [_COMPACTION_ITEM, _TURN_AFTER_COMPACTION]} + full_body: Final = { + "model": model, + "input": [{"role": "user", "content": "Plan a Tokyo trip. " * 500}, _COMPACTION_ITEM, _TURN_AFTER_COMPACTION], + } + + full_counts: Final = await count_input_tokens( + request_body=full_body, raw_body=json.dumps(full_body).encode(), models=(model,) + ) + compacted_counts: Final = await count_input_tokens( + request_body=compacted_body, raw_body=json.dumps(compacted_body).encode(), models=(model,) + ) + + assert full_counts[model] == compacted_counts[model] + assert count_input_tokens_for_model(request_body=full_body, model=model) == compacted_counts[model] diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index bd1dbaa86d9..228da2aa441 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -6534,6 +6534,71 @@ def test_compaction_input_item_replays_as_a_compaction_block_only_when_anthropic ] +def _replayed_compacted_answer(model: str, context_management: object) -> list[tuple[object, object, object]]: + async def session_lookup(previous_response_id: str) -> ChatCompletionSession: + return ChatCompletionSession( + messages=[ + Message( + role="assistant", + content="Day one, Asakusa.", + provider_specific_fields={"compaction_blocks": [_COMPACTION_BLOCK]}, + ) + ], + litellm_session_id=None, + ) + + request: Final = asyncio.run( + LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( + previous_response_id="resp_compacted", + litellm_completion_request={ + "model": model, + "messages": [{"role": "user", "content": "And Osaka?"}], + **({"context_management": context_management} if context_management else {}), + }, + instructions=None, + session_lookup=session_lookup, + ) + ) + return [ + ( + message.get("role"), + message.get("content"), + (message.get("provider_specific_fields") or {}).get("compaction_blocks"), + ) + for message in request["messages"] + ] + + +_ANSWER_AFTER_SUMMARY: Final = [ + ("user", _COMPACTION_BLOCK["content"], None), + ("assistant", "Day one, Asakusa.", None), + ("user", "And Osaka?", None), +] + + +@pytest.mark.parametrize( + ("model", "context_management", "replayed"), + [ + ( + "anthropic/claude-sonnet-5-5", + _COMPACTION_REQUESTED["context_management"], + [("assistant", "Day one, Asakusa.", [_COMPACTION_BLOCK]), ("user", "And Osaka?", None)], + ), + ("anthropic/claude-sonnet-5-5", None, _ANSWER_AFTER_SUMMARY), + ( + "bedrock/claude_platform/claude-sonnet-5-5", + _COMPACTION_REQUESTED["context_management"], + _ANSWER_AFTER_SUMMARY, + ), + ("gpt-5.6", _COMPACTION_REQUESTED["context_management"], _ANSWER_AFTER_SUMMARY), + ], +) +def test_stored_compacted_answer_replays_its_block_only_when_anthropic_compacts( + model: str, context_management: object, replayed: list[tuple[object, object, object]] +) -> None: + assert _replayed_compacted_answer(model, context_management) == replayed + + def test_summary_replay_drops_the_history_its_compaction_item_replaces() -> None: result: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( model="gemini/gemini-3.8-flash", diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index 406a33b5672..c585cdeb239 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -789,7 +789,12 @@ async def test_message_history_replays_real_key_named_tool_payloads() -> None: assert json.loads(tool_message["content"]) == function_output -def _turn_spend_log(turn_input: list[dict[str, str]], answer: str) -> dict[str, object]: +def _turn_spend_log( + turn_input: list[dict[str, str]], answer: str, compaction_blocks: list[dict[str, str]] | None = None +) -> dict[str, object]: + compaction_fields: Final = ( + {"provider_specific_fields": {"compaction_blocks": compaction_blocks}} if compaction_blocks else {} + ) return { "request_id": f"resp_{len(turn_input)}", "session_id": "session-compacted", @@ -797,7 +802,7 @@ def _turn_spend_log(turn_input: list[dict[str, str]], answer: str) -> dict[str, "response": { "id": f"chatcmpl-{len(turn_input)}", "object": "chat.completion", - "choices": [{"index": 0, "message": {"role": "assistant", "content": answer}}], + "choices": [{"index": 0, "message": {"role": "assistant", "content": answer, **compaction_fields}}], }, } @@ -842,3 +847,35 @@ async def test_previous_response_history_starts_at_the_newest_turn_that_sent_a_c ("user", "And Osaka?"), ("assistant", "Dotonbori."), ] + + +@pytest.mark.asyncio +async def test_previous_response_history_starts_at_the_newest_answer_that_compacted(): + compaction_blocks: Final = [{"type": "compaction", "content": "Turn one planned Tokyo."}] + spend_logs: Final = [ + _turn_spend_log([{"role": "user", "content": "Hi."}], "Hello."), + _turn_spend_log( + [{"role": "developer", "content": "Answer briefly."}, {"role": "user", "content": "Plan a Tokyo trip."}], + "Day one, Asakusa.", + compaction_blocks, + ), + _turn_spend_log([{"role": "user", "content": "And Osaka?"}], "Dotonbori."), + ] + + with patch.object( + ResponsesSessionHandler, + "get_all_spend_logs_for_previous_response_id", + new_callable=AsyncMock, + return_value=spend_logs, + ): + session: Final = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + "resp_1" + ) + + assert [_role_and_content(message) for message in session["messages"]] == [ + ("developer", "Answer briefly."), + ("assistant", "Day one, Asakusa."), + ("user", "And Osaka?"), + ("assistant", "Dotonbori."), + ] + assert getattr(session["messages"][1], "provider_specific_fields", None) == {"compaction_blocks": compaction_blocks} From 52774f08a5a701ddb83cc7d319fb6f8346076a46 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 19:59:19 -0700 Subject: [PATCH 06/21] fix(responses): keep a stored compaction summary whole in spend logs --- .../spend_tracking/spend_tracking_utils.py | 19 ++++++- .../test_spend_tracking_utils.py | 57 +++++++++++++++++++ 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 5b02d6a6afd..b6b0b376bfe 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -63,6 +63,7 @@ from litellm.proxy._types import ( from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.proxy.utils import PrismaClient, hash_token +from litellm.responses.litellm_completion_transformation.compaction_items import decode_compaction_item from litellm.types.router import DeploymentTypedDict, LiteLLM_Params from litellm.types.utils import ( BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, @@ -1159,11 +1160,21 @@ def _is_spend_log_content(parent: Mapping[str, object], key: str) -> bool: ) +def _is_compaction_summary(parent: Mapping[str, object], key: str, keep_compaction_blocks: bool) -> bool: + if parent.get("type") != "compaction": + return False + value: Final = parent.get(key) + if key == "content": + return keep_compaction_blocks and isinstance(value, str) + return key == "encrypted_content" and decode_compaction_item(value) is not None + + def _sanitize_request_body_for_spend_logs_payload( request_body: Mapping[str, object], visited: set | None = None, max_string_length_prompt_in_db: int | None = None, mask_credentials: bool = True, + keep_compaction_blocks: bool = False, ) -> dict: """ Recursively sanitize request body to prevent logging large base64 strings or other large values. @@ -1193,7 +1204,7 @@ def _sanitize_request_body_for_spend_logs_payload( def _sanitize_value(value: object, mask_credentials: bool) -> object: if isinstance(value, Mapping): return _sanitize_request_body_for_spend_logs_payload( - value, visited, max_string_length_prompt_in_db, mask_credentials + value, visited, max_string_length_prompt_in_db, mask_credentials, keep_compaction_blocks ) elif isinstance(value, list): return [_sanitize_value(item, mask_credentials) for item in value] @@ -1234,6 +1245,8 @@ def _sanitize_request_body_for_spend_logs_payload( return { k: REDACTED_BY_LITELM_STRING if mask_credentials and _is_request_body_credential(k, v) + else v + if _is_compaction_summary(request_body, k, keep_compaction_blocks) else _sanitize_value(v, mask_credentials and not _is_spend_log_content(request_body, k)) for k, v in request_body.items() if k not in _SENSITIVE_REQUEST_BODY_KEYS @@ -1871,7 +1884,9 @@ def _get_response_for_spend_logs_payload( response_obj = _convert_to_json_serializable_dict(response_obj) response_obj = perform_redaction(model_call_details={}, result=response_obj) - sanitized_wrapper: Final = _sanitize_request_body_for_spend_logs_payload({"response": response_obj}) + sanitized_wrapper: Final = _sanitize_request_body_for_spend_logs_payload( + {"response": response_obj}, keep_compaction_blocks=True + ) sanitized_response: Final = sanitized_wrapper.get("response", response_obj) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index f5f690c4bfd..8badf606c4b 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -979,6 +979,63 @@ def test_get_response_for_spend_logs_payload_truncates_large_base64(mock_should_ assert parsed["data"][0]["other_field"] == "value" +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") +def test_get_response_for_spend_logs_payload_keeps_a_long_compaction_summary_whole(mock_should_store): + from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB + + mock_should_store.return_value = True + filler = "detail " * (MAX_STRING_LENGTH_PROMPT_IN_DB // 7) + summary = f"{filler}The cat is named Biscuit. {filler}" + answer = f"{filler}OK{filler}" + payload = cast( + StandardLoggingPayload, + { + "response": { + "choices": [ + { + "message": { + "role": "assistant", + "content": answer, + "provider_specific_fields": { + "compaction_blocks": [{"type": "compaction", "content": summary}] + }, + } + } + ] + } + }, + ) + + message = json.loads(_get_response_for_spend_logs_payload(payload))["choices"][0]["message"] + assert message["provider_specific_fields"]["compaction_blocks"] == [{"type": "compaction", "content": summary}] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in message["content"] + + +def test_sanitize_request_body_keeps_a_long_claude_compaction_item_whole(): + from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB + + filler = "detail " * (MAX_STRING_LENGTH_PROMPT_IN_DB // 7) + claude_item = json.dumps({"type": "compaction", "content": f"{filler}The cat is named Biscuit. {filler}"}) + opaque_item = "gAAAA" + "x" * (MAX_STRING_LENGTH_PROMPT_IN_DB * 2) + + sanitized = _sanitize_request_body_for_spend_logs_payload( + { + "input": [ + {"type": "compaction", "id": "cmp_1", "encrypted_content": claude_item}, + {"type": "compaction", "id": "cmp_2", "encrypted_content": opaque_item}, + {"role": "assistant", "content": [{"type": "compaction", "content": f"{filler}Biscuit{filler}"}]}, + {"role": "user", "content": f"{filler}What is my cat called?{filler}"}, + ] + } + ) + + claude_stored, opaque_stored, messages_block_stored, message_stored = sanitized["input"] + assert claude_stored["encrypted_content"] == claude_item + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in opaque_stored["encrypted_content"] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in messages_block_stored["content"][0]["content"] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in message_stored["content"] + + @patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_response_for_spend_logs_payload_strips_null_bytes(mock_should_store): """Regression for PostgreSQL 22P05: NUL bytes must be stripped from response.""" From 51a46a1db210b51fab73b877073fa9cbca53ba7e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 22:03:46 -0700 Subject: [PATCH 07/21] fix(responses): keep stored compaction items summary-only, redact compaction_delta, and finish compaction before reasoning A spend-log row now stores a top-level compaction input item as its summary alone, so extra keys in a client-built item no longer reach the database whole. Message redaction also nulls the streamed compaction_delta, and the streaming bridge sends the compaction item's done event before the reasoning item that follows it. Adds the audit's integration tests for compaction items on the wire. --- litellm/litellm_core_utils/redact_messages.py | 8 +- .../spend_tracking/spend_tracking_utils.py | 49 +- .../compaction_items.py | 7 + .../streaming_iterator.py | 1 + .../test_responses_compaction_items_wire.py | 822 ++++++++++++++++++ .../test_redact_messages.py | 1 + .../test_spend_tracking_utils.py | 38 +- .../test_anthropic_responses_bridge.py | 66 ++ 8 files changed, 967 insertions(+), 25 deletions(-) create mode 100644 tests/integration/providers/test_responses_compaction_items_wire.py diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 7247ba2e4be..3b50b9af24c 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -150,7 +150,13 @@ def _redacted_provider_specific_fields(fields: Mapping[str, object]) -> Mapping[ **({"reasoning_content": REDACTED_BY_LITELLM} if fields.get("reasoning_content") is not None else {}), **{ key: None - for key in ("thinking_blocks", "reasoningContentBlocks", "citations", "compaction_blocks") + for key in ( + "thinking_blocks", + "reasoningContentBlocks", + "citations", + "compaction_blocks", + "compaction_delta", + ) if key in fields }, } diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index b6b0b376bfe..940574988dc 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -63,7 +63,7 @@ from litellm.proxy._types import ( from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.proxy.utils import PrismaClient, hash_token -from litellm.responses.litellm_completion_transformation.compaction_items import decode_compaction_item +from litellm.responses.litellm_completion_transformation.compaction_items import summary_only_compaction_item from litellm.types.router import DeploymentTypedDict, LiteLLM_Params from litellm.types.utils import ( BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, @@ -1160,13 +1160,13 @@ def _is_spend_log_content(parent: Mapping[str, object], key: str) -> bool: ) -def _is_compaction_summary(parent: Mapping[str, object], key: str, keep_compaction_blocks: bool) -> bool: - if parent.get("type") != "compaction": - return False - value: Final = parent.get(key) - if key == "content": - return keep_compaction_blocks and isinstance(value, str) - return key == "encrypted_content" and decode_compaction_item(value) is not None +def _is_compaction_block_summary(parent: Mapping[str, object], key: str, keep_compaction_blocks: bool) -> bool: + return ( + keep_compaction_blocks + and key == "content" + and parent.get("type") == "compaction" + and isinstance(parent.get(key), str) + ) def _sanitize_request_body_for_spend_logs_payload( @@ -1246,7 +1246,7 @@ def _sanitize_request_body_for_spend_logs_payload( k: REDACTED_BY_LITELM_STRING if mask_credentials and _is_request_body_credential(k, v) else v - if _is_compaction_summary(request_body, k, keep_compaction_blocks) + if _is_compaction_block_summary(request_body, k, keep_compaction_blocks) else _sanitize_value(v, mask_credentials and not _is_spend_log_content(request_body, k)) for k, v in request_body.items() if k not in _SENSITIVE_REQUEST_BODY_KEYS @@ -1765,6 +1765,34 @@ def _stored_request_without_requester_ip_address(request_body: Mapping[object, o } +def _with_whole_compaction_summary(item: object, stored_item: object) -> object: + summary_only: Final = summary_only_compaction_item(item) + if summary_only is None: + return stored_item + try: + stored_fields: Final = _OBJECT_MAPPING.validate_python(stored_item) + except ValidationError: + return stored_item + return {**stored_fields, "encrypted_content": summary_only} + + +def _sanitize_proxy_server_request_body(request_body: Mapping[str, object]) -> Mapping[object, object]: + sanitized: Final = _OBJECT_MAPPING.validate_python(_sanitize_request_body_for_spend_logs_payload(request_body)) + items: Final = request_body.get("input") + stored_items: Final = sanitized.get("input") + if not isinstance(items, list) or not isinstance(stored_items, list): + return sanitized + return { + **sanitized, + "input": [ + _with_whole_compaction_summary(item, stored_item) + for item, stored_item in zip( + _OBJECT_LIST.validate_python(items), _OBJECT_LIST.validate_python(stored_items) + ) + ], + } + + def _get_proxy_server_request_for_spend_logs_payload( metadata: dict, litellm_params: Mapping[str, object], @@ -1809,8 +1837,7 @@ def _get_proxy_server_request_for_spend_logs_payload( _request_body = _convert_mapping_to_json_serializable(without_classifier_audit(_request_body)) perform_redaction(model_call_details=_request_body, result=None) - _request_body = _sanitize_request_body_for_spend_logs_payload(_request_body) - stored_request_body: Final = _OBJECT_MAPPING.validate_python(_request_body) + stored_request_body: Final = _sanitize_proxy_server_request_body(_request_body) _request_body_json_str: Final = safe_dumps( _stored_request_without_requester_ip_address(stored_request_body) if drop_requester_ip_address diff --git a/litellm/responses/litellm_completion_transformation/compaction_items.py b/litellm/responses/litellm_completion_transformation/compaction_items.py index 746b6261a58..a99b43c5295 100644 --- a/litellm/responses/litellm_completion_transformation/compaction_items.py +++ b/litellm/responses/litellm_completion_transformation/compaction_items.py @@ -140,6 +140,13 @@ def compaction_summary(item: object) -> str | None: return None if block is None else block.summary +def summary_only_compaction_item(item: object) -> str | None: + summary: Final = compaction_summary(item) + if summary is None: + return None + return json.dumps({"type": "compaction", "content": summary}, separators=(",", ":")) + + def with_compaction_summary(encrypted_content: str, summary: str) -> str: original: Final = _without_affinity_wrap(encrypted_content) block: Final = decode_compaction_item(original) diff --git a/litellm/responses/litellm_completion_transformation/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 6918c2634ce..b5dddb77fb9 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -1170,6 +1170,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): or mint_reasoning_item_id() ) reasoning_item_id = self._cached_reasoning_item_id + self._queue_compaction_item_done() # Create text.done event first with its own sequence number self._sequence_number += 1 diff --git a/tests/integration/providers/test_responses_compaction_items_wire.py b/tests/integration/providers/test_responses_compaction_items_wire.py new file mode 100644 index 00000000000..bd4f56da124 --- /dev/null +++ b/tests/integration/providers/test_responses_compaction_items_wire.py @@ -0,0 +1,822 @@ +import json +import os +import re +import uuid +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Final, Literal, TypeAlias +from urllib.parse import urlsplit + +import httpx +import openai +import pytest +import yaml +from integration._support import claude_code as cc +from integration._support.client import ( + Gateway, + Scenario, + eventually, + gateway_from_environment, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from integration._support.worker_readiness import wait_for_worker_readiness +from pydantic import JsonValue, TypeAdapter + +from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD +from litellm.responses.utils import ResponsesAPIRequestUtils + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_CLAUDE: Final = cc.OPUS +_GEMINI: Final = "gemini-3.8-flash" +_OPENAI: Final = "gpt-5.6" +_PROXY_WORKERS: Final = max(1, int(os.environ.get("INTEGRATION_PROXY_WORKERS", "1"))) +_CACHE_BUST: Final[Mapping[str, JsonValue]] = {"cache": {"no-cache": True}} +_COMPACTING: Final[list[JsonValue]] = [{"type": "compaction", "compact_threshold": 150000}] +_COMPACT_BETA: Final = "compact-2026-01-12" +_CITY: Final = "Tokyo" +_MASKED_CITY: Final = "[CITY]" +_MASK_INPUT: Final = "compaction-mask-input" +_MASK_OUTPUT: Final = "compaction-mask-output" +_BLOCK_OUTPUT: Final = "compaction-block-output" +_REDACT_HEADER: Final = {"x-litellm-enable-message-redaction": "true"} +_ASK: Final = re.compile(r"ask-([0-9a-f]{32})") + +Route: TypeAlias = Literal["claude-compacting", "claude-plain", "gemini", "openai", "openai-compact", "openai-opaque"] + + +def _new_marker() -> str: + return uuid.uuid4().hex + + +def _ask(marker: str, *, compact: bool = False) -> str: + return f"ask-{marker} compact-now-{marker}" if compact else f"ask-{marker}" + + +def _answer(marker: str) -> str: + return f"answer-{marker}" + + +def _summary(marker: str) -> str: + return f"kept-{marker} {'a' * 1500} middle-{marker} {'b' * 1500} the user lives in {_CITY}" + + +def _block(summary: str) -> dict[str, JsonValue]: + return {"type": "compaction", "content": summary} + + +def _compaction_item(summary: str) -> dict[str, JsonValue]: + return { + "type": "compaction", + "id": f"cmp_{uuid.uuid4().hex}", + "encrypted_content": json.dumps(_block(summary), separators=(",", ":")), + } + + +def _newest_marker(raw: bytes) -> str: + found: Final = _ASK.findall(raw.decode()) + assert found, raw[:400] + return str(found[-1]) + + +def _anthropic_usage() -> dict[str, JsonValue]: + return {"input_tokens": 20, "output_tokens": 5} + + +def _anthropic_block_frames(index: int, block: Mapping[str, JsonValue]) -> tuple[bytes, ...]: + def start(content_block: Mapping[str, JsonValue]) -> bytes: + return cc.sse_frame( + "content_block_start", {"type": "content_block_start", "index": index, "content_block": dict(content_block)} + ) + + def delta(fields: Mapping[str, JsonValue]) -> bytes: + return cc.sse_frame( + "content_block_delta", {"type": "content_block_delta", "index": index, "delta": dict(fields)} + ) + + stop: Final = cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": index}) + if block["type"] == "compaction": + summary: Final = string_value(block["content"]) + half: Final = len(summary) // 2 + return ( + start({"type": "compaction", "content": ""}), + delta({"type": "compaction_delta", "content": summary[:half]}), + delta({"type": "compaction_delta", "content": summary[half:]}), + stop, + ) + return (start({"type": "text", "text": ""}), delta({"type": "text_delta", "text": block["text"]}), stop) + + +def _anthropic_stream(identity: str, content: Sequence[Mapping[str, JsonValue]]) -> tuple[bytes, ...]: + start: Final = cc.sse_frame( + "message_start", + { + "type": "message_start", + "message": { + "id": identity, + "type": "message", + "role": "assistant", + "model": _CLAUDE, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 20, "output_tokens": 1}, + }, + }, + ) + blocks: Final = tuple( + frame for index, block in enumerate(content) for frame in _anthropic_block_frames(index, block) + ) + end: Final = cc.sse_frame( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ) + return (start, *blocks, end, cc.sse_frame("message_stop", {"type": "message_stop"})) + + +def _claude_reply(body: Mapping[str, JsonValue], raw: bytes) -> Reply: + marker: Final = _newest_marker(raw) + compaction: Final = (_block(_summary(marker)),) if f"compact-now-{marker}".encode() in raw else () + content: Final = (*compaction, {"type": "text", "text": _answer(marker)}) + identity: Final = f"msg_{uuid.uuid4().hex}" + if body.get("stream") is True: + return Reply(content_type="text/event-stream", chunks=_anthropic_stream(identity, content)) + return Reply(body=cc.message_reply(identity, _CLAUDE, content, _anthropic_usage())) + + +def _gemini_reply(raw: bytes) -> Reply: + frame: Final = { + "responseId": f"gemini-{uuid.uuid4().hex}", + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": _answer(_newest_marker(raw))}]}, + "index": 0, + "finishReason": "STOP", + } + ], + "usageMetadata": {"promptTokenCount": 20, "candidatesTokenCount": 5, "totalTokenCount": 25}, + "modelVersion": _GEMINI, + } + return Reply(body=json.dumps(frame).encode()) + + +_OPENAI_USAGE: Final = { + "input_tokens": 20, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 5, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 25, +} + + +def _openai_reply(raw: bytes) -> Reply: + tag: Final = uuid.uuid4().hex + message: Final = { + "id": f"msg_{tag}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": _answer(_newest_marker(raw)), "annotations": []}], + } + response: Final = { + "id": f"resp_{tag}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": _OPENAI, + "output": [message], + "usage": _OPENAI_USAGE, + } + return Reply(body=json.dumps(response).encode()) + + +def _openai_compact_reply() -> Reply: + compacted: Final = { + "id": f"resp_compact_{uuid.uuid4().hex}", + "created_at": 1, + "object": "response.compaction", + "output": [], + "usage": _OPENAI_USAGE, + } + return Reply(body=json.dumps(compacted).encode()) + + +def _vendor(request: Request) -> Reply: + if request.method == "GET": + return Reply(body=json.dumps({"object": "list", "data": [{"id": _OPENAI, "object": "model"}]}).encode()) + path: Final = urlsplit(request.target).path + if path.endswith("/v1/messages/count_tokens"): + return Reply(body=json.dumps({"input_tokens": _word_count(_JSON_OBJECT.validate_json(request.body))}).encode()) + if path.endswith("/v1/messages"): + return _claude_reply(_JSON_OBJECT.validate_json(request.body), request.body) + if path.endswith(":generateContent"): + return _gemini_reply(request.body) + if path.endswith("/responses/compact"): + return _openai_compact_reply() + assert path.endswith("/responses"), request.target + return _openai_reply(request.body) + + +def _provider_calls(wire: Wire, marker: str) -> tuple[Request, ...]: + needle: Final = f"ask-{marker}".encode() + return tuple(request for request in wire.drain() if request.method == "POST" and needle in request.body) + + +def _only_call(wire: Wire, marker: str) -> tuple[Request, dict[str, JsonValue]]: + calls: Final = _provider_calls(wire, marker) + assert len(calls) == 1, [call.target for call in calls] + return calls[0], _JSON_OBJECT.validate_json(calls[0].body) + + +def _objects(value: JsonValue) -> tuple[dict[str, JsonValue], ...]: + assert isinstance(value, list), value + return tuple(object_value(item) for item in value) + + +def _message_texts(message: Mapping[str, JsonValue], key: str = "content") -> tuple[str, ...]: + content: Final = message.get(key) + if isinstance(content, str): + return (content,) + if not isinstance(content, list): + return () + return tuple(string_value(part["text"]) for part in _objects(content) if isinstance(part.get("text"), str)) + + +def _anthropic_texts(body: Mapping[str, JsonValue], role: str) -> tuple[str, ...]: + messages: Final = tuple(message for message in _objects(body["messages"]) if message["role"] == role) + return tuple(text for message in messages for text in _message_texts(message)) + + +def _word_count(body: Mapping[str, JsonValue]) -> int: + texts: Final = (*_anthropic_texts(body, "user"), *_anthropic_texts(body, "assistant")) + return sum(len(text.split()) for text in texts) + + +def _anthropic_block_types(body: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]: + contents: Final = tuple(message["content"] for message in _objects(body["messages"])) + return tuple(block.get("type") for content in contents if isinstance(content, list) for block in _objects(content)) + + +def _gemini_texts(body: Mapping[str, JsonValue], role: str) -> tuple[str, ...]: + turns: Final = tuple(turn for turn in _objects(body["contents"]) if turn["role"] == role) + return tuple(text for turn in turns for text in _message_texts(turn, "parts")) + + +def _openai_user_texts(body: Mapping[str, JsonValue]) -> tuple[str, ...]: + items: Final = tuple(item for item in _objects(body["input"]) if item.get("role") == "user") + return tuple(text for item in items for text in _message_texts(item)) + + +def _openai_item_types(body: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]: + return tuple(item.get("type") for item in _objects(body["input"])) + + +def _register(scenario: Scenario, wire: Wire, provider: Literal["claude", "gemini", "openai"]) -> str: + match provider: + case "claude": + model = scenario.model(model=f"anthropic/{_CLAUDE}", api_base=wire.url, api_key=cc.ANTHROPIC_API_KEY) + case "gemini": + model = scenario.model(model=f"gemini/{_GEMINI}", api_base=wire.url, api_key="synthetic-gemini-key") + case "openai": + model = scenario.model(model=f"openai/{_OPENAI}", api_base=f"{wire.url}/v1", api_key="synthetic-openai-key") + wait_for_worker_readiness(scenario.gateway, model, 1, _PROXY_WORKERS) + return model + + +def _openai_client(gateway: Gateway) -> openai.OpenAI: + return openai.OpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=60), + ) + + +def _decoded(encrypted_content: str) -> dict[str, JsonValue]: + return _JSON_OBJECT.validate_json(encrypted_content) + + +def _responses( + gateway: Gateway, body: Mapping[str, JsonValue], headers: Mapping[str, str] | None = None +) -> dict[str, JsonValue]: + response: Final = gateway.request("POST", "/v1/responses", {**body, **_CACHE_BUST}, headers=headers) + assert response.status_code == 200, response.text + return _JSON_OBJECT.validate_json(response.content) + + +def _sse_events(text: str) -> tuple[dict[str, JsonValue], ...]: + payloads: Final = tuple( + line.removeprefix("data:").strip() for line in text.splitlines() if line.startswith("data:") + ) + return tuple(_JSON_OBJECT.validate_json(payload) for payload in payloads if payload != "[DONE]") + + +def _stream(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + headers: Final = {"Authorization": f"Bearer {gateway.key}"} + with gateway.client.stream("POST", path, json={**body, **_CACHE_BUST, "stream": True}, headers=headers) as response: + text: Final = response.read().decode() + assert response.status_code == 200, text + return _sse_events(text) + + +def _output_items(response: Mapping[str, JsonValue], item_type: str) -> tuple[dict[str, JsonValue], ...]: + return tuple(item for item in _objects(response["output"]) if item["type"] == item_type) + + +def _listed(value: JsonValue) -> list[JsonValue]: + assert isinstance(value, list), value + return value + + +def _only(items: Sequence[dict[str, JsonValue]]) -> dict[str, JsonValue]: + assert len(items) == 1, items + return items[0] + + +@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream")) +def test_claude_compaction_block_is_returned_as_a_compaction_output_item(gateway: Gateway, stream: bool) -> None: + marker: Final = _new_marker() + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, "claude") + client: Final = _openai_client(gateway) + request: Final = { + "input": _ask(marker, compact=True), + "extra_body": {"context_management": _COMPACTING, **_CACHE_BUST}, + } + if not stream: + reply: Final = client.responses.create(model=model, **request) + assert [item.type for item in reply.output] == ["compaction", "message"] + compaction: Final = reply.output[0] + assert compaction.type == "compaction" + assert _decoded(compaction.encrypted_content) == _block(_summary(marker)) + assert reply.output_text == _answer(marker) + else: + events: Final = tuple(client.responses.create(model=model, stream=True, **request)) + added: Final = [e for e in events if e.type == "response.output_item.added" and e.item.type == "compaction"] + done: Final = [e for e in events if e.type == "response.output_item.done" and e.item.type == "compaction"] + completed: Final = [e for e in events if e.type == "response.completed"] + assert len(added) == 1 and len(done) == 1 and len(completed) == 1, [e.type for e in events] + assert added[0].item.encrypted_content == "" + assert done[0].output_index == added[0].output_index + assert done[0].item.id == added[0].item.id + assert _decoded(done[0].item.encrypted_content) == _block(_summary(marker)) + finished: Final = [item for item in completed[0].response.output if item.type == "compaction"] + assert [(item.id, item.encrypted_content) for item in finished] == [ + (done[0].item.id, done[0].item.encrypted_content) + ] + message_indexes: Final = { + e.output_index for e in events if e.type == "response.output_item.added" and e.item.type == "message" + } + assert message_indexes and added[0].output_index not in message_indexes + assert completed[0].response.output_text == _answer(marker) + call, body = _only_call(wire, marker) + assert _COMPACT_BETA in call.headers.get("anthropic-beta", ""), dict(call.headers) + + +def _history(marker: str, compaction: Mapping[str, JsonValue]) -> list[JsonValue]: + return [ + {"role": "developer", "content": f"rules-{marker}"}, + {"role": "user", "content": f"old-{marker}"}, + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": f"old-answer-{marker}", "annotations": []}], + }, + dict(compaction), + { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": f"recap-{marker}", "annotations": []}], + }, + {"role": "user", "content": _ask(marker)}, + ] + + +_ROUTES: Final[tuple[Route, ...]] = ( + "claude-compacting", + "claude-plain", + "gemini", + "openai", + "openai-compact", + "openai-opaque", +) + + +def _route_provider(route: Route) -> Literal["claude", "gemini", "openai"]: + if route.startswith("claude"): + return "claude" + return "gemini" if route == "gemini" else "openai" + + +def _assert_reached( + route: Route, marker: str, sent: Mapping[str, JsonValue], call: Request, body: Mapping[str, JsonValue] +) -> None: + summary: Final = _summary(marker) + dumped: Final = json.dumps(body) + if route == "openai-opaque": + assert urlsplit(call.target).path.endswith("/responses"), call.target + assert sent in _objects(body["input"]), body["input"] + assert f"old-{marker}" in dumped + return + if route != "claude-compacting": + assert f"old-{marker}" not in dumped and f"old-answer-{marker}" not in dumped, dumped[:2000] + assert f"rules-{marker}" in dumped + assert f"recap-{marker}" in dumped + match route: + case "claude-compacting": + assistant: Final = tuple( + message for message in _objects(body["messages"]) if message["role"] == "assistant" + ) + blocks: Final = tuple( + block + for message in assistant + if isinstance(message["content"], list) + for block in _objects(message["content"]) + ) + assert _block(summary) in blocks, body["messages"] + assert summary not in _anthropic_texts(body, "user") + assert f"rules-{marker}" in json.dumps(body["system"]) + assert _COMPACT_BETA in call.headers.get("anthropic-beta", ""), dict(call.headers) + case "claude-plain": + assert "compaction" not in _anthropic_block_types(body), body["messages"] + assert summary in _anthropic_texts(body, "user"), body["messages"] + case "gemini": + assert summary in _gemini_texts(body, "user"), body["contents"] + case "openai" | "openai-compact": + suffix: Final = "/responses/compact" if route == "openai-compact" else "/responses" + assert urlsplit(call.target).path.endswith(suffix), call.target + assert "compaction" not in _openai_item_types(body), body["input"] + assert summary in _openai_user_texts(body), body["input"] + + +@pytest.mark.parametrize("route", _ROUTES) +def test_compaction_item_reaches_each_route_in_the_shape_it_reads(gateway: Gateway, route: Route) -> None: + marker: Final = _new_marker() + compaction: Final = ( + {"type": "compaction", "id": f"cmp_{marker}", "encrypted_content": f"gAAAAA-opaque-{marker}"} + if route == "openai-opaque" + else _compaction_item(_summary(marker)) + ) + compacting: Final[dict[str, JsonValue]] = ( + {"context_management": _COMPACTING} if route == "claude-compacting" else {} + ) + dropping: Final[dict[str, JsonValue]] = ( + {"context_management": _COMPACTING, "drop_params": True} if route == "gemini" else {} + ) + path: Final = "/v1/responses/compact" if route == "openai-compact" else "/v1/responses" + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, _route_provider(route)) + response: Final = gateway.request( + "POST", + path, + {"model": model, "input": _history(marker, compaction), **compacting, **dropping, **_CACHE_BUST}, + ) + assert response.status_code == 200, response.text + if route != "openai-compact": + assert _answer(marker) in response.text, response.text + call, body = _only_call(wire, marker) + _assert_reached(route, marker, compaction, call, body) + + +def _input_tokens(gateway: Gateway, model: str, items: list[JsonValue]) -> int: + counted: Final = gateway.post("/v1/responses/input_tokens", {"model": model, "input": items}) + tokens: Final = counted["input_tokens"] + assert isinstance(tokens, int), counted + return tokens + + +def test_input_tokens_count_from_the_newest_compaction(gateway: Gateway) -> None: + marker: Final = _new_marker() + summary: Final = _summary(marker) + old_history: Final = " ".join(f"history-{index}" for index in range(600)) + compacted: Final[list[JsonValue]] = [ + {"role": "user", "content": old_history}, + _compaction_item(summary), + {"role": "user", "content": _ask(marker)}, + ] + equivalent: Final[list[JsonValue]] = [ + {"role": "user", "content": summary}, + {"role": "user", "content": _ask(marker)}, + ] + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, "claude") + assert _input_tokens(gateway, model, compacted) == _input_tokens(gateway, model, equivalent) + counted: Final = tuple( + request.body for request in wire.drain() if urlsplit(request.target).path.endswith("/count_tokens") + ) + assert len(counted) == 2, counted + assert all(b"history-" not in body for body in counted), counted[0][:2000] + + +def test_chat_stream_reports_the_whole_compaction_block(gateway: Gateway) -> None: + marker: Final = _new_marker() + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, "claude") + chunks: Final = _stream( + gateway, + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _ask(marker, compact=True)}], + "context_management": {"edits": [{"type": "compact_20260112"}]}, + }, + ) + reported: Final = tuple( + fields["compaction_blocks"] + for chunk in chunks + for choice in _objects(chunk.get("choices", [])) + if isinstance(fields := object_value(choice.get("delta", {})).get("provider_specific_fields"), dict) + and "compaction_blocks" in fields + ) + assert reported, chunks + assert _objects(reported[-1])[-1] == _block(_summary(marker)), reported[-1] + assert len(_provider_calls(wire, marker)) == 1 + + +def _guardrail_peer(request: Request) -> Reply: + payload: Final = _JSON_OBJECT.validate_json(request.body) + texts: Final = payload.get("texts") or [] + assert isinstance(texts, list), payload + if urlsplit(request.target).path.startswith("/block/"): + found: Final = any(isinstance(text, str) and _CITY in text for text in texts) + verdict: Final = {"action": "BLOCKED", "blocked_reason": "city in summary"} if found else {"action": "NONE"} + return Reply(body=json.dumps(verdict).encode()) + masked: Final = [text.replace(_CITY, _MASKED_CITY) if isinstance(text, str) else text for text in texts] + return Reply(body=json.dumps({"action": "GUARDRAIL_INTERVENED", "texts": masked}).encode()) + + +def _guardrail(name: str, mode: str, api_base: str, **fields: JsonValue) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": mode, + "default_on": False, + "api_base": api_base, + "api_key": "synthetic-guardrail-key", + **fields, + }, + } + + +def _deployment(name: str, model: str, api_base: str, **fields: JsonValue) -> dict[str, JsonValue]: + return {"model_name": name, "litellm_params": {"model": model, "api_base": api_base, **fields}} + + +_STORED_CLAUDE: Final = "compaction-stored-claude" +_SMALL_CLAUDE: Final = "compaction-small-claude" +_STORED_GEMINI: Final = "compaction-stored-gemini" +_STORED_OPENAI: Final = "compaction-stored-openai" +_SMALL_CONTEXT: Final = 2000 + + +def _stored_config(directory: Path, provider_url: str, guardrail_url: str) -> Path: + stock: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + claude: Final = {"api_key": cc.ANTHROPIC_API_KEY} + config: Final[dict[str, JsonValue]] = { + **stock, + "model_list": [ + _deployment(_STORED_CLAUDE, f"anthropic/{_CLAUDE}", provider_url, **claude), + { + **_deployment(_SMALL_CLAUDE, f"anthropic/{_CLAUDE}", provider_url, **claude), + "model_info": {"max_input_tokens": _SMALL_CONTEXT}, + }, + _deployment(_STORED_GEMINI, f"gemini/{_GEMINI}", provider_url, api_key="synthetic-gemini-key"), + _deployment(_STORED_OPENAI, f"openai/{_OPENAI}", f"{provider_url}/v1", api_key="synthetic-openai-key"), + ], + "general_settings": { + **object_value(stock["general_settings"]), + "store_prompts_in_spend_logs": True, + "disable_responses_id_security": True, + }, + "router_settings": {**object_value(stock["router_settings"]), "num_retries": 0, "enable_pre_call_checks": True}, + "guardrails": [ + _guardrail(_MASK_INPUT, "pre_call", f"{guardrail_url}/input"), + _guardrail(_MASK_OUTPUT, "post_call", f"{guardrail_url}/output"), + _guardrail(_BLOCK_OUTPUT, "post_call", f"{guardrail_url}/block", streaming_buffer_until_moderated=True), + ], + } + path: Final = directory / "compaction-stored.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@dataclass(frozen=True, slots=True) +class _StoredRig: + gateway: Gateway + provider: Wire + guardrail: Wire + + +@pytest.fixture(scope="module") +def stored(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_StoredRig]: + directory: Final = tmp_path_factory.mktemp("compaction-stored") + with ( + wire_server(_vendor) as provider, + wire_server(_guardrail_peer) as guardrail, + gateway_from_environment() as gateway, + ): + config: Final = _stored_config(directory, provider.url, guardrail.url) + with owned_proxy(gateway, directory, {}, config=config, workers=2) as owned: + yield _StoredRig(owned, provider, guardrail) + + +def _spend_row(response_id: str) -> dict[str, JsonValue]: + decoded: Final = ResponsesAPIRequestUtils.decode_responses_api_response_id(response_id).get("response_id") + request_id: Final = decoded if isinstance(decoded, str) else response_id + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, proxy_server_request, response, status FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (request_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["status"] == "success", rows[0] + return rows[0] + + +_REPLAY_ROUTES: Final = ( + pytest.param(_STORED_CLAUDE, True, id="claude-compacting"), + pytest.param(_STORED_CLAUDE, False, id="claude-plain"), + pytest.param(_STORED_GEMINI, True, id="gemini"), +) + + +@pytest.mark.parametrize(("follow_up_model", "compacting"), _REPLAY_ROUTES) +def test_previous_response_id_after_a_compaction_turn_replays_the_whole_summary( + stored: _StoredRig, follow_up_model: str, compacting: bool +) -> None: + first: Final = _new_marker() + second: Final = _new_marker() + summary: Final = _summary(first) + long_text: Final = f"long-{first} {'z' * 3000} tail-{first}" + compaction_request: Final[dict[str, JsonValue]] = {"context_management": _COMPACTING} + opened: Final = _responses( + stored.gateway, + { + "model": _STORED_CLAUDE, + "input": [{"role": "user", "content": long_text}, {"role": "user", "content": _ask(first, compact=True)}], + **compaction_request, + }, + ) + assert _decoded(string_value(_only(_output_items(opened, "compaction"))["encrypted_content"])) == _block(summary) + row: Final = _spend_row(string_value(opened["id"])) + assert summary in json.dumps(row["response"]), json.dumps(row["response"])[:600] + stored_request: Final = json.dumps(row["proxy_server_request"]) + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in stored_request and "z" * 3000 not in stored_request + _only_call(stored.provider, first) + + follow_up: Final = _responses( + stored.gateway, + { + "model": follow_up_model, + "input": _ask(second), + "previous_response_id": opened["id"], + **(compaction_request if compacting else {}), + **({"drop_params": True} if follow_up_model == _STORED_GEMINI else {}), + }, + ) + assert _answer(second) in json.dumps(follow_up) + call, body = _only_call(stored.provider, second) + dumped: Final = json.dumps(body) + assert f"long-{first}" not in dumped and f"ask-{first}" not in dumped, dumped[:2000] + assert _answer(first) in dumped + if follow_up_model == _STORED_GEMINI: + assert summary in _gemini_texts(body, "user"), body["contents"] + elif compacting: + replayed: Final = _objects(body["messages"])[0] + assert replayed["role"] == "assistant", body["messages"] + assert _objects(replayed["content"])[0] == _block(summary), replayed + assert _COMPACT_BETA in call.headers.get("anthropic-beta", ""), dict(call.headers) + else: + assert "compaction" not in _anthropic_block_types(body), body["messages"] + assert summary in _anthropic_texts(body, "user"), body["messages"] + + +def test_router_pre_call_check_counts_from_the_newest_compaction(stored: _StoredRig) -> None: + compacted_marker: Final = _new_marker() + control_marker: Final = _new_marker() + old_history: Final = " ".join(f"history-{index}" for index in range(3 * _SMALL_CONTEXT)) + + def send(marker: str, items: list[JsonValue]) -> httpx.Response: + return stored.gateway.request( + "POST", + "/v1/responses", + {"model": _SMALL_CLAUDE, "input": [*items, {"role": "user", "content": _ask(marker)}], **_CACHE_BUST}, + ) + + compacted: Final = send( + compacted_marker, + [{"role": "user", "content": old_history}, _compaction_item(f"short summary {compacted_marker}")], + ) + assert compacted.status_code == 200, compacted.text + _only_call(stored.provider, compacted_marker) + control: Final = send(control_marker, [{"role": "user", "content": old_history}]) + assert control.status_code == 400, control.text + assert _provider_calls(stored.provider, control_marker) == () + + +_GUARDED_ROUTES: Final = ( + pytest.param(_STORED_OPENAI, id="openai"), + pytest.param(_STORED_CLAUDE, id="claude-compacting"), +) + + +@pytest.mark.parametrize("model", _GUARDED_ROUTES) +def test_pre_call_guardrail_rewrites_the_compaction_summary(stored: _StoredRig, model: str) -> None: + marker: Final = _new_marker() + summary: Final = _summary(marker) + masked: Final = summary.replace(_CITY, _MASKED_CITY) + compacting: Final[dict[str, JsonValue]] = {"context_management": _COMPACTING} if model == _STORED_CLAUDE else {} + _responses( + stored.gateway, + { + "model": model, + "input": [_compaction_item(summary), {"role": "user", "content": _ask(marker)}], + "guardrails": [_MASK_INPUT], + **compacting, + }, + ) + scanned: Final = tuple( + _JSON_OBJECT.validate_json(request.body) + for request in stored.guardrail.drain() + if f"kept-{marker}".encode() in request.body + ) + assert any(summary in _listed(payload["texts"]) for payload in scanned), scanned + call, body = _only_call(stored.provider, marker) + assert _CITY not in json.dumps(body), json.dumps(body)[:2000] + if model == _STORED_OPENAI: + assert masked in _openai_user_texts(body), body["input"] + else: + assert _objects(_objects(body["messages"])[0]["content"])[0] == _block(masked), body["messages"] + + +def test_post_call_guardrail_rewrites_the_returned_compaction_summary(stored: _StoredRig) -> None: + marker: Final = _new_marker() + returned: Final = _responses( + stored.gateway, + { + "model": _STORED_CLAUDE, + "input": _ask(marker, compact=True), + "context_management": _COMPACTING, + "guardrails": [_MASK_OUTPUT], + }, + ) + masked: Final = _block(_summary(marker).replace(_CITY, _MASKED_CITY)) + assert _decoded(string_value(_only(_output_items(returned, "compaction"))["encrypted_content"])) == masked + _only_call(stored.provider, marker) + + +def test_buffered_stream_guardrail_blocks_on_the_compaction_summary(stored: _StoredRig) -> None: + marker: Final = _new_marker() + body: Final[dict[str, JsonValue]] = { + "model": _STORED_CLAUDE, + "input": _ask(marker, compact=True), + "context_management": _COMPACTING, + "guardrails": [_BLOCK_OUTPUT], + "stream": True, + **_CACHE_BUST, + } + headers: Final = {"Authorization": f"Bearer {stored.gateway.key}"} + with stored.gateway.client.stream("POST", "/v1/responses", json=body, headers=headers) as response: + text: Final = response.read().decode() + assert response.status_code == 200, text + peer_calls: Final = tuple( + (urlsplit(request.target).path, _JSON_OBJECT.validate_json(request.body).get("texts")) + for request in stored.guardrail.drain() + if marker.encode() in request.body + ) + assert any(path.startswith("/block/") and _summary(marker) in _listed(texts) for path, texts in peer_calls), [ + (path, [str(value)[:80] for value in _listed(texts or [])]) for path, texts in peer_calls + ] + event_types: Final = tuple(event["type"] for event in _sse_events(text)) + assert "response.failed" in event_types, event_types + assert "response.completed" not in event_types, event_types + assert f"kept-{marker}" not in text and _CITY not in text, text[:2000] + _only_call(stored.provider, marker) + + +def test_redacted_spend_row_holds_no_compaction_summary(stored: _StoredRig) -> None: + marker: Final = _new_marker() + returned: Final = _responses( + stored.gateway, + {"model": _STORED_CLAUDE, "input": _ask(marker, compact=True), "context_management": _COMPACTING}, + headers=_REDACT_HEADER, + ) + assert _decoded(string_value(_only(_output_items(returned, "compaction"))["encrypted_content"])) == _block( + _summary(marker) + ) + row: Final = _spend_row(string_value(returned["id"])) + assert f"kept-{marker}" not in json.dumps(row["response"]), json.dumps(row["response"])[:2000] + _only_call(stored.provider, marker) diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 7e1bce8e46b..b98bf974f52 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -56,6 +56,7 @@ def _provider_fields() -> dict[str, object]: "citations": [{"cited_text": SECRET}], "thought_signatures": [SIGNATURE], "compaction_blocks": [{"type": "compaction", "content": SECRET}], + "compaction_delta": {"type": "compaction_delta", "content": SECRET}, } diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 8badf606c4b..aa1fdb2a37a 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -1011,29 +1011,41 @@ def test_get_response_for_spend_logs_payload_keeps_a_long_compaction_summary_who assert LITELLM_TRUNCATED_PAYLOAD_FIELD in message["content"] -def test_sanitize_request_body_keeps_a_long_claude_compaction_item_whole(): +@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") +def test_stored_request_keeps_only_the_summary_of_a_long_claude_compaction_input_item(mock_should_store): from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB + mock_should_store.return_value = True filler = "detail " * (MAX_STRING_LENGTH_PROMPT_IN_DB // 7) - claude_item = json.dumps({"type": "compaction", "content": f"{filler}The cat is named Biscuit. {filler}"}) + summary = f"{filler}The cat is named Biscuit. {filler}" + claude_item = json.dumps({"type": "compaction", "content": summary, "padding": filler * 10}) opaque_item = "gAAAA" + "x" * (MAX_STRING_LENGTH_PROMPT_IN_DB * 2) + request_body = { + "input": [ + {"type": "compaction", "id": "cmp_1", "encrypted_content": claude_item}, + {"type": "compaction", "id": "cmp_2", "encrypted_content": opaque_item}, + {"role": "assistant", "content": [{"type": "compaction", "content": f"{filler}Biscuit{filler}"}]}, + {"role": "user", "content": f"{filler}What is my cat called?{filler}"}, + ], + "metadata": {"note": {"type": "compaction", "encrypted_content": claude_item}}, + } - sanitized = _sanitize_request_body_for_spend_logs_payload( - { - "input": [ - {"type": "compaction", "id": "cmp_1", "encrypted_content": claude_item}, - {"type": "compaction", "id": "cmp_2", "encrypted_content": opaque_item}, - {"role": "assistant", "content": [{"type": "compaction", "content": f"{filler}Biscuit{filler}"}]}, - {"role": "user", "content": f"{filler}What is my cat called?{filler}"}, - ] - } + stored = json.loads( + _get_proxy_server_request_for_spend_logs_payload( + metadata={}, litellm_params={"proxy_server_request": {"body": request_body}} + ) ) - claude_stored, opaque_stored, messages_block_stored, message_stored = sanitized["input"] - assert claude_stored["encrypted_content"] == claude_item + claude_stored, opaque_stored, messages_block_stored, message_stored = stored["input"] + assert {**claude_stored, "encrypted_content": json.loads(claude_stored["encrypted_content"])} == { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": {"type": "compaction", "content": summary}, + } assert LITELLM_TRUNCATED_PAYLOAD_FIELD in opaque_stored["encrypted_content"] assert LITELLM_TRUNCATED_PAYLOAD_FIELD in messages_block_stored["content"][0]["content"] assert LITELLM_TRUNCATED_PAYLOAD_FIELD in message_stored["content"] + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in stored["metadata"]["note"]["encrypted_content"] @patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 72e539a1987..59ac4a633ae 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -671,6 +671,72 @@ def test_streamed_responses_over_anthropic_emit_the_compaction_item_and_send_it_ assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] +@pytest.mark.asyncio +async def test_async_streamed_compaction_item_is_done_before_the_reasoning_that_follows_it() -> None: + anthropic: Final = httpx.MockTransport( + lambda _request: httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=_anthropic_sse( + { + "type": "message_start", + "message": { + "id": "msg_compacted_thinking", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [], + "usage": {"input_tokens": 9, "output_tokens": 1}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "compaction", "content": None}}, + { + "type": "content_block_delta", + "index": 0, + "delta": {"type": "compaction_delta", "content": _COMPACTION_BLOCK["content"]}, + }, + {"type": "content_block_stop", "index": 0}, + {"type": "content_block_start", "index": 1, "content_block": {"type": "thinking", "thinking": ""}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "thinking_delta", "thinking": "Recall."}}, + {"type": "content_block_delta", "index": 1, "delta": {"type": "signature_delta", "signature": "sig"}}, + {"type": "content_block_stop", "index": 1}, + {"type": "content_block_start", "index": 2, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 2, "delta": {"type": "text_delta", "text": "The."}}, + {"type": "content_block_stop", "index": 2}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 4}}, + {"type": "message_stop"}, + ), + ) + ) + + stream: Final = await litellm.aresponses( + model=_MODEL, + input=list(_FIRST_INPUT), + context_management=list(_COMPACT_AT_50K), + stream=True, + api_key="mock_api_key", + client=AsyncHTTPHandler(transport=anthropic), + ) + events: Final = [event async for event in stream] + done_items: Final = [ + (event.output_index, event.item.type) + for event in events + if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE + ] + reasoning_done: Final = next( + index + for index, event in enumerate(events) + if event.type == ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE and event.item.type == "reasoning" + ) + sequence_numbers: Final = [ + event.sequence_number for event in events[: reasoning_done + 1] if getattr(event, "sequence_number", None) + ] + + assert done_items == [(0, "compaction"), (1, "reasoning"), (2, "message")] + assert [item.type for item in events[-1].response.output] == ["compaction", "reasoning", "message"] + assert sequence_numbers == sorted(set(sequence_numbers)) + + def _compaction_stream_events(*summaries: str) -> tuple[Mapping[str, object], ...]: return tuple( chain.from_iterable( From 2c6fd893b955413a9bb84a7c79fccbc5d4ddbfbe Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 23:16:54 -0700 Subject: [PATCH 08/21] test(responses): cover stored compaction items, compaction_delta redaction, and compaction done order Adds integration cells for the summary-only spend row of a caller-sent compaction item and its previous_response_id replay, a redacted chat stream row that holds no streamed compaction_delta, and a compaction then thinking stream whose done events follow output order. --- .../test_responses_compaction_items_wire.py | 168 +++++++++++++++--- 1 file changed, 146 insertions(+), 22 deletions(-) diff --git a/tests/integration/providers/test_responses_compaction_items_wire.py b/tests/integration/providers/test_responses_compaction_items_wire.py index bd4f56da124..34a681d9e74 100644 --- a/tests/integration/providers/test_responses_compaction_items_wire.py +++ b/tests/integration/providers/test_responses_compaction_items_wire.py @@ -4,6 +4,7 @@ import re import uuid from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass +from itertools import chain from pathlib import Path from typing import Final, Literal, TypeAlias from urllib.parse import urlsplit @@ -59,6 +60,14 @@ def _ask(marker: str, *, compact: bool = False) -> str: return f"ask-{marker} compact-now-{marker}" if compact else f"ask-{marker}" +def _think(marker: str) -> str: + return f"think-now-{marker}" + + +def _thinking(marker: str) -> dict[str, JsonValue]: + return {"type": "thinking", "thinking": f"thought-{marker}", "signature": f"signature-{marker}"} + + def _answer(marker: str) -> str: return f"answer-{marker}" @@ -101,16 +110,25 @@ def _anthropic_block_frames(index: int, block: Mapping[str, JsonValue]) -> tuple ) stop: Final = cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": index}) - if block["type"] == "compaction": - summary: Final = string_value(block["content"]) - half: Final = len(summary) // 2 - return ( - start({"type": "compaction", "content": ""}), - delta({"type": "compaction_delta", "content": summary[:half]}), - delta({"type": "compaction_delta", "content": summary[half:]}), - stop, - ) - return (start({"type": "text", "text": ""}), delta({"type": "text_delta", "text": block["text"]}), stop) + match block["type"]: + case "compaction": + summary: Final = string_value(block["content"]) + half: Final = len(summary) // 2 + return ( + start({"type": "compaction", "content": ""}), + delta({"type": "compaction_delta", "content": summary[:half]}), + delta({"type": "compaction_delta", "content": summary[half:]}), + stop, + ) + case "thinking": + return ( + start({"type": "thinking", "thinking": ""}), + delta({"type": "thinking_delta", "thinking": block["thinking"]}), + delta({"type": "signature_delta", "signature": block["signature"]}), + stop, + ) + case _: + return (start({"type": "text", "text": ""}), delta({"type": "text_delta", "text": block["text"]}), stop) def _anthropic_stream(identity: str, content: Sequence[Mapping[str, JsonValue]]) -> tuple[bytes, ...]: @@ -147,7 +165,8 @@ def _anthropic_stream(identity: str, content: Sequence[Mapping[str, JsonValue]]) def _claude_reply(body: Mapping[str, JsonValue], raw: bytes) -> Reply: marker: Final = _newest_marker(raw) compaction: Final = (_block(_summary(marker)),) if f"compact-now-{marker}".encode() in raw else () - content: Final = (*compaction, {"type": "text", "text": _answer(marker)}) + thinking: Final = (_thinking(marker),) if _think(marker).encode() in raw else () + content: Final = (*compaction, *thinking, {"type": "text", "text": _answer(marker)}) identity: Final = f"msg_{uuid.uuid4().hex}" if body.get("stream") is True: return Reply(content_type="text/event-stream", chunks=_anthropic_stream(identity, content)) @@ -321,14 +340,31 @@ def _sse_events(text: str) -> tuple[dict[str, JsonValue], ...]: return tuple(_JSON_OBJECT.validate_json(payload) for payload in payloads if payload != "[DONE]") -def _stream(gateway: Gateway, path: str, body: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: - headers: Final = {"Authorization": f"Bearer {gateway.key}"} - with gateway.client.stream("POST", path, json={**body, **_CACHE_BUST, "stream": True}, headers=headers) as response: +def _stream( + gateway: Gateway, path: str, body: Mapping[str, JsonValue], headers: Mapping[str, str] | None = None +) -> tuple[dict[str, JsonValue], ...]: + sent_headers: Final = {"Authorization": f"Bearer {gateway.key}", **(headers or {})} + with gateway.client.stream( + "POST", path, json={**body, **_CACHE_BUST, "stream": True}, headers=sent_headers + ) as response: text: Final = response.read().decode() assert response.status_code == 200, text return _sse_events(text) +def _delta_provider_fields(chunk: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + return tuple( + fields + for choice in _objects(chunk.get("choices", [])) + if isinstance(fields := object_value(choice.get("delta", {})).get("provider_specific_fields"), dict) + ) + + +def _streamed_compaction_blocks(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]: + fields: Final = chain.from_iterable(_delta_provider_fields(chunk) for chunk in chunks) + return tuple(field["compaction_blocks"] for field in fields if "compaction_blocks" in field) + + def _output_items(response: Mapping[str, JsonValue], item_type: str) -> tuple[dict[str, JsonValue], ...]: return tuple(item for item in _objects(response["output"]) if item["type"] == item_type) @@ -518,6 +554,41 @@ def test_input_tokens_count_from_the_newest_compaction(gateway: Gateway) -> None assert all(b"history-" not in body for body in counted), counted[0][:2000] +def test_streamed_compaction_item_is_done_before_the_reasoning_after_it(gateway: Gateway) -> None: + marker: Final = _new_marker() + with wire_server(_vendor) as wire, gateway.scenario() as scenario: + model: Final = _register(scenario, wire, "claude") + events: Final = _stream( + gateway, + "/v1/responses", + { + "model": model, + "input": f"{_ask(marker, compact=True)} {_think(marker)}", + "context_management": _COMPACTING, + }, + ) + done: Final = tuple( + (position, event["output_index"], object_value(event["item"])["type"]) + for position, event in enumerate(events) + if event["type"] == "response.output_item.done" + ) + completed: Final = _only([event for event in events if event["type"] == "response.completed"]) + assert [(index, kind) for _, index, kind in done] == [(0, "compaction"), (1, "reasoning"), (2, "message")], [ + event["type"] for event in events + ] + assert [item["type"] for item in _objects(object_value(completed["response"])["output"])] == [ + "compaction", + "reasoning", + "message", + ] + reasoning_done: Final = done[1][0] + sequence_numbers: Final = [ + number for event in events[: reasoning_done + 1] if isinstance(number := event.get("sequence_number"), int) + ] + assert sequence_numbers == sorted(set(sequence_numbers)), sequence_numbers + _only_call(wire, marker) + + def test_chat_stream_reports_the_whole_compaction_block(gateway: Gateway) -> None: marker: Final = _new_marker() with wire_server(_vendor) as wire, gateway.scenario() as scenario: @@ -531,13 +602,7 @@ def test_chat_stream_reports_the_whole_compaction_block(gateway: Gateway) -> Non "context_management": {"edits": [{"type": "compact_20260112"}]}, }, ) - reported: Final = tuple( - fields["compaction_blocks"] - for chunk in chunks - for choice in _objects(chunk.get("choices", [])) - if isinstance(fields := object_value(choice.get("delta", {})).get("provider_specific_fields"), dict) - and "compaction_blocks" in fields - ) + reported: Final = _streamed_compaction_blocks(chunks) assert reported, chunks assert _objects(reported[-1])[-1] == _block(_summary(marker)), reported[-1] assert len(_provider_calls(wire, marker)) == 1 @@ -636,7 +701,8 @@ def _spend_row(response_id: str) -> dict[str, JsonValue]: request_id: Final = decoded if isinstance(decoded, str) else response_id rows: Final = eventually( lambda: read_rows( - 'SELECT request_id, proxy_server_request, response, status FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + 'SELECT request_id, proxy_server_request, response, metadata, status FROM "LiteLLM_SpendLogs"' + " WHERE request_id=%s", (request_id,), ), lambda values: len(values) == 1, @@ -820,3 +886,61 @@ def test_redacted_spend_row_holds_no_compaction_summary(stored: _StoredRig) -> N row: Final = _spend_row(string_value(returned["id"])) assert f"kept-{marker}" not in json.dumps(row["response"]), json.dumps(row["response"])[:2000] _only_call(stored.provider, marker) + + +def test_stored_compaction_input_item_keeps_its_whole_summary_and_nothing_else(stored: _StoredRig) -> None: + first: Final = _new_marker() + second: Final = _new_marker() + summary: Final = _summary(first) + padded_item: Final = { + **_compaction_item(summary), + "encrypted_content": json.dumps({**_block(summary), "padding": f"pad-{first} {'p' * 3000}"}), + } + opened: Final = _responses( + stored.gateway, + { + "model": _STORED_CLAUDE, + "input": [padded_item, {"role": "user", "content": _ask(first)}], + "context_management": _COMPACTING, + }, + ) + stored_input: Final = _objects( + object_value(_spend_row(string_value(opened["id"]))["proxy_server_request"])["input"] + ) + assert _decoded(string_value(stored_input[0]["encrypted_content"])) == _block(summary), stored_input[0] + _only_call(stored.provider, first) + + _responses( + stored.gateway, + { + "model": _STORED_CLAUDE, + "input": _ask(second), + "previous_response_id": opened["id"], + "context_management": _COMPACTING, + }, + ) + _, body = _only_call(stored.provider, second) + assert summary in _anthropic_texts(body, "user"), body["messages"] + assert f"pad-{first}" not in json.dumps(body), body["messages"] + + +def test_redacted_chat_stream_spend_row_holds_no_compaction_delta(stored: _StoredRig) -> None: + marker: Final = _new_marker() + summary: Final = _summary(marker) + chunks: Final = _stream( + stored.gateway, + "/v1/chat/completions", + { + "model": _STORED_CLAUDE, + "messages": [{"role": "user", "content": _ask(marker, compact=True)}], + "context_management": {"edits": [{"type": "compact_20260112"}]}, + }, + headers=_REDACT_HEADER, + ) + streamed_blocks: Final = _streamed_compaction_blocks(chunks) + assert streamed_blocks and _objects(streamed_blocks[-1])[-1] == _block(summary), chunks + secrets: Final = (f"kept-{marker}", summary[len(summary) // 2 :]) + logged: Final = {column: json.dumps(value) for column, value in _spend_row(string_value(chunks[0]["id"])).items()} + leaked: Final = {column: text for column, text in logged.items() if any(secret in text for secret in secrets)} + assert not leaked, {column: text[max(0, text.find(secrets[0]) - 600) :][:1500] for column, text in leaked.items()} + _only_call(stored.provider, marker) From ef4d736d3baf05e6a0a94e2a883da112b5ffb5ff Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 00:23:00 -0700 Subject: [PATCH 09/21] fix(spend_logs): honor the redaction header on failure rows for litellm_metadata routes A failed /v1/responses or /v1/messages call keeps its proxy request headers in litellm_metadata, but the failure row read redaction headers from metadata only, so x-litellm-enable-message-redaction never redacted those rows. --- .../proxy/hooks/proxy_track_cost_callback.py | 32 +++++++++++---- .../hooks/test_proxy_track_cost_callback.py | 40 +++++++++++++++++++ 2 files changed, 65 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 5c31bce6f33..d3c52269d54 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -4,6 +4,8 @@ from collections.abc import Callable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Final, Protocol, cast +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import batch_cost_is_final @@ -87,6 +89,9 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( ) +_OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping) + + def _proxy_stamped_used_client_oauth_token( request_data: Mapping[str, object], request_route: str | None ) -> bool | None: @@ -112,6 +117,18 @@ def _failure_snapshot_with_guardrails_recorded_after_it( return {**failure_snapshot, "guardrail_information": recorded_guardrail_information} +def _failure_metadata_lifted_from_litellm_metadata( + litellm_metadata: ObjectMapping, *, guardrail_info_already_set: bool +) -> ObjectMapping: + guardrail_info: Final = litellm_metadata.get("standard_logging_guardrail_information") + request_headers: Final = litellm_metadata.get("headers") + keeps_guardrail_info: Final = guardrail_info is not None and not guardrail_info_already_set + return { + **({"standard_logging_guardrail_information": guardrail_info} if keeps_guardrail_info else {}), + **({"headers": request_headers} if isinstance(request_headers, dict) else {}), + } + + def _proxy_spend_writer() -> DBSpendUpdateWriter: from litellm.proxy.proxy_server import proxy_logging_obj @@ -226,13 +243,14 @@ class _ProxyDBLogger(CustomLogger): existing_metadata.update(_metadata) litellm_metadata_bucket: Final = request_data.get("litellm_metadata") - if ( - isinstance(litellm_metadata_bucket, dict) - and "standard_logging_guardrail_information" not in existing_metadata - ): - guardrail_info: Final = litellm_metadata_bucket.get("standard_logging_guardrail_information") - if guardrail_info is not None: - existing_metadata["standard_logging_guardrail_information"] = guardrail_info + existing_metadata.update( + _failure_metadata_lifted_from_litellm_metadata( + _OBJECT_MAPPING.validate_python(litellm_metadata_bucket) + if isinstance(litellm_metadata_bucket, dict) + else {}, + guardrail_info_already_set="standard_logging_guardrail_information" in existing_metadata, + ) + ) if "litellm_params" not in request_data: request_data["litellm_params"] = {} diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 833af4c03a3..f676914fe88 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -161,6 +161,46 @@ async def test_async_post_call_failure_hook_does_not_clobber_guardrail_info_in_m assert metadata["standard_logging_guardrail_information"] == metadata_bucket_info +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_headers, code_word_is_stored", + [({"x-litellm-enable-message-redaction": "true"}, False), ({}, True)], +) +async def test_async_post_call_failure_hook_honors_the_redaction_header_on_a_litellm_metadata_route( + monkeypatch: pytest.MonkeyPatch, request_headers: dict[str, str], code_word_is_stored: bool +): + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + writer: Final = MagicMock(spec=DBSpendUpdateWriter) + writer.update_database = AsyncMock() + logger: Final = ProxyDBLogger(spend_writer=lambda: writer) + request_input: Final = "The code word is VERMILION-42. Reply OK." + + await logger.async_post_call_failure_hook( + request_data={ + "model": "gpt-5.6", + "input": request_input, + "metadata": {"caller_tag": "kept"}, + "litellm_metadata": {"headers": request_headers}, + "proxy_server_request": { + "url": "http://localhost/v1/responses", + "body": {"model": "gpt-5.6", "input": request_input, "temperature": 7}, + }, + }, + original_exception=Exception("temperature must be at most 2"), + user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"), + ) + + row: Final = get_logging_payload( + kwargs=writer.update_database.call_args.kwargs["kwargs"], + response_obj={}, + start_time=datetime.now(), + end_time=datetime.now(), + ) + stored_request: Final = json.loads(row["proxy_server_request"]) + assert stored_request["temperature"] == 7 + assert ("VERMILION-42" in row["proxy_server_request"]) is code_word_is_stored + + @pytest.mark.asyncio @pytest.mark.parametrize( "used_client_oauth_token, custom_llm_provider, expected", From 3220b2dfa15f018038010f6b741a18ee87784670 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:08:38 -0700 Subject: [PATCH 10/21] fix(spend_logs): lift only the redaction opt-in header into failure rows The failure hook on litellm_metadata routes copied every real request header into the caller's own metadata dict, so the stored request body gained the headers, a caller's metadata.headers was replaced, and a disable header loosened global redaction on failed calls. It now copies only the enable headers into a new litellm_params metadata dict, and only when the caller sent no metadata.headers, so no failure row is less redacted than on main --- litellm/litellm_core_utils/redact_messages.py | 9 ++-- .../proxy/hooks/proxy_track_cost_callback.py | 39 +++++++++++++----- .../hooks/test_proxy_track_cost_callback.py | 41 ++++++++++++++++--- 3 files changed, 66 insertions(+), 23 deletions(-) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 3b50b9af24c..0cff0a0b0ef 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -40,6 +40,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObject = Any +MESSAGE_REDACTION_ENABLE_HEADERS: Final = ("litellm-enable-message-redaction", "x-litellm-enable-message-redaction") + def redact_message_input_output_from_custom_logger( litellm_logging_obj: LiteLLMLoggingObject, result, custom_logger: CustomLogger @@ -458,13 +460,8 @@ def should_redact_message_logging(model_call_details: dict) -> bool: # User explicitly disabled redaction via header return False - possible_enable_headers: Final = [ - "litellm-enable-message-redaction", # old header. maintain backwards compatibility - "x-litellm-enable-message-redaction", # new header - ] - is_redaction_enabled_via_header = False - for header in possible_enable_headers: + for header in MESSAGE_REDACTION_ENABLE_HEADERS: if bool(request_headers.get(header, False)): is_redaction_enabled_via_header = True break diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index d3c52269d54..ab39a84d9a4 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -19,6 +19,7 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost +from litellm.litellm_core_utils.redact_messages import MESSAGE_REDACTION_ENABLE_HEADERS from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( get_key_object, @@ -117,16 +118,24 @@ def _failure_snapshot_with_guardrails_recorded_after_it( return {**failure_snapshot, "guardrail_information": recorded_guardrail_information} -def _failure_metadata_lifted_from_litellm_metadata( +def _guardrail_info_lifted_from_litellm_metadata( litellm_metadata: ObjectMapping, *, guardrail_info_already_set: bool ) -> ObjectMapping: guardrail_info: Final = litellm_metadata.get("standard_logging_guardrail_information") + if guardrail_info is None or guardrail_info_already_set: + return {} + return {"standard_logging_guardrail_information": guardrail_info} + + +def _redaction_opt_in_lifted_from_litellm_metadata( + litellm_metadata: ObjectMapping, *, caller_set_headers: bool +) -> ObjectMapping: request_headers: Final = litellm_metadata.get("headers") - keeps_guardrail_info: Final = guardrail_info is not None and not guardrail_info_already_set - return { - **({"standard_logging_guardrail_information": guardrail_info} if keeps_guardrail_info else {}), - **({"headers": request_headers} if isinstance(request_headers, dict) else {}), - } + if caller_set_headers or not isinstance(request_headers, dict): + return {} + headers: Final = _OBJECT_MAPPING.validate_python(request_headers) + opt_in_headers: Final = {name: headers[name] for name in MESSAGE_REDACTION_ENABLE_HEADERS if name in headers} + return {"headers": opt_in_headers} if opt_in_headers else {} def _proxy_spend_writer() -> DBSpendUpdateWriter: @@ -243,11 +252,14 @@ class _ProxyDBLogger(CustomLogger): existing_metadata.update(_metadata) litellm_metadata_bucket: Final = request_data.get("litellm_metadata") + litellm_metadata: Final[ObjectMapping] = ( + _OBJECT_MAPPING.validate_python(litellm_metadata_bucket) + if isinstance(litellm_metadata_bucket, dict) + else {} + ) existing_metadata.update( - _failure_metadata_lifted_from_litellm_metadata( - _OBJECT_MAPPING.validate_python(litellm_metadata_bucket) - if isinstance(litellm_metadata_bucket, dict) - else {}, + _guardrail_info_lifted_from_litellm_metadata( + litellm_metadata, guardrail_info_already_set="standard_logging_guardrail_information" in existing_metadata, ) ) @@ -265,7 +277,12 @@ class _ProxyDBLogger(CustomLogger): request_data["litellm_params"]["proxy_server_request"] = ( request_data.get("proxy_server_request") or existing_litellm_params.get("proxy_server_request") or {} ) - request_data["litellm_params"]["metadata"] = existing_metadata + request_data["litellm_params"]["metadata"] = { + **existing_metadata, + **_redaction_opt_in_lifted_from_litellm_metadata( + litellm_metadata, caller_set_headers="headers" in existing_metadata + ), + } # Preserve model name and custom_llm_provider if "model" not in request_data: diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index f676914fe88..a470813fad4 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -9,6 +9,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest +import litellm from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import SpendLogsPayload, UserAPIKeyAuth @@ -163,27 +164,53 @@ async def test_async_post_call_failure_hook_does_not_clobber_guardrail_info_in_m @pytest.mark.asyncio @pytest.mark.parametrize( - "request_headers, code_word_is_stored", - [({"x-litellm-enable-message-redaction": "true"}, False), ({}, True)], + "request_headers, caller_metadata, redaction_is_global, code_word_is_stored", + [ + ({"x-litellm-enable-message-redaction": "true"}, {"user_id": "claude-code-session"}, False, False), + ({"litellm-enable-message-redaction": "true"}, {}, False, False), + ({"user-agent": "claude-cli"}, {"user_id": "claude-code-session"}, False, True), + ({"user-agent": "claude-cli"}, {"headers": {"x-litellm-enable-message-redaction": "true"}}, False, False), + ({"litellm-disable-message-redaction": "true"}, {"user_id": "claude-code-session"}, True, False), + ], + ids=[ + "enable_header_redacts", + "legacy_enable_header_redacts", + "no_header_stores_prompt", + "caller_metadata_headers_still_apply", + "disable_header_never_loosens_global_redaction", + ], ) async def test_async_post_call_failure_hook_honors_the_redaction_header_on_a_litellm_metadata_route( - monkeypatch: pytest.MonkeyPatch, request_headers: dict[str, str], code_word_is_stored: bool + monkeypatch: pytest.MonkeyPatch, + request_headers: dict[str, str], + caller_metadata: dict[str, object], + redaction_is_global: bool, + code_word_is_stored: bool, ): monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + monkeypatch.setattr(litellm, "turn_off_message_logging", redaction_is_global) writer: Final = MagicMock(spec=DBSpendUpdateWriter) writer.update_database = AsyncMock() logger: Final = ProxyDBLogger(spend_writer=lambda: writer) request_input: Final = "The code word is VERMILION-42. Reply OK." + body_metadata: Final = dict(caller_metadata) + litellm_metadata_headers: Final = dict(request_headers) await logger.async_post_call_failure_hook( request_data={ "model": "gpt-5.6", "input": request_input, - "metadata": {"caller_tag": "kept"}, - "litellm_metadata": {"headers": request_headers}, + "metadata": body_metadata, + "litellm_metadata": {"headers": litellm_metadata_headers}, "proxy_server_request": { "url": "http://localhost/v1/responses", - "body": {"model": "gpt-5.6", "input": request_input, "temperature": 7}, + "body": { + "model": "gpt-5.6", + "input": request_input, + "temperature": 7, + "metadata": body_metadata, + "litellm_metadata": {"headers": litellm_metadata_headers}, + }, }, }, original_exception=Exception("temperature must be at most 2"), @@ -199,6 +226,8 @@ async def test_async_post_call_failure_hook_honors_the_redaction_header_on_a_lit stored_request: Final = json.loads(row["proxy_server_request"]) assert stored_request["temperature"] == 7 assert ("VERMILION-42" in row["proxy_server_request"]) is code_word_is_stored + assert stored_request["metadata"].get("headers") == caller_metadata.get("headers") + assert stored_request["litellm_metadata"]["headers"] == request_headers @pytest.mark.asyncio From c5c7c114886842bf44cb764b2c6b9be544e53e92 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 02:40:50 -0700 Subject: [PATCH 11/21] test(responses): cover failure-row redaction headers and caller metadata on litellm_metadata routes --- .../test_responses_compaction_items_wire.py | 129 ++++++++++++++++-- 1 file changed, 120 insertions(+), 9 deletions(-) diff --git a/tests/integration/providers/test_responses_compaction_items_wire.py b/tests/integration/providers/test_responses_compaction_items_wire.py index 34a681d9e74..6441a8c68cf 100644 --- a/tests/integration/providers/test_responses_compaction_items_wire.py +++ b/tests/integration/providers/test_responses_compaction_items_wire.py @@ -47,6 +47,10 @@ _MASK_INPUT: Final = "compaction-mask-input" _MASK_OUTPUT: Final = "compaction-mask-output" _BLOCK_OUTPUT: Final = "compaction-block-output" _REDACT_HEADER: Final = {"x-litellm-enable-message-redaction": "true"} +_UNREDACT_HEADER: Final = {"litellm-disable-message-redaction": "true"} +_SCRIPTED_REJECTION: Final = json.dumps( + {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted rejection"}} +).encode() _ASK: Final = re.compile(r"ask-([0-9a-f]{32})") Route: TypeAlias = Literal["claude-compacting", "claude-plain", "gemini", "openai", "openai-compact", "openai-opaque"] @@ -60,6 +64,10 @@ def _ask(marker: str, *, compact: bool = False) -> str: return f"ask-{marker} compact-now-{marker}" if compact else f"ask-{marker}" +def _fail(marker: str) -> str: + return f"fail-now-{marker}" + + def _think(marker: str) -> str: return f"think-now-{marker}" @@ -164,6 +172,8 @@ def _anthropic_stream(identity: str, content: Sequence[Mapping[str, JsonValue]]) def _claude_reply(body: Mapping[str, JsonValue], raw: bytes) -> Reply: marker: Final = _newest_marker(raw) + if _fail(marker).encode() in raw: + return Reply(status=400, body=_SCRIPTED_REJECTION) compaction: Final = (_block(_summary(marker)),) if f"compact-now-{marker}".encode() in raw else () thinking: Final = (_thinking(marker),) if _think(marker).encode() in raw else () content: Final = (*compaction, *thinking, {"type": "text", "text": _answer(marker)}) @@ -645,7 +655,7 @@ _STORED_OPENAI: Final = "compaction-stored-openai" _SMALL_CONTEXT: Final = 2000 -def _stored_config(directory: Path, provider_url: str, guardrail_url: str) -> Path: +def _stored_config(directory: Path, provider_url: str, guardrail_url: str, **litellm_settings: JsonValue) -> Path: stock: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) claude: Final = {"api_key": cc.ANTHROPIC_API_KEY} config: Final[dict[str, JsonValue]] = { @@ -664,6 +674,7 @@ def _stored_config(directory: Path, provider_url: str, guardrail_url: str) -> Pa "store_prompts_in_spend_logs": True, "disable_responses_id_security": True, }, + "litellm_settings": {**object_value(stock["litellm_settings"]), **litellm_settings}, "router_settings": {**object_value(stock["router_settings"]), "num_retries": 0, "enable_pre_call_checks": True}, "guardrails": [ _guardrail(_MASK_INPUT, "pre_call", f"{guardrail_url}/input"), @@ -683,22 +694,28 @@ class _StoredRig: guardrail: Wire -@pytest.fixture(scope="module") -def stored(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_StoredRig]: - directory: Final = tmp_path_factory.mktemp("compaction-stored") +def _stored_rig(directory: Path, **litellm_settings: JsonValue) -> Iterator[_StoredRig]: with ( wire_server(_vendor) as provider, wire_server(_guardrail_peer) as guardrail, gateway_from_environment() as gateway, ): - config: Final = _stored_config(directory, provider.url, guardrail.url) + config: Final = _stored_config(directory, provider.url, guardrail.url, **litellm_settings) with owned_proxy(gateway, directory, {}, config=config, workers=2) as owned: yield _StoredRig(owned, provider, guardrail) -def _spend_row(response_id: str) -> dict[str, JsonValue]: - decoded: Final = ResponsesAPIRequestUtils.decode_responses_api_response_id(response_id).get("response_id") - request_id: Final = decoded if isinstance(decoded, str) else response_id +@pytest.fixture(scope="module") +def stored(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_StoredRig]: + yield from _stored_rig(tmp_path_factory.mktemp("compaction-stored")) + + +@pytest.fixture(scope="module") +def redacting(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_StoredRig]: + yield from _stored_rig(tmp_path_factory.mktemp("compaction-redacting"), turn_off_message_logging=True) + + +def _logged_row(request_id: str, status: Literal["success", "failure"]) -> dict[str, JsonValue]: rows: Final = eventually( lambda: read_rows( 'SELECT request_id, proxy_server_request, response, metadata, status FROM "LiteLLM_SpendLogs"' @@ -708,10 +725,15 @@ def _spend_row(response_id: str) -> dict[str, JsonValue]: lambda values: len(values) == 1, seconds=70, ) - assert rows[0]["status"] == "success", rows[0] + assert rows[0]["status"] == status, rows[0] return rows[0] +def _spend_row(response_id: str) -> dict[str, JsonValue]: + decoded: Final = ResponsesAPIRequestUtils.decode_responses_api_response_id(response_id).get("response_id") + return _logged_row(decoded if isinstance(decoded, str) else response_id, "success") + + _REPLAY_ROUTES: Final = ( pytest.param(_STORED_CLAUDE, True, id="claude-compacting"), pytest.param(_STORED_CLAUDE, False, id="claude-plain"), @@ -944,3 +966,92 @@ def test_redacted_chat_stream_spend_row_holds_no_compaction_delta(stored: _Store leaked: Final = {column: text for column, text in logged.items() if any(secret in text for secret in secrets)} assert not leaked, {column: text[max(0, text.find(secrets[0]) - 600) :][:1500] for column, text in leaked.items()} _only_call(stored.provider, marker) + + +def _failing_request(path: str, marker: str) -> dict[str, JsonValue]: + ask: Final = f"{_ask(marker)} {_fail(marker)}" + if path == "/v1/responses": + return { + "model": _STORED_CLAUDE, + "input": [_compaction_item(_summary(marker)), {"role": "user", "content": ask}], + } + return {"model": _STORED_CLAUDE, "max_tokens": 64, "messages": [{"role": "user", "content": ask}]} + + +def _failed_call_row( + rig: _StoredRig, + path: str, + marker: str, + headers: Mapping[str, str], + key: str | None = None, + **fields: JsonValue, +) -> dict[str, JsonValue]: + body: Final = {**_failing_request(path, marker), **fields} + response: Final = rig.gateway.request("POST", path, body, headers=headers, key=key) + assert response.status_code == 400, response.text + _only_call(rig.provider, marker) + return _logged_row(response.headers["x-litellm-call-id"], "failure") + + +_LEGACY_REDACT_HEADER: Final = {"litellm-enable-message-redaction": "true"} +_REDACTED_FAILURES: Final = ( + pytest.param("/v1/responses", _REDACT_HEADER, id="responses"), + pytest.param("/v1/responses", _LEGACY_REDACT_HEADER, id="responses-legacy-header"), + pytest.param("/v1/messages", _REDACT_HEADER, id="messages"), + pytest.param("/v1/chat/completions", _REDACT_HEADER, id="chat"), +) + + +@pytest.mark.parametrize(("path", "headers"), _REDACTED_FAILURES) +def test_failed_call_spend_row_holds_no_request_text_with_the_redaction_header( + stored: _StoredRig, path: str, headers: Mapping[str, str] +) -> None: + marker: Final = _new_marker() + logged: Final = json.dumps(_failed_call_row(stored, path, marker, headers)) + assert marker not in logged, logged[max(0, logged.find(marker) - 600) :][:1500] + + +def test_failed_responses_call_spend_row_keeps_the_whole_compaction_summary_without_the_header( + stored: _StoredRig, +) -> None: + marker: Final = _new_marker() + row: Final = _failed_call_row(stored, "/v1/responses", marker, {}) + stored_input: Final = _objects(object_value(row["proxy_server_request"])["input"]) + assert _decoded(string_value(stored_input[0]["encrypted_content"])) == _block(_summary(marker)), stored_input[0] + assert f"{_ask(marker)} {_fail(marker)}" in _message_texts(stored_input[1]), stored_input[1] + + +def test_failed_responses_call_stores_the_caller_metadata_without_the_request_headers(stored: _StoredRig) -> None: + marker: Final = _new_marker() + row: Final = _failed_call_row(stored, "/v1/responses", marker, _REDACT_HEADER, metadata={"caller": marker}) + stored_metadata: Final = object_value(object_value(row["proxy_server_request"])["metadata"]) + assert stored_metadata.get("caller") == marker, stored_metadata + assert "headers" not in stored_metadata, stored_metadata + + +def test_failed_responses_call_keeps_the_caller_metadata_headers_over_the_redaction_header(stored: _StoredRig) -> None: + marker: Final = _new_marker() + caller_headers: Final = {"caller": marker} + row: Final = _failed_call_row(stored, "/v1/responses", marker, _REDACT_HEADER, metadata={"headers": caller_headers}) + stored_request: Final = object_value(row["proxy_server_request"]) + assert object_value(stored_request["metadata"])["headers"] == caller_headers, stored_request["metadata"] + assert _ask(marker) in json.dumps(stored_request), json.dumps(stored_request)[:1500] + + +_GLOBALLY_REDACTED_FAILURES: Final = ( + pytest.param("/v1/responses", _UNREDACT_HEADER, False, id="responses-disable-header"), + pytest.param("/v1/chat/completions", _UNREDACT_HEADER, True, id="chat-disable-header"), + pytest.param("/v1/responses", {}, False, id="responses-no-header"), +) + + +@pytest.mark.parametrize(("path", "headers", "kept"), _GLOBALLY_REDACTED_FAILURES) +def test_opted_out_key_failed_call_spend_row_takes_the_disable_header_over_global_redaction_only_on_chat( + redacting: _StoredRig, path: str, headers: Mapping[str, str], kept: bool +) -> None: + marker: Final = _new_marker() + with redacting.gateway.scenario() as scenario: + key: Final = scenario.key(metadata={"allow_client_message_redaction_opt_out": True}) + row: Final = _failed_call_row(redacting, path, marker, headers, key) + stored_request: Final = json.dumps(row["proxy_server_request"]) + assert (_ask(marker) in stored_request) is kept, stored_request[:1500] From 2b6b3558325e96da58d9e7e7c82c53e970d911ca Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 14:01:57 -0700 Subject: [PATCH 12/21] fix(responses): keep earlier-turn instructions after compaction and count compaction items by their summary A previous_response_id rebuild that starts at a compaction now keeps the system and developer messages of every earlier turn. Token estimates for raw Responses input count a compaction item as its summary text instead of the serialized item. --- .../guardrail_translation/handler.py | 6 +- .../hooks/parallel_request_limiter_v3.py | 4 +- litellm/proxy/spend_tracking/input_tokens.py | 6 +- .../compaction_items.py | 18 +- .../session_handler.py | 41 ++++- tests/unit/proxy/hooks/test_tpm_concurrent.py | 16 ++ .../proxy/spend_tracking/test_input_tokens.py | 23 +++ .../test_session_handler.py | 170 ++++++++++-------- 8 files changed, 189 insertions(+), 95 deletions(-) diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index bda2eab26b4..19c5750cb5e 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -1520,11 +1520,9 @@ class OpenAIResponsesHandler(BaseTranslation): message_items: Final = tuple( item for item in output_items if stream_item_field(item, "type") not in _OUTPUT_TOOL_CALL_ITEM_TYPES ) + message_parts: Final = chain.from_iterable(stream_item_items(item, "content") for item in message_items) part_texts: Final = tuple( - text - for item in message_items - for part in stream_item_items(item, "content") - if isinstance(text := stream_item_field(part, "text"), str) and text + text for part in message_parts if isinstance(text := stream_item_field(part, "text"), str) and text ) summaries: Final = tuple(summary for item in output_items if (summary := _output_compaction_summary(item))) return StreamingScanKey( diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 0a8b0c9ec05..e92c78cff25 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -71,7 +71,7 @@ from litellm.proxy.hooks.batch_enqueued_tokens import ( canonical_provider_batch_id, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit -from litellm.responses.litellm_completion_transformation.compaction_items import input_since_newest_compaction +from litellm.responses.litellm_completion_transformation.compaction_items import billable_input_since_newest_compaction from litellm.router_utils.add_retry_fallback_headers import ( ensure_response_additional_headers, response_has_hidden_params, @@ -1228,7 +1228,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) messages, prompt, raw_input = selected_fields input_text: Final = ( - input_since_newest_compaction(raw_input) if call_type in RESPONSES_API_CALL_TYPES else raw_input + billable_input_since_newest_compaction(raw_input) if call_type in RESPONSES_API_CALL_TYPES else raw_input ) total_chars: Final = ( diff --git a/litellm/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 520b5661438..927152296e9 100644 --- a/litellm/proxy/spend_tracking/input_tokens.py +++ b/litellm/proxy/spend_tracking/input_tokens.py @@ -20,8 +20,8 @@ from typing import Final import litellm from litellm._logging import verbose_proxy_logger from litellm.responses.litellm_completion_transformation.compaction_items import ( + billable_input_since_newest_compaction, holds_compaction_summary, - input_since_newest_compaction, ) from litellm.rust_bridge import runtime from litellm.rust_bridge.catalog import Route, RouteContext @@ -145,7 +145,9 @@ def count_input_tokens_for_model(request_body: dict, model: str) -> int | None: if "prompt" in request_body: return _count_text_tokens(model=model, text=request_body.get("prompt")) if "input" in request_body: - return _count_text_tokens(model=model, text=input_since_newest_compaction(request_body.get("input"))) + return _count_text_tokens( + model=model, text=billable_input_since_newest_compaction(request_body.get("input")) + ) if "query" in request_body or "documents" in request_body: query_tokens: Final = _count_text_tokens(model=model, text=request_body.get("query")) document_tokens: Final = _count_text_tokens( diff --git a/litellm/responses/litellm_completion_transformation/compaction_items.py b/litellm/responses/litellm_completion_transformation/compaction_items.py index a99b43c5295..d0f65bc1bd6 100644 --- a/litellm/responses/litellm_completion_transformation/compaction_items.py +++ b/litellm/responses/litellm_completion_transformation/compaction_items.py @@ -169,12 +169,15 @@ def _is_instruction_item(item: object) -> bool: return fields is not None and fields.get("role") in _INSTRUCTION_ROLES -def instruction_items(input: object) -> tuple[object, ...]: +def input_items(input: object) -> tuple[object, ...]: try: - items: Final = _OBJECT_LIST.validate_python(input) + return tuple(_OBJECT_LIST.validate_python(input)) except ValidationError: return () - return tuple(item for item in items if _is_instruction_item(item)) + + +def instruction_items(input: object) -> tuple[object, ...]: + return tuple(item for item in input_items(input) if _is_instruction_item(item)) def _items_from_newest_compaction(input: Sequence[_Item]) -> tuple[_Item, ...]: @@ -188,12 +191,17 @@ def without_items_before_compaction(input: str | ResponseInputParam) -> str | Re return list(_items_from_newest_compaction(input)) if isinstance(input, list) else input -def input_since_newest_compaction(input: object) -> object: +def _summary_or_item(item: object) -> object: + summary: Final = compaction_summary(item) + return item if summary is None else summary + + +def billable_input_since_newest_compaction(input: object) -> object: try: items: Final = _OBJECT_LIST.validate_python(input) except ValidationError: return input - return list(_items_from_newest_compaction(items)) + return [_summary_or_item(item) for item in _items_from_newest_compaction(items)] def _as_summary_message(item: ResponseInputItemParam) -> ResponseInputItemParam: diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 917779a6566..4ea0bc319c1 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -2,6 +2,7 @@ import asyncio import json from collections.abc import Mapping, Sequence from dataclasses import dataclass +from itertools import chain from typing import TYPE_CHECKING, Any, Final, cast import litellm @@ -10,6 +11,7 @@ from litellm.constants import REDACTED_BY_LITELLM, REDACTED_TOOL_CALL_ARGUMENTS_ from litellm.litellm_core_utils.cold_storage_handler import ColdStorageHandler from litellm.responses.litellm_completion_transformation.compaction_items import ( holds_compaction_summary, + input_items, instruction_items, response_holds_compaction_block, ) @@ -102,15 +104,35 @@ def _replay_start( def _replayed_request( - index: int, proxy_server_request: Mapping[str, object] | None, start: _ReplayStart + index: int, + proxy_server_request: Mapping[str, object] | None, + start: _ReplayStart, + earlier_instructions: tuple[object, ...], ) -> Mapping[str, object] | None: - if proxy_server_request is None or not start.at_output or index != start.index: + if proxy_server_request is None or index != start.index: return proxy_server_request - return { - **proxy_server_request, - "input": list(instruction_items(proxy_server_request.get("input"))), - "messages": None, - } + stored_input: Final = proxy_server_request.get("input") + if start.at_output: + return { + **proxy_server_request, + "input": [*earlier_instructions, *instruction_items(stored_input)], + "messages": None, + } + if not earlier_instructions: + return proxy_server_request + return {**proxy_server_request, "input": [*earlier_instructions, *input_items(stored_input)]} + + +def _instruction_items_before( + proxy_server_requests: Sequence[Mapping[str, object] | None], start: _ReplayStart +) -> tuple[object, ...]: + return tuple( + chain.from_iterable( + instruction_items(request.get("input")) + for request in proxy_server_requests[: start.index] + if request is not None + ) + ) def _stored_instructions(proxy_server_request: Mapping[str, object] | None) -> str | None: @@ -158,11 +180,14 @@ class ResponsesSessionHandler: for spend_log in all_spend_logs ] start: Final = _replay_start(all_spend_logs, proxy_server_requests) + earlier_instructions: Final = _instruction_items_before(proxy_server_requests, start) for index in range(start.index, len(all_spend_logs)): chat_completion_message_history = ( ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload( spend_log=all_spend_logs[index], - proxy_server_request_dict=_replayed_request(index, proxy_server_requests[index], start), + proxy_server_request_dict=_replayed_request( + index, proxy_server_requests[index], start, earlier_instructions + ), chat_completion_message_history=chat_completion_message_history, item_cache_control=item_cache_control, ) diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index 9419ebbca4e..3d1e494e36c 100644 --- a/tests/unit/proxy/hooks/test_tpm_concurrent.py +++ b/tests/unit/proxy/hooks/test_tpm_concurrent.py @@ -3759,3 +3759,19 @@ def test_reservation_skips_the_history_a_compaction_item_replaces(rate_limiter): ) == handler._estimate_precise_input_tokens( {**compacted_request, "input": _SUMMARY_INPUT}, model="gpt-4o", call_type="aresponses" ) + + +def test_reservation_counts_a_compaction_item_by_its_summary(rate_limiter): + handler, _cache = rate_limiter + padded_item: Final = { + **_COMPACTED_INPUT[2], + "encrypted_content": json.dumps( + {"type": "compaction", "content": _COMPACTION_SUMMARY, "opaque_state": "x" * 4000} + ), + } + + assert handler._estimate_input_and_output_tokens( + data={"model": "gpt-4o", "input": [padded_item, _NEXT_TURN], "max_output_tokens": 1}, call_type="aresponses" + ) == handler._estimate_input_and_output_tokens( + data={"model": "gpt-4o", "input": _COMPACTED_INPUT[2:], "max_output_tokens": 1}, call_type="aresponses" + ) diff --git a/tests/unit/proxy/spend_tracking/test_input_tokens.py b/tests/unit/proxy/spend_tracking/test_input_tokens.py index bfec1c10212..da8c8d001a7 100644 --- a/tests/unit/proxy/spend_tracking/test_input_tokens.py +++ b/tests/unit/proxy/spend_tracking/test_input_tokens.py @@ -56,3 +56,26 @@ async def test_history_before_a_compaction_item_is_not_counted(model: str) -> No assert full_counts[model] == compacted_counts[model] assert count_input_tokens_for_model(request_body=full_body, model=model) == compacted_counts[model] + + +def _compaction_item_carrying(extra_block_fields: dict[str, str]) -> dict[str, str]: + return { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps( + {"type": "compaction", "content": "Turn one planned Tokyo.", **extra_block_fields} + ), + } + + +@pytest.mark.parametrize("model", [CL100K_MODEL, "claude-sonnet-5-5"]) +def test_a_compaction_item_counts_only_its_summary(model: str) -> None: + plain_body: Final = {"model": model, "input": [_compaction_item_carrying({}), _TURN_AFTER_COMPACTION]} + padded_body: Final = { + "model": model, + "input": [_compaction_item_carrying({"opaque_state": "x" * 4000}), _TURN_AFTER_COMPACTION], + } + + assert count_input_tokens_for_model(request_body=padded_body, model=model) == count_input_tokens_for_model( + request_body=plain_body, model=model + ) diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index d03b23a1e3f..d1deae806f6 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -18,6 +18,7 @@ from litellm.responses.litellm_completion_transformation.session_handler import ResponsesSessionHandler, _normalize_redacted_tool_call_arguments, ) +from litellm.responses.litellm_completion_transformation.transformation import ChatCompletionSession from litellm.responses.mcp.litellm_proxy_mcp_handler import LiteLLM_Proxy_MCP_Handler from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ResponsesAPIResponse @@ -233,11 +234,10 @@ async def test_e2e_cold_storage_successful_retrieval(): patch.object(session_handler, "COLD_STORAGE_HANDLER") as mock_cold_storage, patch("litellm.cold_storage_custom_logger", return_value="s3"), ): - # Setup mocks mock_get_spend_logs.return_value = mock_spend_logs - mock_cold_storage.get_proxy_server_request_from_cold_storage_with_object_key = ( - AsyncMock(return_value=full_proxy_request) + mock_cold_storage.get_proxy_server_request_from_cold_storage_with_object_key = AsyncMock( + return_value=full_proxy_request ) # Call the main function @@ -291,7 +291,6 @@ async def test_e2e_cold_storage_fallback_to_truncated_payload(): ) as mock_get_spend_logs, patch.object(session_handler, "COLD_STORAGE_HANDLER") as mock_cold_storage, ): - # Setup mocks mock_get_spend_logs.return_value = mock_spend_logs @@ -358,41 +357,25 @@ async def test_should_check_cold_storage_for_full_payload(): with patch("litellm.cold_storage_custom_logger", return_value="s3"): # Test case 1: Should return True for truncated content - result1 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_with_truncated_pdf - ) - assert ( - result1 == True - ), "Should return True for proxy request with truncated PDF content" + result1 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf) + assert result1 == True, "Should return True for proxy request with truncated PDF content" # Test case 2: Should return False for regular content - result2 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_regular - ) - assert ( - result2 == False - ), "Should return False for regular proxy request without truncation" + result2 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_regular) + assert result2 == False, "Should return False for regular proxy request without truncation" # Test case 3: Should return True for empty request - result3 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_empty - ) + result3 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_empty) assert result3 == True, "Should return True for empty proxy request" # Test case 4: Should return True for None request - result4 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_none - ) + result4 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_none) assert result4 == True, "Should return True for None proxy request" # Test case 5: Should return False when cold storage is not configured with patch.object(litellm, "cold_storage_custom_logger", None): - result5 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload( - proxy_request_with_truncated_pdf - ) - assert ( - result5 == False - ), "Should return False when cold storage is not configured, even with truncated content" + result5 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf) + assert result5 == False, "Should return False when cold storage is not configured, even with truncated content" @pytest.mark.asyncio @@ -422,9 +405,7 @@ async def test_get_chat_completion_message_history_empty_response_dict(): } ] - with patch.object( - ResponsesSessionHandler, "get_all_spend_logs_for_previous_response_id" - ) as mock_get_spend_logs: + with patch.object(ResponsesSessionHandler, "get_all_spend_logs_for_previous_response_id") as mock_get_spend_logs: mock_get_spend_logs.return_value = mock_spend_logs # Call the function @@ -518,9 +499,7 @@ async def test_message_history_reconstructs_list_shaped_input(): ) as mock_get_spend_logs: mock_get_spend_logs.return_value = mock_spend_logs - result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( - request_id - ) + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(request_id) messages = result["messages"] assert [(message.get("role"), message.get("content")) for message in messages] == [ @@ -549,9 +528,7 @@ async def test_message_history_retries_a_spend_log_the_batch_writer_has_not_flus fake_prisma_client = _FakePrismaClient(results=[[], [spend_log]]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): - result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( - request_id - ) + result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(request_id) messages = result["messages"] assert [(message.get("role"), message.get("content")) for message in messages] == [ @@ -618,14 +595,10 @@ async def test_message_history_looks_up_the_decoded_chat_completion_id(): model_id="e0f302a1412e78470ebb28cbed01fff5f88c0d331c667e9f2ba4b413c6fbd282", response_id=request_id, ) - fake_prisma_client = _FakePrismaClient( - results=[[_spend_log(request_id, "session-a", "Hello.", "Hi.")]] - ) + fake_prisma_client = _FakePrismaClient(results=[[_spend_log(request_id, "session-a", "Hello.", "Hi.")]]) with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client): - await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( - encoded_response_id - ) + await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(encoded_response_id) assert len(fake_prisma_client.db.calls) == 1 assert request_id in fake_prisma_client.db.calls[0] @@ -724,8 +697,9 @@ async def test_session_lookup_does_not_retry_when_spend_logs_are_disabled( """ fake_prisma_client = _FakePrismaClient(results=[]) - with patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client), patch( - "litellm.proxy.proxy_server.disable_spend_logs", True + with ( + patch("litellm.proxy.proxy_server.prisma_client", fake_prisma_client), + patch("litellm.proxy.proxy_server.disable_spend_logs", True), ): spend_logs = await ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( "chatcmpl-does-not-exist" @@ -741,7 +715,11 @@ def test_normalize_redacted_arguments_skips_custom_tool_calls(): content=None, tool_calls=[ {"id": "call_c", "type": "custom", "custom": {"name": "run_code", "input": "print(1)"}}, - {"id": "call_f", "type": "function", "function": {"name": "get_weather", "arguments": "redacted-by-litellm"}}, + { + "id": "call_f", + "type": "function", + "function": {"name": "get_weather", "arguments": "redacted-by-litellm"}, + }, ], ) @@ -796,11 +774,13 @@ async def test_message_history_normalizes_redacted_tool_call_arguments(): } ] - with patch.object( # test-quality-ok: the handler has no DI seam for the spend-log fetch; every test in this file stubs this same boundary - ResponsesSessionHandler, - "get_all_spend_logs_for_previous_response_id", - new_callable=AsyncMock, - ) as mock_get_spend_logs: + with ( + patch.object( # test-quality-ok: the handler has no DI seam for the spend-log fetch; every test in this file stubs this same boundary + ResponsesSessionHandler, + "get_all_spend_logs_for_previous_response_id", + new_callable=AsyncMock, + ) as mock_get_spend_logs + ): mock_get_spend_logs.return_value = mock_spend_logs result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( @@ -902,6 +882,15 @@ def _role_and_content(message: object) -> tuple[object, object]: return getattr(message, "role", None), getattr(message, "content", None) +async def _history_from_stored_turns( + spend_logs: tuple[dict[str, object], ...], previous_response_id: str +) -> ChatCompletionSession: + with patch("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(db=_StoredSpendLogs(spend_logs))): + return await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( + previous_response_id + ) + + @pytest.mark.asyncio async def test_previous_response_history_starts_at_the_newest_turn_that_sent_a_compaction_item(): first_turn: Final = [{"role": "user", "content": "Plan a Tokyo trip."}] @@ -916,20 +905,12 @@ async def test_previous_response_history_starts_at_the_newest_turn_that_sent_a_c compaction_item, {"role": "user", "content": "And Osaka?"}, ] - spend_logs: Final = [ + spend_logs: Final = ( _turn_spend_log(first_turn, "Day one, Asakusa."), _turn_spend_log(second_turn, "Dotonbori."), - ] + ) - with patch.object( - ResponsesSessionHandler, - "get_all_spend_logs_for_previous_response_id", - new_callable=AsyncMock, - return_value=spend_logs, - ): - session: Final = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( - "resp_4" - ) + session: Final = await _history_from_stored_turns(spend_logs, "resp_4") assert [_role_and_content(message) for message in session["messages"]] == [ ("user", "Turn one planned Tokyo."), @@ -941,7 +922,7 @@ async def test_previous_response_history_starts_at_the_newest_turn_that_sent_a_c @pytest.mark.asyncio async def test_previous_response_history_starts_at_the_newest_answer_that_compacted(): compaction_blocks: Final = [{"type": "compaction", "content": "Turn one planned Tokyo."}] - spend_logs: Final = [ + spend_logs: Final = ( _turn_spend_log([{"role": "user", "content": "Hi."}], "Hello."), _turn_spend_log( [{"role": "developer", "content": "Answer briefly."}, {"role": "user", "content": "Plan a Tokyo trip."}], @@ -949,17 +930,9 @@ async def test_previous_response_history_starts_at_the_newest_answer_that_compac compaction_blocks, ), _turn_spend_log([{"role": "user", "content": "And Osaka?"}], "Dotonbori."), - ] + ) - with patch.object( - ResponsesSessionHandler, - "get_all_spend_logs_for_previous_response_id", - new_callable=AsyncMock, - return_value=spend_logs, - ): - session: Final = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id( - "resp_1" - ) + session: Final = await _history_from_stored_turns(spend_logs, "resp_1") assert [_role_and_content(message) for message in session["messages"]] == [ ("developer", "Answer briefly."), @@ -970,6 +943,53 @@ async def test_previous_response_history_starts_at_the_newest_answer_that_compac assert getattr(session["messages"][1], "provider_specific_fields", None) == {"compaction_blocks": compaction_blocks} +_COMPACTION_ITEM_SENT: Final = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Turn one planned Tokyo."}), +} + + +@pytest.mark.parametrize( + "compacting_turn", + [ + pytest.param( + _turn_spend_log( + [{"role": "user", "content": "Plan a Tokyo trip."}], + "Day one, Asakusa.", + [{"type": "compaction", "content": "Turn one planned Tokyo."}], + ), + id="answer_compacted", + ), + pytest.param( + _turn_spend_log( + [_COMPACTION_ITEM_SENT, {"role": "user", "content": "Plan a Tokyo trip."}], "Day one, Asakusa." + ), + id="input_sent_a_compaction_item", + ), + ], +) +@pytest.mark.asyncio +async def test_previous_response_history_keeps_instructions_from_turns_the_compaction_covers( + compacting_turn: dict[str, object], +): + spend_logs: Final = ( + _turn_spend_log( + [{"role": "developer", "content": "Use metric units."}, {"role": "user", "content": "Hi."}], "Hello." + ), + compacting_turn, + _turn_spend_log([{"role": "user", "content": "And Osaka?"}], "Dotonbori."), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_1") + + roles_and_contents: Final = [_role_and_content(message) for message in session["messages"]] + assert roles_and_contents[0] == ("developer", "Use metric units.") + assert ("user", "Hi.") not in roles_and_contents + assert ("assistant", "Hello.") not in roles_and_contents + assert roles_and_contents[-2:] == [("user", "And Osaka?"), ("assistant", "Dotonbori.")] + + class _SqliteBackedPrismaDB: def __init__(self, rows: list[tuple[str, str, str, str, str | None, str | None]]) -> None: self._connection = sqlite3.connect(":memory:") @@ -981,7 +1001,9 @@ class _SqliteBackedPrismaDB: self._connection.executemany('INSERT INTO "LiteLLM_SpendLogs" VALUES (?, ?, ?, ?, ?, ?)', rows) async def query_raw(self, query: str, *args: object) -> list[dict[str, object]]: - sqlite_query: Final = query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each(?))").replace("$1", "?") + sqlite_query: Final = query.replace("= ANY($1::text[])", "IN (SELECT value FROM json_each(?))").replace( + "$1", "?" + ) sqlite_args: Final = tuple(json.dumps(arg) if isinstance(arg, list) else arg for arg in args) return [dict(row) for row in self._connection.execute(sqlite_query, sqlite_args)] From a46745198cbb45b552828ff0900e6710b686073b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 18:07:30 -0700 Subject: [PATCH 13/21] fix(responses): drop the stored history a caller-sent compaction item replaces on previous_response_id --- .../handler.py | 2 + .../transformation.py | 14 ++++- .../test_responses_compaction_items_wire.py | 11 ++-- .../test_spend_tracking_utils.py | 10 ++- .../test_anthropic_responses_bridge.py | 61 +++++++++++++++++++ 5 files changed, 86 insertions(+), 12 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 5489d2223ac..e6b6b36f1dc 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -9,6 +9,7 @@ from pydantic import TypeAdapter import litellm from litellm.responses.additional_tools import hoist_additional_tools +from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary from litellm.responses.litellm_completion_transformation.streaming_iterator import ( LiteLLMCompletionStreamingIterator, ) @@ -120,6 +121,7 @@ class LiteLLMCompletionTransformationHandler: instructions=responses_api_request.get("instructions"), request_kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), previous_response_client_ids=previous_response_client_ids, + input_replaces_history=holds_compaction_summary(request_input), ) acompletion_args: Final = {} diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index acc5af778e3..c70d5f69a28 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -316,6 +316,14 @@ def _without_stored_tool_search_uses(message: SessionMessage, search_call_ids: f return message +def _instruction_messages(messages: Sequence[SessionMessage]) -> tuple[SessionMessage, ...]: + return tuple( + message + for message in messages + if isinstance(message, Mapping) and message.get("role") in ("system", "developer") + ) + + def _stored_compaction_summary(message: SessionMessage) -> str | None: return provider_compaction_summary(message.provider_specific_fields) if isinstance(message, Message) else None @@ -1092,6 +1100,7 @@ class LiteLLMCompletionResponsesConfig: session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None = None, request_kwargs: Mapping[str, object] = MappingProxyType({}), previous_response_client_ids: tuple[str, ...] = (), + input_replaces_history: bool = False, ) -> dict: """ Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages @@ -1115,8 +1124,11 @@ class LiteLLMCompletionResponsesConfig: ) _messages: Final = litellm_completion_request.get("messages") or [] route: Final = CompletionRoute(model=model, custom_llm_provider=provider) + session_history: Final = chat_completion_session.get("messages") or [] stored_messages: Final = LiteLLMCompletionResponsesConfig._with_compaction_blocks_the_route_takes( - chat_completion_session.get("messages") or [], route, request_fields.get("context_management") + _instruction_messages(session_history) if input_replaces_history else session_history, + route, + request_fields.get("context_management"), ) replayed_messages: Final = ( LiteLLMCompletionResponsesConfig._without_tool_search_uses_off_anthropic(stored_messages, route) diff --git a/tests/integration/providers/test_responses_compaction_items_wire.py b/tests/integration/providers/test_responses_compaction_items_wire.py index 6441a8c68cf..e975ee24cb3 100644 --- a/tests/integration/providers/test_responses_compaction_items_wire.py +++ b/tests/integration/providers/test_responses_compaction_items_wire.py @@ -157,7 +157,7 @@ def _anthropic_stream(identity: str, content: Sequence[Mapping[str, JsonValue]]) }, ) blocks: Final = tuple( - frame for index, block in enumerate(content) for frame in _anthropic_block_frames(index, block) + chain.from_iterable(_anthropic_block_frames(index, block) for index, block in enumerate(content)) ) end: Final = cc.sse_frame( "message_delta", @@ -283,7 +283,7 @@ def _message_texts(message: Mapping[str, JsonValue], key: str = "content") -> tu def _anthropic_texts(body: Mapping[str, JsonValue], role: str) -> tuple[str, ...]: messages: Final = tuple(message for message in _objects(body["messages"]) if message["role"] == role) - return tuple(text for message in messages for text in _message_texts(message)) + return tuple(chain.from_iterable(map(_message_texts, messages))) def _word_count(body: Mapping[str, JsonValue]) -> int: @@ -293,17 +293,18 @@ def _word_count(body: Mapping[str, JsonValue]) -> int: def _anthropic_block_types(body: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]: contents: Final = tuple(message["content"] for message in _objects(body["messages"])) - return tuple(block.get("type") for content in contents if isinstance(content, list) for block in _objects(content)) + blocks: Final = chain.from_iterable(_objects(content) for content in contents if isinstance(content, list)) + return tuple(block.get("type") for block in blocks) def _gemini_texts(body: Mapping[str, JsonValue], role: str) -> tuple[str, ...]: turns: Final = tuple(turn for turn in _objects(body["contents"]) if turn["role"] == role) - return tuple(text for turn in turns for text in _message_texts(turn, "parts")) + return tuple(chain.from_iterable(_message_texts(turn, "parts") for turn in turns)) def _openai_user_texts(body: Mapping[str, JsonValue]) -> tuple[str, ...]: items: Final = tuple(item for item in _objects(body["input"]) if item.get("role") == "user") - return tuple(text for item in items for text in _message_texts(item)) + return tuple(chain.from_iterable(map(_message_texts, items))) def _openai_item_types(body: Mapping[str, JsonValue]) -> tuple[JsonValue, ...]: diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 7f32831dedd..63a164166fc 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -1055,11 +1055,10 @@ def test_get_response_for_spend_logs_payload_truncates_large_base64(mock_should_ assert parsed["data"][0]["other_field"] == "value" -@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") -def test_get_response_for_spend_logs_payload_keeps_a_long_compaction_summary_whole(mock_should_store): +def test_get_response_for_spend_logs_payload_keeps_a_long_compaction_summary_whole(monkeypatch: pytest.MonkeyPatch): from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB - mock_should_store.return_value = True + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") filler = "detail " * (MAX_STRING_LENGTH_PROMPT_IN_DB // 7) summary = f"{filler}The cat is named Biscuit. {filler}" answer = f"{filler}OK{filler}" @@ -1087,11 +1086,10 @@ def test_get_response_for_spend_logs_payload_keeps_a_long_compaction_summary_who assert LITELLM_TRUNCATED_PAYLOAD_FIELD in message["content"] -@patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") -def test_stored_request_keeps_only_the_summary_of_a_long_claude_compaction_input_item(mock_should_store): +def test_stored_request_keeps_only_the_summary_of_a_long_claude_compaction_input_item(monkeypatch: pytest.MonkeyPatch): from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB - mock_should_store.return_value = True + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") filler = "detail " * (MAX_STRING_LENGTH_PROMPT_IN_DB // 7) summary = f"{filler}The cat is named Biscuit. {filler}" claude_item = json.dumps({"type": "compaction", "content": summary, "padding": filler * 10}) diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 1dc783675ee..18ee62f0853 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -51,6 +51,7 @@ async def test_async_response_api_handler_merges_trace_id_without_error(): instructions: str | None = None, request_kwargs: Mapping[str, object] | None = None, previous_response_client_ids: tuple[str, ...] = (), + input_replaces_history: bool = False, ) -> dict[str, object]: litellm_completion_request["litellm_trace_id"] = "session-trace" return litellm_completion_request @@ -866,6 +867,66 @@ async def test_router_with_encrypted_content_affinity_sends_the_compaction_block assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] +_DOCUMENT_BEFORE_COMPACTION: Final = "The design document pasted before the compaction." +_TURN_BEFORE_COMPACTION: Final = { + "request_id": "chatcmpl-before-compaction", + "call_type": "aresponses", + "session_id": "session-compacted", + "proxy_server_request": { + "model": _MODEL, + "input": [ + {"role": "developer", "content": "Use metric units."}, + {"role": "user", "content": _DOCUMENT_BEFORE_COMPACTION}, + ], + }, + "response": { + "id": "chatcmpl-before-compaction", + "object": "chat.completion", + "created": 0, + "model": _MODEL, + "choices": [{"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "Read it."}}], + }, +} + + +class _RecordingRawAnthropicMessages: + def __init__(self) -> None: + self.body: str = "" + + def __call__(self, request: httpx.Request) -> httpx.Response: + self.body = request.content.decode() + return _anthropic_message({"type": "text", "text": "dog."}) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("context_management", [None, list(_COMPACT_AT_50K)], ids=["summary_text", "native_block"]) +async def test_previous_response_id_with_a_compaction_item_in_input_replays_only_the_stored_instructions( + context_management: list[Mapping[str, object]] | None, +) -> None: + anthropic: Final = _RecordingRawAnthropicMessages() + compaction_item: Final = { + "type": "compaction", + "id": "cmp_caller", + "encrypted_content": json.dumps(_COMPACTION_BLOCK), + } + + with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient([_TURN_BEFORE_COMPACTION])): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-before-compaction", + input=[compaction_item, *_NEXT_INPUT], + context_management=context_management, + api_key="sk-ant-fake", + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + ) + + assert _DOCUMENT_BEFORE_COMPACTION not in anthropic.body + assert "Read it." not in anthropic.body + assert "Use metric units." in anthropic.body + assert _COMPACTION_BLOCK["content"] in anthropic.body + assert _NEXT_INPUT[0]["content"] in anthropic.body + + def test_responses_over_anthropic_send_a_compaction_item_as_summary_text_when_the_request_does_not_compact( respx_mock: respx.MockRouter, ) -> None: From 8f75bb0a6463799a1eb6641f9419c4c44b0e8469 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 18:57:25 -0700 Subject: [PATCH 14/21] fix(responses): count only message items with a string role as compaction instructions --- .../compaction_items.py | 5 ++++- .../proxy/spend_tracking/test_input_tokens.py | 15 ++++++++----- .../test_session_handler.py | 22 +++++++++++++++++++ 3 files changed, 36 insertions(+), 6 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/compaction_items.py b/litellm/responses/litellm_completion_transformation/compaction_items.py index d0f65bc1bd6..0d1ef495de5 100644 --- a/litellm/responses/litellm_completion_transformation/compaction_items.py +++ b/litellm/responses/litellm_completion_transformation/compaction_items.py @@ -166,7 +166,10 @@ def holds_compaction_summary(input: object) -> bool: def _is_instruction_item(item: object) -> bool: fields: Final = _json_object(item) - return fields is not None and fields.get("role") in _INSTRUCTION_ROLES + if fields is None or fields.get("type", "message") != "message": + return False + role: Final = fields.get("role") + return isinstance(role, str) and role in _INSTRUCTION_ROLES def input_items(input: object) -> tuple[object, ...]: diff --git a/tests/unit/proxy/spend_tracking/test_input_tokens.py b/tests/unit/proxy/spend_tracking/test_input_tokens.py index da8c8d001a7..71f4dec2345 100644 --- a/tests/unit/proxy/spend_tracking/test_input_tokens.py +++ b/tests/unit/proxy/spend_tracking/test_input_tokens.py @@ -38,14 +38,19 @@ _COMPACTION_ITEM: Final = { _TURN_AFTER_COMPACTION: Final = {"role": "user", "content": "And Osaka?"} +@pytest.mark.parametrize( + "earlier_item", + [ + pytest.param({"role": "user", "content": "Plan a Tokyo trip. " * 500}, id="user_turn"), + pytest.param({"role": ["system"], "content": "Plan a Tokyo trip."}, id="list_role"), + pytest.param({"role": {"name": "system"}, "content": "Plan a Tokyo trip."}, id="dict_role"), + ], +) @pytest.mark.parametrize("model", [CL100K_MODEL, "claude-sonnet-5-5"]) @pytest.mark.asyncio -async def test_history_before_a_compaction_item_is_not_counted(model: str) -> None: +async def test_history_before_a_compaction_item_is_not_counted(model: str, earlier_item: dict[str, object]) -> None: compacted_body: Final = {"model": model, "input": [_COMPACTION_ITEM, _TURN_AFTER_COMPACTION]} - full_body: Final = { - "model": model, - "input": [{"role": "user", "content": "Plan a Tokyo trip. " * 500}, _COMPACTION_ITEM, _TURN_AFTER_COMPACTION], - } + full_body: Final = {"model": model, "input": [earlier_item, _COMPACTION_ITEM, _TURN_AFTER_COMPACTION]} full_counts: Final = await count_input_tokens( request_body=full_body, raw_body=json.dumps(full_body).encode(), models=(model,) diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index d1deae806f6..3102c878c95 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -990,6 +990,28 @@ async def test_previous_response_history_keeps_instructions_from_turns_the_compa assert roles_and_contents[-2:] == [("user", "And Osaka?"), ("assistant", "Dotonbori.")] +@pytest.mark.asyncio +async def test_previous_response_history_drops_a_tool_output_carrying_a_system_role_the_compaction_covers(): + spend_logs: Final = ( + _turn_spend_log( + [ + {"role": "user", "content": "Look up the code."}, + {"type": "function_call", "call_id": "call_1", "name": "lookup", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_1", "output": "4242", "role": "system"}, + ], + "Found it.", + ), + _turn_spend_log([_COMPACTION_ITEM_SENT, {"role": "user", "content": "Plan a Tokyo trip."}], "Day one."), + _turn_spend_log([{"role": "user", "content": "And Osaka?"}], "Dotonbori."), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_1") + + contents: Final = [str(_role_and_content(message)[1]) for message in session["messages"]] + assert not any("4242" in content for content in contents) + assert contents[-2:] == ["And Osaka?", "Dotonbori."] + + class _SqliteBackedPrismaDB: def __init__(self, rows: list[tuple[str, str, str, str, str | None, str | None]]) -> None: self._connection = sqlite3.connect(":memory:") From ed6b528fcd069200a7807ad64064d6153b447bee Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 19:21:02 -0700 Subject: [PATCH 15/21] test(responses): flatten the last two-for comprehension in the compaction wire test --- .../providers/test_responses_compaction_items_wire.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tests/integration/providers/test_responses_compaction_items_wire.py b/tests/integration/providers/test_responses_compaction_items_wire.py index e975ee24cb3..f309c374645 100644 --- a/tests/integration/providers/test_responses_compaction_items_wire.py +++ b/tests/integration/providers/test_responses_compaction_items_wire.py @@ -485,10 +485,9 @@ def _assert_reached( message for message in _objects(body["messages"]) if message["role"] == "assistant" ) blocks: Final = tuple( - block - for message in assistant - if isinstance(message["content"], list) - for block in _objects(message["content"]) + chain.from_iterable( + _objects(message["content"]) for message in assistant if isinstance(message["content"], list) + ) ) assert _block(summary) in blocks, body["messages"] assert summary not in _anthropic_texts(body, "user") From f008c72eef3c3f15c8e2e92fd02e44eb4c202df8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 16 Oct 2026 21:03:04 -0700 Subject: [PATCH 16/21] fix(responses): replay from a compaction in stored Responses output and SDK compaction items A previous_response_id whose stored row is a Responses output holding a compaction item now starts its replay at that compaction, and the handler dumps SDK compaction objects before checking for a summary. Adds tests for the uncovered compaction paths. --- .../spend_tracking/spend_tracking_utils.py | 8 +- .../compaction_items.py | 8 +- .../handler.py | 3 +- .../session_handler.py | 1 + ...test_openai_responses_guardrail_handler.py | 13 +++ .../test_anthropic_responses_bridge.py | 63 ++++++++++++-- .../test_compaction_items.py | 87 +++++++++++++++++++ .../test_litellm_completion_responses.py | 13 ++- .../test_session_handler.py | 47 +++++++++- .../test_streaming_iterator_transformation.py | 34 ++++++++ .../test_responses_api_request_body.py | 16 +++- 11 files changed, 267 insertions(+), 26 deletions(-) create mode 100644 tests/unit/responses/litellm_completion_transformation/test_compaction_items.py diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index cf204c23bf6..61775d47857 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -1789,13 +1789,9 @@ def _stored_request_without_requester_ip_address(request_body: Mapping[object, o def _with_whole_compaction_summary(item: object, stored_item: object) -> object: summary_only: Final = summary_only_compaction_item(item) - if summary_only is None: + if summary_only is None or not isinstance(stored_item, Mapping): return stored_item - try: - stored_fields: Final = _OBJECT_MAPPING.validate_python(stored_item) - except ValidationError: - return stored_item - return {**stored_fields, "encrypted_content": summary_only} + return {**_OBJECT_MAPPING.validate_python(stored_item), "encrypted_content": summary_only} def _sanitize_proxy_server_request_body(request_body: Mapping[str, object]) -> Mapping[object, object]: diff --git a/litellm/responses/litellm_completion_transformation/compaction_items.py b/litellm/responses/litellm_completion_transformation/compaction_items.py index 0d1ef495de5..75d957ddbca 100644 --- a/litellm/responses/litellm_completion_transformation/compaction_items.py +++ b/litellm/responses/litellm_completion_transformation/compaction_items.py @@ -21,8 +21,6 @@ _INSTRUCTION_ROLES: Final = frozenset(("system", "developer")) def _json_object(value: object) -> Mapping[str, object] | None: - if not isinstance(value, Mapping): - return None try: return _JSON_OBJECT.validate_python(value) except ValidationError: @@ -70,8 +68,12 @@ def _choice_provider_specific_fields(choice: object) -> object: def response_holds_compaction_block(response: object) -> bool: fields: Final = _json_object(response) + if fields is None: + return False + if fields.get("object") == "response": + return holds_compaction_summary(fields.get("output")) try: - choices: Final = _OBJECT_LIST.validate_python(None if fields is None else fields.get("choices")) + choices: Final = _OBJECT_LIST.validate_python(fields.get("choices")) except ValidationError: return False return any(provider_compaction_summary(_choice_provider_specific_fields(choice)) is not None for choice in choices) diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index e6b6b36f1dc..c9d800421f7 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -15,6 +15,7 @@ from litellm.responses.litellm_completion_transformation.streaming_iterator impo ) from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, + with_compaction_items_as_params, ) from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( @@ -121,7 +122,7 @@ class LiteLLMCompletionTransformationHandler: instructions=responses_api_request.get("instructions"), request_kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), previous_response_client_ids=previous_response_client_ids, - input_replaces_history=holds_compaction_summary(request_input), + input_replaces_history=holds_compaction_summary(with_compaction_items_as_params(request_input)), ) acompletion_args: Final = {} diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 4ea0bc319c1..1abc00c477a 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -269,6 +269,7 @@ class ResponsesSessionHandler: responses_api_request={}, replay_reasoning=True, item_cache_control=item_cache_control, + only_since_compaction=True, ) ) elif isinstance(_response_output, dict) and _response_output and _response_output != {}: diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index f5efaa4106b..7b0dc1f9922 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -4201,6 +4201,19 @@ class TestOpenAIResponsesHandlerStreamingScanKey: tool_calls = guardrail.seen_inputs[0].get("tool_calls") or [] assert [call["function"]["arguments"] for call in tool_calls] == ['{"city": "Paris"}'], tool_calls + @pytest.mark.asyncio + async def test_a_tool_call_finished_mid_stream_reaches_the_guardrail_without_texts(self): + guardrail: Final = MockRecordingGuardrail(guardrail_name="test") + await OpenAIResponsesHandler().process_output_streaming_response( + responses_so_far=[self._finished_function_call(1, "fc_1", "Paris")[-1]], + guardrail_to_apply=guardrail, + request_data={}, + ) + assert [ + (sorted(inputs), [call["function"]["arguments"] for call in inputs.get("tool_calls") or []]) + for inputs in guardrail.seen_inputs + ] == [(["tool_calls"], ['{"city": "Paris"}'])] + class TestCompactionOutputSummaryReachesOutputGuardrails: SUMMARY: Final = "The user lives at 1 Secret Lane." diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 18ee62f0853..1c34b92e209 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -1,15 +1,17 @@ import json from collections.abc import Mapping, Sequence from itertools import chain -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx +from openai.types.responses.response_compaction_item import ResponseCompactionItem from pydantic import BaseModel import litellm +from litellm import CustomLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, @@ -900,15 +902,24 @@ class _RecordingRawAnthropicMessages: @pytest.mark.asyncio @pytest.mark.parametrize("context_management", [None, list(_COMPACT_AT_50K)], ids=["summary_text", "native_block"]) +@pytest.mark.parametrize( + "compaction_item", + [ + pytest.param( + {"type": "compaction", "id": "cmp_caller", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + id="param_dict", + ), + pytest.param( + ResponseCompactionItem(type="compaction", id="cmp_caller", encrypted_content=json.dumps(_COMPACTION_BLOCK)), + id="sdk_output_item", + ), + ], +) async def test_previous_response_id_with_a_compaction_item_in_input_replays_only_the_stored_instructions( context_management: list[Mapping[str, object]] | None, + compaction_item: Mapping[str, object] | ResponseCompactionItem, ) -> None: anthropic: Final = _RecordingRawAnthropicMessages() - compaction_item: Final = { - "type": "compaction", - "id": "cmp_caller", - "encrypted_content": json.dumps(_COMPACTION_BLOCK), - } with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient([_TURN_BEFORE_COMPACTION])): await litellm.aresponses( @@ -927,6 +938,46 @@ async def test_previous_response_id_with_a_compaction_item_in_input_replays_only assert _NEXT_INPUT[0]["content"] in anthropic.body +class _RecordingCustomLLM(CustomLLM): + def __init__(self) -> None: + super().__init__() + self.sent_messages: list[object] = [] + + def completion(self, *args: object, **kwargs: object) -> ModelResponse: + self.sent_messages.extend(cast(list[object], kwargs["messages"])) + return litellm.completion( + model="gpt-5.6", messages=[{"role": "user", "content": "Hi."}], mock_response="dog." + ) + + +def test_responses_over_a_custom_provider_send_a_compaction_item_as_summary_text( + monkeypatch: pytest.MonkeyPatch, +) -> None: + custom_llm: Final = _RecordingCustomLLM() + monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list)) + monkeypatch.setattr(litellm, "_custom_providers", list(litellm._custom_providers)) + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "compaction-test-llm", "custom_handler": custom_llm}] + ) + compaction_item: Final = { + "type": "compaction", + "id": "cmp_earlier", + "encrypted_content": json.dumps(_COMPACTION_BLOCK), + } + + litellm.responses( + model="compaction-test-llm/my-model", + input=[*_FIRST_INPUT, {"role": "assistant", "content": "The."}, compaction_item, *_NEXT_INPUT], + context_management=list(_COMPACT_AT_50K), + drop_params=True, + ) + + assert custom_llm.sent_messages == [ + {"role": "user", "content": _COMPACTION_BLOCK["content"]}, + *_NEXT_INPUT, + ] + + def test_responses_over_anthropic_send_a_compaction_item_as_summary_text_when_the_request_does_not_compact( respx_mock: respx.MockRouter, ) -> None: diff --git a/tests/unit/responses/litellm_completion_transformation/test_compaction_items.py b/tests/unit/responses/litellm_completion_transformation/test_compaction_items.py new file mode 100644 index 00000000000..e0ef3fbb552 --- /dev/null +++ b/tests/unit/responses/litellm_completion_transformation/test_compaction_items.py @@ -0,0 +1,87 @@ +import json +from typing import Final + +import pytest + +from litellm.responses.litellm_completion_transformation.compaction_items import ( + compaction_output_item, + holds_compaction_summary, + response_holds_compaction_block, + with_compaction_summary, +) +from litellm.responses.utils import ResponsesAPIRequestUtils + +_BLOCK: Final = {"type": "compaction", "content": "Turn one planned Tokyo."} +_ITEM: Final = {"type": "compaction", "id": "cmp_1", "encrypted_content": json.dumps(_BLOCK)} + + +@pytest.mark.parametrize( + "newer_entry", + [ + pytest.param("not a block", id="not_a_mapping"), + pytest.param({"type": "compaction", "content": None}, id="summary_not_streamed_yet"), + pytest.param({"type": "compaction", "content": ""}, id="empty_summary"), + pytest.param({"type": "text", "content": "Hi."}, id="not_a_compaction"), + ], +) +def test_compaction_output_item_carries_the_newest_block_that_holds_a_summary(newer_entry: object) -> None: + item: Final = compaction_output_item({"compaction_blocks": [_BLOCK, newer_entry]}) + + assert item is not None + assert json.loads(item.encrypted_content) == _BLOCK + + +@pytest.mark.parametrize( + ("response", "holds"), + [ + pytest.param( + {"choices": [{"message": {"role": "assistant", "provider_specific_fields": {"compaction_blocks": [_BLOCK]}}}]}, + True, + id="chat_answer_compacted", + ), + pytest.param({"choices": [{"message": {"role": "assistant", "content": "Hi."}}]}, False, id="chat_answer"), + pytest.param({"object": "response", "output": [_ITEM]}, True, id="responses_output_compacted"), + pytest.param( + {"object": "response", "output": [{"type": "compaction", "id": "cmp_openai", "encrypted_content": "gAAAA"}]}, + False, + id="responses_output_with_an_opaque_compaction", + ), + pytest.param({"choices": "not a list"}, False, id="malformed_choices"), + pytest.param(None, False, id="no_stored_response"), + ], +) +def test_response_holds_compaction_block_reads_chat_and_responses_shaped_rows(response: object, holds: bool) -> None: + assert response_holds_compaction_block(response) is holds + + +@pytest.mark.parametrize( + ("encrypted_content", "holds"), + [ + pytest.param(json.dumps(_BLOCK), True, id="litellm_block"), + pytest.param( + ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id(json.dumps(_BLOCK), "deployment-1"), + True, + id="affinity_wrapped_block", + ), + pytest.param("gAAAAABopaque", False, id="openai_token"), + pytest.param(None, False, id="missing"), + ], +) +def test_only_a_compaction_item_litellm_can_decode_holds_a_summary(encrypted_content: object, holds: bool) -> None: + assert holds_compaction_summary([{"type": "compaction", "id": "cmp_1", "encrypted_content": encrypted_content}]) is ( + holds + ) + + +def test_with_compaction_summary_rewrites_the_summary_and_keeps_the_affinity_wrap() -> None: + wrapped: Final = ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id(json.dumps(_BLOCK), "deployment-1") + + rewritten: Final = with_compaction_summary(wrapped, "[REDACTED]") + model_id, original = ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id(rewritten) + + assert model_id == "deployment-1" + assert json.loads(original) == {**_BLOCK, "content": "[REDACTED]"} + + +def test_with_compaction_summary_leaves_a_token_it_cannot_read_unchanged() -> None: + assert with_compaction_summary("gAAAAABopaque", "[REDACTED]") == "gAAAAABopaque" diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index eead6053208..d646d597d3d 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -6578,7 +6578,9 @@ def _replayed_compacted_answer(model: str, context_management: object) -> list[t role="assistant", content="Day one, Asakusa.", provider_specific_fields={"compaction_blocks": [_COMPACTION_BLOCK]}, - ) + ), + {"role": "user", "content": "Day two?"}, + Message(role="assistant", content="Ueno."), ], litellm_session_id=None, ) @@ -6605,10 +6607,15 @@ def _replayed_compacted_answer(model: str, context_management: object) -> list[t ] +_TURNS_AFTER_THE_COMPACTED_ANSWER: Final = [ + ("user", "Day two?", None), + ("assistant", "Ueno.", None), + ("user", "And Osaka?", None), +] _ANSWER_AFTER_SUMMARY: Final = [ ("user", _COMPACTION_BLOCK["content"], None), ("assistant", "Day one, Asakusa.", None), - ("user", "And Osaka?", None), + *_TURNS_AFTER_THE_COMPACTED_ANSWER, ] @@ -6618,7 +6625,7 @@ _ANSWER_AFTER_SUMMARY: Final = [ ( "anthropic/claude-sonnet-5-5", _COMPACTION_REQUESTED["context_management"], - [("assistant", "Day one, Asakusa.", [_COMPACTION_BLOCK]), ("user", "And Osaka?", None)], + [("assistant", "Day one, Asakusa.", [_COMPACTION_BLOCK]), *_TURNS_AFTER_THE_COMPACTED_ANSWER], ), ("anthropic/claude-sonnet-5-5", None, _ANSWER_AFTER_SUMMARY), ( diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index 3102c878c95..4964ef53c53 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -859,7 +859,7 @@ async def test_message_history_replays_real_key_named_tool_payloads() -> None: def _turn_spend_log( - turn_input: list[dict[str, str]], answer: str, compaction_blocks: list[dict[str, str]] | None = None + turn_input: str | list[dict[str, str]], answer: str, compaction_blocks: list[dict[str, str]] | None = None ) -> dict[str, object]: compaction_fields: Final = ( {"provider_specific_fields": {"compaction_blocks": compaction_blocks}} if compaction_blocks else {} @@ -919,11 +919,17 @@ async def test_previous_response_history_starts_at_the_newest_turn_that_sent_a_c ] +@pytest.mark.parametrize( + "first_turn_input", + [pytest.param([{"role": "user", "content": "Hi."}], id="item_list"), pytest.param("Hi.", id="plain_string")], +) @pytest.mark.asyncio -async def test_previous_response_history_starts_at_the_newest_answer_that_compacted(): +async def test_previous_response_history_starts_at_the_newest_answer_that_compacted( + first_turn_input: str | list[dict[str, str]], +): compaction_blocks: Final = [{"type": "compaction", "content": "Turn one planned Tokyo."}] spend_logs: Final = ( - _turn_spend_log([{"role": "user", "content": "Hi."}], "Hello."), + _turn_spend_log(first_turn_input, "Hello."), _turn_spend_log( [{"role": "developer", "content": "Answer briefly."}, {"role": "user", "content": "Plan a Tokyo trip."}], "Day one, Asakusa.", @@ -1012,6 +1018,41 @@ async def test_previous_response_history_drops_a_tool_output_carrying_a_system_r assert contents[-2:] == ["And Osaka?", "Dotonbori."] + +def _output_message(text: str) -> dict[str, object]: + return {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]} + + +@pytest.mark.asyncio +async def test_previous_response_history_starts_at_a_compaction_item_in_a_stored_responses_output(): + responses_shaped_turn: Final = { + "request_id": "resp_mcp", + "session_id": "session-compacted", + "proxy_server_request": json.dumps({"input": [{"role": "user", "content": "Plan a Tokyo trip."}]}), + "response": { + "id": "resp_mcp", + "object": "response", + "output": [_output_message("Checking the map."), _COMPACTION_ITEM_SENT, _output_message("Day one.")], + }, + } + spend_logs: Final = ( + _turn_spend_log([{"role": "user", "content": "Hi."}], "Hello."), + responses_shaped_turn, + _turn_spend_log([{"role": "user", "content": "And Osaka?"}], "Dotonbori."), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "resp_1") + + replayed: Final = json.dumps([_role_and_content(message) for message in session["messages"]], default=str) + assert "Hello." not in replayed + assert "Plan a Tokyo trip." not in replayed + assert "Checking the map." not in replayed + assert "Turn one planned Tokyo." in replayed + assert [str(_role_and_content(message)[1]) for message in session["messages"]][-2:] == [ + "And Osaka?", + "Dotonbori.", + ] + class _SqliteBackedPrismaDB: def __init__(self, rows: list[tuple[str, str, str, str, str | None, str | None]]) -> None: self._connection = sqlite3.connect(":memory:") diff --git a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py index 1a3430224d1..ed90acbd2fa 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py +++ b/tests/unit/responses/litellm_completion_transformation/test_streaming_iterator_transformation.py @@ -1524,3 +1524,37 @@ async def test_caller_function_named_like_the_search_tool_streams_as_function_ca assert [(item["type"], item["name"]) for item in output] == [("function_call", "tool_search_tool_regex")] assert "response.function_call_arguments.done" in [event["type"] for event in events] + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_a_compaction_block_without_its_summary_streams_no_compaction_item(sync_mode: bool): + deltas: Final = ( + Delta( + role="assistant", + content="", + provider_specific_fields={"compaction_blocks": [{"type": "compaction", "content": None}]}, + ), + Delta(content="Ok."), + ) + chunks: Final = tuple( + ModelResponseStream( + id=CHAT_COMPLETION_ID, + created=1748575031, + model="claude-sonnet-5-5", + object="chat.completion.chunk", + choices=[ + StreamingChoices(index=0, delta=delta, finish_reason="stop" if index == len(deltas) - 1 else None) + ], + ) + for index, delta in enumerate(deltas) + ) + iterator: Final = _build_iterator(chunks) + events: Final = ( + [event.model_dump(exclude_none=True) for event in iterator] + if sync_mode + else [event.model_dump(exclude_none=True) async for event in iterator] + ) + + assert [event["item"]["type"] for event in events if "item" in event] == ["message", "message"] + assert [item["type"] for item in events[-1]["response"]["output"]] == ["message"] diff --git a/tests/unit/responses/test_responses_api_request_body.py b/tests/unit/responses/test_responses_api_request_body.py index e06c88ba24e..50b655d5681 100644 --- a/tests/unit/responses/test_responses_api_request_body.py +++ b/tests/unit/responses/test_responses_api_request_body.py @@ -6,6 +6,7 @@ in expected_responses_api_request/. import copy import json +from collections.abc import Awaitable, Callable from pathlib import Path from importlib import import_module from typing import Final @@ -1052,6 +1053,13 @@ _NEXT_TURN: Final = {"role": "user", "content": "What is my cat's name?"} @pytest.mark.asyncio +@pytest.mark.parametrize( + ("call", "upstream_url"), + [ + pytest.param(litellm.aresponses, "https://api.openai.com/v1/responses", id="create"), + pytest.param(litellm.acompact_responses, "https://api.openai.com/v1/responses/compact", id="compact"), + ], +) @pytest.mark.parametrize( ("compaction_item", "sent_input"), [ @@ -1069,17 +1077,17 @@ _NEXT_TURN: Final = {"role": "user", "content": "What is my cat's name?"} async def test_aresponses_sends_a_claude_compaction_item_to_openai_as_its_summary( monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter, + call: Callable[..., Awaitable[object]], + upstream_url: str, compaction_item: dict[str, str], sent_input: list[dict[str, str]], ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) litellm.in_memory_llm_clients_cache.flush_cache() - upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + upstream: Final = respx_mock.post(upstream_url).mock( return_value=httpx.Response(200, json=_minimal_responses_api_payload("resp_compacted", "gpt-5.6")) ) - await litellm.aresponses( - model="openai/gpt-5.6", api_key="fake-api-key", input=[*_HISTORY, compaction_item, _NEXT_TURN] - ) + await call(model="openai/gpt-5.6", api_key="fake-api-key", input=[*_HISTORY, compaction_item, _NEXT_TURN]) assert json.loads(upstream.calls[0].request.read())["input"] == sent_input From 27d27276c404c2ad9ec15908ee90d3af2d60dc79 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 14:27:32 -0700 Subject: [PATCH 17/21] fix(responses): keep the stored tool call a function_call_output answers when a compaction item replaces history --- .../transformation.py | 91 ++++++++++++++++--- .../test_anthropic_responses_bridge.py | 66 ++++++++++++++ 2 files changed, 146 insertions(+), 11 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index c70d5f69a28..54fc6598dde 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -234,6 +234,7 @@ TOOL_CALLS_CACHE: Final = InMemoryCache() _ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_MessageT = TypeVar("_MessageT") _CACHED_CONTENT_ADAPTER: Final = TypeAdapter(ChatCompletionCachedContent) _DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]]) _TEXT_ADAPTER: Final = TypeAdapter(str) @@ -324,22 +325,88 @@ def _instruction_messages(messages: Sequence[SessionMessage]) -> tuple[SessionMe ) +def _message_field(message: object, key: str) -> object: + if isinstance(message, Mapping): + return _STR_KEY_DICT_ADAPTER.validate_python(message).get(key) + return getattr(message, key, None) + + +def _str_field(message: object, key: str) -> str | None: + value: Final = _message_field(message, key) + return value if isinstance(value, str) else None + + +def _tool_output_call_id(message: object) -> str | None: + return _str_field(message, "tool_call_id") if _message_field(message, "role") == "tool" else None + + +def _assistant_tool_calls(message: object) -> tuple[object, ...]: + if _message_field(message, "role") != "assistant": + return () + return tuple(_OBJECT_LIST_ADAPTER.validate_python(_message_field(message, "tool_calls") or [])) + + +def _answered_tool_call_ids(messages: Sequence[object]) -> frozenset[str]: + called: Final = frozenset( + _str_field(call, "id") for call in chain.from_iterable(map(_assistant_tool_calls, messages)) + ) + return frozenset( + call_id for call_id in map(_tool_output_call_id, messages) if call_id is not None and call_id not in called + ) + + +def _stored_turn_with_calls(message: SessionMessage, call_ids: frozenset[str]) -> SessionMessage | None: + calls: Final = [call for call in _assistant_tool_calls(message) if _str_field(call, "id") in call_ids] + if not calls: + return None + if isinstance(message, Message): + return message.model_copy( + update={"tool_calls": calls, "provider_specific_fields": _provider_fields_without_compaction(message)} + ) + return cast( # cast-ok: the stored assistant turn keeps every field but the calls the new input leaves unanswered + SessionMessage, {**_STR_KEY_DICT_ADAPTER.validate_python(message), "tool_calls": calls} + ) + + +def _with_stored_calls_the_input_answers( + new_messages: Sequence[_MessageT], session_history: Sequence[SessionMessage] +) -> tuple[_MessageT | SessionMessage, ...]: + answered: Final = _answered_tool_call_ids(new_messages) + stored_calls: Final = tuple( + turn for turn in (_stored_turn_with_calls(message, answered) for message in session_history) if turn is not None + ) + if not stored_calls: + return tuple(new_messages) + first_output: Final = next( + index for index, message in enumerate(new_messages) if _tool_output_call_id(message) in answered + ) + return (*new_messages[:first_output], *stored_calls, *new_messages[first_output:]) + + def _stored_compaction_summary(message: SessionMessage) -> str | None: return provider_compaction_summary(message.provider_specific_fields) if isinstance(message, Message) else None +def _provider_fields_without_compaction(message: Message) -> Mapping[str, object] | None: + if not message.provider_specific_fields: + return None + other_fields: Final = { + key: value + for key, value in _STR_KEY_DICT_ADAPTER.validate_python(message.provider_specific_fields).items() + if key != "compaction_blocks" + } + return other_fields or None + + def _with_stored_compaction_as_summary(message: SessionMessage) -> tuple[SessionMessage, ...]: if not isinstance(message, Message): return (message,) summary: Final = provider_compaction_summary(message.provider_specific_fields) if summary is None: return (message,) - other_fields: Final = { - key: value - for key, value in _STR_KEY_DICT_ADAPTER.validate_python(message.provider_specific_fields).items() - if key != "compaction_blocks" - } - answer: Final = message.model_copy(update={"provider_specific_fields": other_fields or None}) + answer: Final = message.model_copy( + update={"provider_specific_fields": _provider_fields_without_compaction(message)} + ) summary_message: Final = ChatCompletionUserMessage(role="user", content=summary) return (summary_message, answer) if answer.content or answer.tool_calls else (summary_message,) @@ -484,9 +551,6 @@ def _block_with_cache_control(block: object, rewrite: _CacheControlRewrite) -> o return _fields_with_cache_control(_STR_KEY_DICT_ADAPTER.validate_python(block), rewrite) -_MessageT = TypeVar("_MessageT") - - def _with_cache_control(message: _MessageT, rewrite: _CacheControlRewrite) -> _MessageT: if not isinstance(message, dict): return message @@ -1162,11 +1226,16 @@ class LiteLLMCompletionResponsesConfig: # Store original _messages before combining for safety check original_new_messages: Final = _messages.copy() if _messages else [] - combined_messages = [*leading_system_messages, *session_messages, *_messages[instructions_end:]] + new_messages: Final = ( + _with_stored_calls_the_input_answers(_messages[instructions_end:], session_history) + if input_replaces_history + else _messages[instructions_end:] + ) + combined_messages = [*leading_system_messages, *session_messages, *new_messages] # Fix: Ensure tool_results have corresponding tool_calls in previous assistant message # Pass tools parameter to help reconstruct tool_calls if not in cache - tools: Final = litellm_completion_request.get("tools") or [] + tools: Final = _OBJECT_LIST_ADAPTER.validate_python(litellm_completion_request.get("tools") or []) combined_messages = LiteLLMCompletionResponsesConfig._ensure_tool_results_have_corresponding_tool_calls( messages=combined_messages, tools=tools ) diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 1c34b92e209..afc0bf644de 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -978,6 +978,72 @@ def test_responses_over_a_custom_provider_send_a_compaction_item_as_summary_text ] +_WEATHER_CALL: Final = { + "id": "toolu_weather", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, +} +_TOOL_TURN_BEFORE_COMPACTION: Final = { + **_TURN_BEFORE_COMPACTION, + "request_id": "chatcmpl-tool-turn", + "response": { + "id": "chatcmpl-tool-turn", + "object": "chat.completion", + "created": 0, + "model": _MODEL, + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": {"role": "assistant", "content": None, "tool_calls": [_WEATHER_CALL]}, + } + ], + }, +} + + +def _content_blocks(message: Mapping[str, object]) -> list[Mapping[str, object]]: + content: Final = message["content"] + return content if isinstance(content, list) else [{"type": "text", "text": content}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("context_management", [None, list(_COMPACT_AT_50K)], ids=["summary_text", "native_block"]) +async def test_previous_response_id_with_a_compaction_item_keeps_the_stored_call_its_tool_output_answers( + context_management: list[Mapping[str, object]] | None, +) -> None: + anthropic: Final = _RecordingRawAnthropicMessages() + compaction_item: Final = { + "type": "compaction", + "id": "cmp_caller", + "encrypted_content": json.dumps(_COMPACTION_BLOCK), + } + tool_output: Final = {"type": "function_call_output", "call_id": "toolu_weather", "output": "18C and sunny"} + + with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient([_TOOL_TURN_BEFORE_COMPACTION])): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-tool-turn", + input=[compaction_item, tool_output], + context_management=context_management, + api_key="sk-ant-fake", + client=AsyncHTTPHandler(transport=httpx.MockTransport(anthropic)), + ) + + messages: Final = json.loads(anthropic.body)["messages"] + call_turn: Final = next( + index + for index, message in enumerate(messages) + if any(block["type"] == "tool_use" for block in _content_blocks(message)) + ) + tool_use: Final = next(block for block in _content_blocks(messages[call_turn]) if block["type"] == "tool_use") + assert (tool_use["id"], tool_use["input"]) == ("toolu_weather", {"city": "Paris"}) + assert _content_blocks(messages[call_turn + 1])[0]["type"] == "tool_result" + assert _content_blocks(messages[call_turn + 1])[0]["tool_use_id"] == "toolu_weather" + assert _COMPACTION_BLOCK["content"] in json.dumps(messages[: call_turn + 1]) + assert _DOCUMENT_BEFORE_COMPACTION not in anthropic.body + + def test_responses_over_anthropic_send_a_compaction_item_as_summary_text_when_the_request_does_not_compact( respx_mock: respx.MockRouter, ) -> None: From e276c80e77a329f49b492c2b19778fdb9edce5cc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 16:33:04 -0700 Subject: [PATCH 18/21] fix(responses): keep a stored tool call on every turn after a compaction item and bound the stored summary A previous_response_id rebuild that starts at a turn holding a compaction item now sends again each stored call a replayed tool output answers, so a retry or a later turn no longer loses the turn-1 call. The start is picked only from rows up to the referenced response. A spend row keeps a compaction summary up to the larger of 64K chars and MAX_STRING_LENGTH_PROMPT_IN_DB instead of with no bound. --- litellm/constants.py | 1 + .../spend_tracking/spend_tracking_utils.py | 82 ++++++-------- .../compaction_items.py | 5 +- .../session_handler.py | 83 ++++++++++---- .../transformation.py | 4 +- .../test_spend_tracking_utils.py | 107 ++++++++++++++++-- .../test_session_handler.py | 68 ++++++++++- 7 files changed, 265 insertions(+), 85 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 39363cc02e4..19d283e7379 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -579,6 +579,7 @@ NON_LLM_CONNECTION_TIMEOUT: Final = get_env_int( ) # timeout for adjacent services (e.g. jwt auth) MAX_EXCEPTION_MESSAGE_LENGTH: Final = get_env_int("MAX_EXCEPTION_MESSAGE_LENGTH", 2000) MAX_STRING_LENGTH_PROMPT_IN_DB: Final = get_env_int("MAX_STRING_LENGTH_PROMPT_IN_DB", 2048) +MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB: Final = 64 * 1024 BEDROCK_MAX_POLICY_SIZE: Final = get_env_int("BEDROCK_MAX_POLICY_SIZE", 75) # One entry per distinct AWS credential-argument set. Per-user cost attribution passes the attributed # identity as aws_session_name, so this bounds how many attributed identities keep a cached STS session. diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 61775d47857..f694145cbb4 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -25,6 +25,7 @@ from litellm.constants import ( LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, + MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB, MCP_SPEND_LOG_MODEL_PREFIX, NON_INFERENCE_CALL_TYPES, REDACTED_BY_LITELM_STRING, @@ -67,7 +68,10 @@ from litellm.proxy._types import ( from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.proxy.utils import PrismaClient, hash_token -from litellm.responses.litellm_completion_transformation.compaction_items import summary_only_compaction_item +from litellm.responses.litellm_completion_transformation.compaction_items import ( + compaction_summary, + summary_only_compaction_content, +) from litellm.types.router import DeploymentTypedDict, LiteLLM_Params from litellm.types.utils import ( BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, @@ -1183,11 +1187,25 @@ def _is_spend_log_content(parent: Mapping[str, object], key: str) -> bool: def _is_compaction_block_summary(parent: Mapping[str, object], key: str, keep_compaction_blocks: bool) -> bool: + return keep_compaction_blocks and key == "content" and parent.get("type") == "compaction" + + +def _truncated_for_db(value: str, max_length: int) -> str: + if len(value) <= max_length: + return value + start_chars: Final = int(max_length * 0.35) + end_chars: Final = int(max_length * 0.65) return ( - keep_compaction_blocks - and key == "content" - and parent.get("type") == "compaction" - and isinstance(parent.get(key), str) + f"{value[:start_chars]}" + f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {len(value) - start_chars - end_chars} chars. " + f"{LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..." + f"{value[len(value) - end_chars :]}" + ) + + +def _compaction_summary_for_db(summary: str) -> str: + return _truncated_for_db( + summary, max(_get_max_string_length_prompt_in_db(), MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB) ) @@ -1207,11 +1225,6 @@ def _sanitize_request_body_for_spend_logs_payload( SensitiveDataMasker classifies as credentials with REDACTED_BY_LITELM_STRING, except inside tool payloads and logprobs. """ - from litellm.constants import ( - LITELLM_TRUNCATED_PAYLOAD_FIELD, - LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, - ) - if visited is None: visited = set() if max_string_length_prompt_in_db is None: @@ -1231,44 +1244,14 @@ def _sanitize_request_body_for_spend_logs_payload( elif isinstance(value, list): return [_sanitize_value(item, mask_credentials) for item in value] elif isinstance(value, str): - if len(value) > max_string_length_prompt_in_db: - # Keep 35% from beginning and 65% from end (end is usually more important) - # This split ensures we keep more context from the end of conversations - start_ratio: Final = 0.35 - end_ratio: Final = 0.65 - - # Calculate character distribution - start_chars: Final = int(max_string_length_prompt_in_db * start_ratio) - end_chars = int(max_string_length_prompt_in_db * end_ratio) - - # Ensure we don't exceed the total limit - total_keep: Final = start_chars + end_chars - if total_keep > max_string_length_prompt_in_db: - end_chars = max_string_length_prompt_in_db - start_chars - - # If the string length is less than what we want to keep, just truncate normally - if len(value) <= max_string_length_prompt_in_db: - return value - - # Calculate how many characters are being skipped - skipped_chars: Final = len(value) - total_keep - - # Build the truncated string: beginning + truncation marker + end - truncated_value: Final = ( - f"{value[:start_chars]}" - f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {skipped_chars} chars. " - f"{LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..." - f"{value[-end_chars:]}" - ) - return truncated_value - return value + return _truncated_for_db(value, max_string_length_prompt_in_db) return value return { k: REDACTED_BY_LITELM_STRING if mask_credentials and _is_request_body_credential(k, v) - else v - if _is_compaction_block_summary(request_body, k, keep_compaction_blocks) + else _compaction_summary_for_db(v) + if isinstance(v, str) and _is_compaction_block_summary(request_body, k, keep_compaction_blocks) else _sanitize_value(v, mask_credentials and not _is_spend_log_content(request_body, k)) for k, v in request_body.items() if k not in _SENSITIVE_REQUEST_BODY_KEYS @@ -1787,11 +1770,14 @@ def _stored_request_without_requester_ip_address(request_body: Mapping[object, o } -def _with_whole_compaction_summary(item: object, stored_item: object) -> object: - summary_only: Final = summary_only_compaction_item(item) - if summary_only is None or not isinstance(stored_item, Mapping): +def _with_bounded_compaction_summary(item: object, stored_item: object) -> object: + summary: Final = compaction_summary(item) + if summary is None or not isinstance(stored_item, Mapping): return stored_item - return {**_OBJECT_MAPPING.validate_python(stored_item), "encrypted_content": summary_only} + return { + **_OBJECT_MAPPING.validate_python(stored_item), + "encrypted_content": summary_only_compaction_content(_compaction_summary_for_db(summary)), + } def _sanitize_proxy_server_request_body(request_body: Mapping[str, object]) -> Mapping[object, object]: @@ -1803,7 +1789,7 @@ def _sanitize_proxy_server_request_body(request_body: Mapping[str, object]) -> M return { **sanitized, "input": [ - _with_whole_compaction_summary(item, stored_item) + _with_bounded_compaction_summary(item, stored_item) for item, stored_item in zip( _OBJECT_LIST.validate_python(items), _OBJECT_LIST.validate_python(stored_items) ) diff --git a/litellm/responses/litellm_completion_transformation/compaction_items.py b/litellm/responses/litellm_completion_transformation/compaction_items.py index 75d957ddbca..ecab21ccce4 100644 --- a/litellm/responses/litellm_completion_transformation/compaction_items.py +++ b/litellm/responses/litellm_completion_transformation/compaction_items.py @@ -142,10 +142,7 @@ def compaction_summary(item: object) -> str | None: return None if block is None else block.summary -def summary_only_compaction_item(item: object) -> str | None: - summary: Final = compaction_summary(item) - if summary is None: - return None +def summary_only_compaction_content(summary: str) -> str: return json.dumps({"type": "compaction", "content": summary}, separators=(",", ":")) diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 1abc00c477a..9fe416f7f06 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -28,6 +28,7 @@ if TYPE_CHECKING: from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload from litellm.responses.litellm_completion_transformation.transformation import ( ChatCompletionSession, + SessionMessage, ) else: ChatCompletionSession = Any @@ -103,6 +104,25 @@ def _replay_start( ) +def _referenced_request_ids(previous_response_id: str, client_response_ids: tuple[str, ...]) -> tuple[str, ...]: + decoded_response_id: Final = ResponsesAPIRequestUtils.decode_responses_api_response_id(previous_response_id) + response_id: Final = decoded_response_id.get("response_id", previous_response_id) + return tuple(dict.fromkeys((response_id, previous_response_id, *client_response_ids))) + + +def _turns_through_referenced_response( + spend_logs: Sequence["SpendLogsPayload"], referenced_request_ids: tuple[str, ...] +) -> int: + return next( + ( + index + 1 + for index in range(len(spend_logs) - 1, -1, -1) + if spend_logs[index].get("request_id") in referenced_request_ids + ), + len(spend_logs), + ) + + def _replayed_request( index: int, proxy_server_request: Mapping[str, object] | None, @@ -152,6 +172,7 @@ class ResponsesSessionHandler: """ from litellm.responses.litellm_completion_transformation.transformation import ( ChatCompletionSession, + with_stored_calls_the_input_answers, ) verbose_proxy_logger.debug("inside get_chat_completion_message_history_for_previous_response_id") @@ -168,30 +189,31 @@ class ResponsesSessionHandler: if len(all_spend_logs) > 0: litellm_session_id = all_spend_logs[0].get("session_id") - chat_completion_message_history: list[ - AllMessageValues - | GenericChatCompletionMessage - | ChatCompletionMessageToolCall - | ChatCompletionResponseMessage - | Message - ] = [] proxy_server_requests: Final = [ await ResponsesSessionHandler.get_proxy_server_request_from_spend_log(spend_log=spend_log) for spend_log in all_spend_logs ] - start: Final = _replay_start(all_spend_logs, proxy_server_requests) + replayable_turns: Final = _turns_through_referenced_response( + all_spend_logs, _referenced_request_ids(previous_response_id, client_response_ids) + ) + start: Final = _replay_start(all_spend_logs[:replayable_turns], proxy_server_requests[:replayable_turns]) earlier_instructions: Final = _instruction_items_before(proxy_server_requests, start) - for index in range(start.index, len(all_spend_logs)): - chat_completion_message_history = ( - ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload( - spend_log=all_spend_logs[index], - proxy_server_request_dict=_replayed_request( - index, proxy_server_requests[index], start, earlier_instructions - ), - chat_completion_message_history=chat_completion_message_history, - item_cache_control=item_cache_control, - ) - ) + replayed_history: Final = ResponsesSessionHandler._history_of_turns( + all_spend_logs[start.index :], + [ + _replayed_request(index, proxy_server_requests[index], start, earlier_instructions) + for index in range(start.index, len(all_spend_logs)) + ], + item_cache_control, + ) + history_before_replay: Final = ResponsesSessionHandler._history_of_turns( + all_spend_logs[: start.index], proxy_server_requests[: start.index], item_cache_control + ) + chat_completion_message_history: Final = ( + list(with_stored_calls_the_input_answers(replayed_history, history_before_replay)) + if history_before_replay + else replayed_history + ) verbose_proxy_logger.debug( "chat_completion_message_history %s", @@ -210,6 +232,24 @@ class ResponsesSessionHandler: ), ) + @staticmethod + def _history_of_turns( + spend_logs: Sequence["SpendLogsPayload"], + proxy_server_requests: Sequence[Mapping[str, object] | None], + item_cache_control: bool, + ) -> list["SessionMessage"]: + return list( + chain.from_iterable( + ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload( + spend_log=spend_log, + proxy_server_request_dict=proxy_server_request, + chat_completion_message_history=[], + item_cache_control=item_cache_control, + ) + for spend_log, proxy_server_request in zip(spend_logs, proxy_server_requests, strict=True) + ) + ) + @staticmethod def extend_chat_completion_message_with_spend_log_payload( spend_log: "SpendLogsPayload", @@ -416,12 +456,11 @@ class ResponsesSessionHandler: verbose_proxy_logger.debug("decoding response id=%s", previous_response_id) - decoded_response_id: Final = ResponsesAPIRequestUtils.decode_responses_api_response_id(previous_response_id) - response_id: Final = decoded_response_id.get("response_id", previous_response_id) + request_ids: Final = list(_referenced_request_ids(previous_response_id, client_response_ids)) + response_id: Final = request_ids[0] if prisma_client is None: return [] - request_ids: Final = list(dict.fromkeys((response_id, previous_response_id, *client_response_ids))) query: Final = """ WITH matching_session AS ( SELECT session_id, api_key, "user", team_id diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 54fc6598dde..d09abac9bcf 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -368,7 +368,7 @@ def _stored_turn_with_calls(message: SessionMessage, call_ids: frozenset[str]) - ) -def _with_stored_calls_the_input_answers( +def with_stored_calls_the_input_answers( new_messages: Sequence[_MessageT], session_history: Sequence[SessionMessage] ) -> tuple[_MessageT | SessionMessage, ...]: answered: Final = _answered_tool_call_ids(new_messages) @@ -1227,7 +1227,7 @@ class LiteLLMCompletionResponsesConfig: original_new_messages: Final = _messages.copy() if _messages else [] new_messages: Final = ( - _with_stored_calls_the_input_answers(_messages[instructions_end:], session_history) + with_stored_calls_the_input_answers(_messages[instructions_end:], session_history) if input_replaces_history else _messages[instructions_end:] ) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 63a164166fc..6f63dcf640e 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -14,8 +14,22 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm import litellm.constants as litellm_constants import litellm.proxy.spend_tracking.spend_tracking_utils as spend_tracking_utils -from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL -from litellm.constants import CALLER_SESSION_ID_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY +from litellm.constants import ( + LITELLM_TRUNCATED_PAYLOAD_FIELD, + LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, + LITTELM_CLI_SERVICE_ACCOUNT_NAME, + LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, + MAX_SPEND_LOG_MODEL_NAME_LENGTH, + REDACTED_BY_LITELM_STRING, + SESSION_ID_GENERATED_METADATA_KEY, + SESSION_ID_OMITTED_METADATA_KEY, + UNKNOWN_MODEL_SPEND_LOG_MODEL, +) +from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, + INTERNAL_CALL_ORIGIN_METADATA_KEY, + MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, +) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -279,7 +293,10 @@ def test_caller_session_keeps_the_legacy_trace_id_with_correlation_logging_off(m def test_session_displaced_by_a_trace_header_wins_only_with_correlation_logging( monkeypatch: pytest.MonkeyPatch, correlation_in_logs: bool, expected_session_id: str ): - header_chain: Final = {"litellm_session_id": "trace-abc", "metadata": {"trace_id": "trace-abc", "session_id": "trace-abc"}} + header_chain: Final = { + "litellm_session_id": "trace-abc", + "metadata": {"trace_id": "trace-abc", "session_id": "trace-abc"}, + } session_id: Final = _chat_session_id( monkeypatch, header_chain, @@ -293,7 +310,11 @@ def test_session_displaced_by_a_trace_header_wins_only_with_correlation_logging( ("litellm_params", "metadata"), [ pytest.param({"litellm_session_id": "convo\x00x"}, {}, id="body-session"), - pytest.param({"litellm_session_id": "trace-abc"}, {CALLER_SESSION_ID_METADATA_KEY: "convo\x00x"}, id="beside-trace-header"), + pytest.param( + {"litellm_session_id": "trace-abc"}, + {CALLER_SESSION_ID_METADATA_KEY: "convo\x00x"}, + id="beside-trace-header", + ), ], ) def test_caller_session_postgres_cannot_store_keeps_the_trace_id( @@ -307,7 +328,10 @@ def test_caller_session_postgres_cannot_store_keeps_the_trace_id( [ pytest.param({"metadata": {}}, id="no-session"), pytest.param( - {"litellm_session_id": "minted-1", "metadata": {"session_id": "minted-1", SESSION_ID_GENERATED_METADATA_KEY: True}}, + { + "litellm_session_id": "minted-1", + "metadata": {"session_id": "minted-1", SESSION_ID_GENERATED_METADATA_KEY: True}, + }, id="minted-by-generate-policy", ), pytest.param({"litellm_session_id": "job-7", "metadata": {"trace_id": "job-7"}}, id="copied-from-trace-id"), @@ -661,9 +685,7 @@ def test_small_request_no_truncation(): # Verify the content was NOT truncated assert sanitized["messages"][0]["content"] == small_content - assert ( - len(sanitized["messages"][0]["content"]) == MAX_STRING_LENGTH_PROMPT_IN_DB - 100 - ) + assert len(sanitized["messages"][0]["content"]) == MAX_STRING_LENGTH_PROMPT_IN_DB - 100 def test_configurable_string_length_env_var(monkeypatch): @@ -1122,6 +1144,75 @@ def test_stored_request_keeps_only_the_summary_of_a_long_claude_compaction_input assert LITELLM_TRUNCATED_PAYLOAD_FIELD in stored["metadata"]["note"]["encrypted_content"] +def _compaction_summary_past_the_summary_cap() -> str: + from litellm.constants import MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB + + filler: Final = "detail " * (MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB // 7) + return f"The cat is named Biscuit. {filler}The dog is named Rex." + + +def _assert_bounded_like_a_db_string(stored: str, summary: str) -> None: + from litellm.constants import MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB + + assert LITELLM_TRUNCATED_PAYLOAD_FIELD in stored + assert len(stored) < MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB + 500 < len(summary) + 500 + assert stored.startswith("The cat is named Biscuit.") + assert stored.endswith("The dog is named Rex.") + + +@pytest.mark.parametrize("operator_cap_above_the_summary", [False, True]) +def test_spend_row_bounds_a_compaction_summary_at_the_larger_of_the_summary_cap_and_the_operator_cap( + monkeypatch: pytest.MonkeyPatch, operator_cap_above_the_summary: bool +): + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + summary: Final = _compaction_summary_past_the_summary_cap() + if operator_cap_above_the_summary: + monkeypatch.setenv("MAX_STRING_LENGTH_PROMPT_IN_DB", str(len(summary) * 2)) + payload: Final = cast( + StandardLoggingPayload, + { + "response": { + "choices": [ + { + "message": { + "role": "assistant", + "content": "OK", + "provider_specific_fields": { + "compaction_blocks": [{"type": "compaction", "content": summary}] + }, + } + } + ] + } + }, + ) + request_body: Final = { + "input": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": summary}), + } + ] + } + + response_message: Final = json.loads(_get_response_for_spend_logs_payload(payload))["choices"][0]["message"] + stored_item: Final = json.loads( + _get_proxy_server_request_for_spend_logs_payload( + metadata={}, litellm_params={"proxy_server_request": {"body": request_body}} + ) + )["input"][0] + + stored_response_block: Final = response_message["provider_specific_fields"]["compaction_blocks"][0] + stored_request_block: Final = json.loads(stored_item["encrypted_content"]) + assert stored_request_block["type"] == stored_response_block["type"] == "compaction" + if operator_cap_above_the_summary: + assert stored_response_block["content"] == stored_request_block["content"] == summary + return + _assert_bounded_like_a_db_string(stored_response_block["content"], summary) + _assert_bounded_like_a_db_string(stored_request_block["content"], summary) + + @patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_response_for_spend_logs_payload_strips_null_bytes(mock_should_store): """Regression for PostgreSQL 22P05: NUL bytes must be stripped from response.""" diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index 4964ef53c53..531ad436603 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -1,6 +1,7 @@ import json import sqlite3 from collections.abc import Sequence +from itertools import chain from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, patch @@ -1018,7 +1019,6 @@ async def test_previous_response_history_drops_a_tool_output_carrying_a_system_r assert contents[-2:] == ["And Osaka?", "Dotonbori."] - def _output_message(text: str) -> dict[str, object]: return {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]} @@ -1053,6 +1053,72 @@ async def test_previous_response_history_starts_at_a_compaction_item_in_a_stored "Dotonbori.", ] + +def _tool_call_ids(message: object) -> tuple[str, ...]: + tool_calls: Final = message.get("tool_calls") if isinstance(message, dict) else getattr(message, "tool_calls", None) + return tuple(call["id"] if isinstance(call, dict) else call.id for call in tool_calls or ()) + + +_WEATHER_CALL_TURN: Final = { + "request_id": "resp_weather", + "session_id": "session-compacted", + "proxy_server_request": json.dumps({"input": [{"role": "user", "content": "Weather in Paris?"}]}), + "response": { + "id": "chatcmpl-weather", + "object": "chat.completion", + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + } + ], + }, +} + + +def _compacted_weather_answer(request_id: str) -> dict[str, object]: + return { + **_turn_spend_log( + [_COMPACTION_ITEM_SENT, {"type": "function_call_output", "call_id": "call_1", "output": "sunny"}], + "It is sunny.", + ), + "request_id": request_id, + } + + +@pytest.mark.asyncio +async def test_previous_response_history_ignores_a_compaction_item_sent_after_the_referenced_response(): + session: Final = await _history_from_stored_turns( + (_WEATHER_CALL_TURN, _compacted_weather_answer("resp_first_send")), "resp_weather" + ) + + assert _role_and_content(session["messages"][0]) == ("user", "Weather in Paris?") + assert tuple(chain.from_iterable(map(_tool_call_ids, session["messages"]))) == ("call_1",) + + +@pytest.mark.asyncio +async def test_a_turn_after_a_compaction_item_keeps_the_stored_call_its_tool_output_answers(): + session: Final = await _history_from_stored_turns( + (_WEATHER_CALL_TURN, _compacted_weather_answer("resp_compacted")), "resp_compacted" + ) + + roles: Final = [_role_and_content(message)[0] for message in session["messages"]] + assert roles == ["user", "assistant", "tool", "assistant"] + assert _tool_call_ids(session["messages"][1]) == ("call_1",) + assert ("user", "Weather in Paris?") not in map(_role_and_content, session["messages"]) + + class _SqliteBackedPrismaDB: def __init__(self, rows: list[tuple[str, str, str, str, str | None, str | None]]) -> None: self._connection = sqlite3.connect(":memory:") From 6621c117ece8ea3a49adc14c8b9960ed0b48e8cf Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 17:32:00 -0700 Subject: [PATCH 19/21] fix(responses): replay a compacted session only through the referenced response A later sibling turn that sent a compaction item put its summary between a stored tool call and its tool output, so Claude answered 400 on a turn that continued the first of two retried turns, and on a retry sent without the item --- .../session_handler.py | 16 +++++++++++++-- .../test_session_handler.py | 20 +++++++++++++++++++ 2 files changed, 34 insertions(+), 2 deletions(-) diff --git a/litellm/responses/litellm_completion_transformation/session_handler.py b/litellm/responses/litellm_completion_transformation/session_handler.py index 9fe416f7f06..924101b39ff 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -104,6 +104,15 @@ def _replay_start( ) +def _holds_compaction_point( + spend_logs: Sequence["SpendLogsPayload"], proxy_server_requests: Sequence[Mapping[str, object] | None] +) -> bool: + return any( + _compaction_point(index, spend_log, proxy_server_request) is not None + for index, (spend_log, proxy_server_request) in enumerate(zip(spend_logs, proxy_server_requests, strict=True)) + ) + + def _referenced_request_ids(previous_response_id: str, client_response_ids: tuple[str, ...]) -> tuple[str, ...]: decoded_response_id: Final = ResponsesAPIRequestUtils.decode_responses_api_response_id(previous_response_id) response_id: Final = decoded_response_id.get("response_id", previous_response_id) @@ -197,12 +206,15 @@ class ResponsesSessionHandler: all_spend_logs, _referenced_request_ids(previous_response_id, client_response_ids) ) start: Final = _replay_start(all_spend_logs[:replayable_turns], proxy_server_requests[:replayable_turns]) + replayed_turns: Final = ( + replayable_turns if _holds_compaction_point(all_spend_logs, proxy_server_requests) else len(all_spend_logs) + ) earlier_instructions: Final = _instruction_items_before(proxy_server_requests, start) replayed_history: Final = ResponsesSessionHandler._history_of_turns( - all_spend_logs[start.index :], + all_spend_logs[start.index : replayed_turns], [ _replayed_request(index, proxy_server_requests[index], start, earlier_instructions) - for index in range(start.index, len(all_spend_logs)) + for index in range(start.index, replayed_turns) ], item_cache_control, ) diff --git a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py index 531ad436603..0eb1aa26a65 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_session_handler.py +++ b/tests/unit/responses/litellm_completion_transformation/test_session_handler.py @@ -1119,6 +1119,26 @@ async def test_a_turn_after_a_compaction_item_keeps_the_stored_call_its_tool_out assert ("user", "Weather in Paris?") not in map(_role_and_content, session["messages"]) +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("spend_logs", "previous_response_id", "replayed_roles"), + [ + ((_WEATHER_CALL_TURN, _compacted_weather_answer("resp_first_send")), "resp_weather", ["user", "assistant"]), + ( + (_WEATHER_CALL_TURN, _compacted_weather_answer("resp_first_send"), _compacted_weather_answer("resp_retry")), + "resp_first_send", + ["user", "assistant", "tool", "assistant"], + ), + ], +) +async def test_a_compacted_session_replays_no_turn_sent_after_the_referenced_response( + spend_logs: tuple[dict[str, object], ...], previous_response_id: str, replayed_roles: list[str] +): + session: Final = await _history_from_stored_turns(spend_logs, previous_response_id) + + assert [_role_and_content(message)[0] for message in session["messages"]] == replayed_roles + + class _SqliteBackedPrismaDB: def __init__(self, rows: list[tuple[str, str, str, str, str | None, str | None]]) -> None: self._connection = sqlite3.connect(":memory:") From 506339246c52786a556a576d56acab7bd3ab701d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 21:55:10 -0700 Subject: [PATCH 20/21] fix(spend-logs): keep main's stored value for a prompt length limit under 2 At a MAX_STRING_LENGTH_PROMPT_IN_DB of 0 or 1 the shared truncation helper stored only the truncation marker, where main stores the marker followed by the whole value --- litellm/proxy/spend_tracking/spend_tracking_utils.py | 2 +- .../proxy/spend_tracking/test_spend_tracking_utils.py | 10 ++++++++++ 2 files changed, 11 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index cbc0d1c6e72..9e6c383766c 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -1205,7 +1205,7 @@ def _truncated_for_db(value: str, max_length: int) -> str: f"{value[:start_chars]}" f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {len(value) - start_chars - end_chars} chars. " f"{LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..." - f"{value[len(value) - end_chars :]}" + f"{value[-end_chars:]}" ) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 2193f1c52fa..edfb1005ec7 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -668,6 +668,16 @@ def test_large_request_no_truncation_threshold(): assert "skipped" in truncated_content +@pytest.mark.parametrize("max_string_length_prompt_in_db", [0, 1]) +def test_a_limit_under_two_still_stores_the_whole_value_after_the_marker(max_string_length_prompt_in_db: int): + sanitized: Final = _sanitize_request_body_for_spend_logs_payload( + {"model": "claude-sonnet-5-5"}, max_string_length_prompt_in_db=max_string_length_prompt_in_db + ) + + assert sanitized["model"].startswith("...") + assert sanitized["model"].endswith("claude-sonnet-5-5") + + def test_small_request_no_truncation(): """ Test that small strings are not truncated by MAX_STRING_LENGTH_PROMPT_IN_DB From 7cfab05960feb666a02e912550c0fba9812932d6 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 22:42:34 -0700 Subject: [PATCH 21/21] fix(spend-logs): keep main's stored tail for a negative prompt length limit At a negative MAX_STRING_LENGTH_PROMPT_IN_DB the shared truncation helper stored one more character of each value than main, since it dropped main's clamp of the end slice to the limit minus the start --- litellm/proxy/spend_tracking/spend_tracking_utils.py | 2 +- .../proxy/spend_tracking/test_spend_tracking_utils.py | 9 ++++++--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 9e6c383766c..536950aff45 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -1205,7 +1205,7 @@ def _truncated_for_db(value: str, max_length: int) -> str: f"{value[:start_chars]}" f"... ({LITELLM_TRUNCATED_PAYLOAD_FIELD} skipped {len(value) - start_chars - end_chars} chars. " f"{LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE}) ..." - f"{value[-end_chars:]}" + f"{value[-min(end_chars, max_length - start_chars) :]}" ) diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index edfb1005ec7..2201150c74c 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -668,14 +668,17 @@ def test_large_request_no_truncation_threshold(): assert "skipped" in truncated_content -@pytest.mark.parametrize("max_string_length_prompt_in_db", [0, 1]) -def test_a_limit_under_two_still_stores_the_whole_value_after_the_marker(max_string_length_prompt_in_db: int): +@pytest.mark.parametrize( + ("max_string_length_prompt_in_db", "stored_tail"), + [(0, "claude-sonnet-5-5"), (1, "claude-sonnet-5-5"), (-1, "laude-sonnet-5-5"), (-2, "aude-sonnet-5-5")], +) +def test_a_limit_under_two_keeps_mains_tail_after_the_marker(max_string_length_prompt_in_db: int, stored_tail: str): sanitized: Final = _sanitize_request_body_for_spend_logs_payload( {"model": "claude-sonnet-5-5"}, max_string_length_prompt_in_db=max_string_length_prompt_in_db ) assert sanitized["model"].startswith("...") - assert sanitized["model"].endswith("claude-sonnet-5-5") + assert sanitized["model"].endswith(f") ...{stored_tail}") def test_small_request_no_truncation():