From a8113b0dd721804c752a6f1952367fb642205224 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 08:50:19 +0000 Subject: [PATCH 1/2] fix(responses)!: round-trip Anthropic compaction blocks as compaction items (#46754) * 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 * fix(responses): replay compaction blocks natively only when the request compacts * 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. * 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. * fix(responses): start previous_response_id replay at a stored compaction and count reservations from it * fix(responses): keep a stored compaction summary whole in spend logs * 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. * 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. * 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. * 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 * test(responses): cover failure-row redaction headers and caller metadata on litellm_metadata routes * 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. * fix(responses): drop the stored history a caller-sent compaction item replaces on previous_response_id * fix(responses): count only message items with a string role as compaction instructions * test(responses): flatten the last two-for comprehension in the compaction wire test * 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. * fix(responses): keep the stored tool call a function_call_output answers when a compaction item replaces history * 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. * 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 * 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 * 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 * fix(responses): keep a compaction item when a guardrail rewrites its summary through structured_messages * test(responses): run the compaction wire tests on the worker database under isolated workers * fix(router): keep the encrypted content affinity pin out of provider requests * fix(proxy): inspect compaction summaries in guardrails, strip a client-sent affinity pin, and price failures from the newest compaction * fix(proxy): count a failed chat or /v1/messages call on its whole history, not from a compaction-typed message * fix(proxy): trim a failed call's input at the newest compaction only when the client called /v1/responses * fix(proxy): trim a failed call's input at the newest compaction when the proxy dispatched it as a Responses call * test(proxy): run the dispatched-call-type failure test through the real request setup * fix(proxy): count a failed Cursor stream's input from the newest compaction item The Cursor route maps to no call type, so the stream generator's failure hook fell back to the route and counted the whole history. Pass the processor's route_type. * fix(proxy): count a broken Cursor stream's prompt from the newest compaction item * fix(responses): keep a tool call a later output answers when trimming before a compaction item * fix(responses): move a compaction-split tool call's turn to right before its output on both compaction routes * perf(responses): look up moved turn items in a set so a long compaction-split turn trims in linear time * test(responses): measure the compaction-split turn move as a CodSpeed benchmark instead of a clock read in unit tests * fix(responses): move only calls still open at the compaction, and bill a disconnected stream from the newest compaction item * fix(responses): restore only the newest stored call when a reused call_id is answered after a compaction item * fix(spend logs): bound a Responses compaction item's summary like a chat block instead of cutting its encrypted_content * test(proxy): price a failed dispatched call through a real ProxyLogging instead of mocked hooks * fix(responses): finish a stream whose Claude compaction failed and group answered call ids by turn once * fix(responses): count input_tokens after a compaction item the way inference sends it --- litellm/constants.py | 1 + litellm/litellm_core_utils/redact_messages.py | 57 +- litellm/llms/anthropic/chat/handler.py | 14 +- .../guardrail_translation/handler.py | 82 +- litellm/proxy/_lazy_openapi_snapshot.json | 3 + litellm/proxy/common_request_processing.py | 6 + litellm/proxy/guardrails/_content_utils.py | 17 +- .../hooks/parallel_request_limiter_v3.py | 7 +- .../proxy/hooks/proxy_track_cost_callback.py | 49 +- litellm/proxy/litellm_pre_call_utils.py | 1 + litellm/proxy/proxy_server.py | 2 + .../proxy/response_api_endpoints/endpoints.py | 21 +- litellm/proxy/spend_tracking/input_tokens.py | 11 +- .../spend_tracking/spend_tracking_utils.py | 126 +- litellm/proxy/utils.py | 31 +- .../compaction_items.py | 329 +++++ .../count_request.py | 21 +- .../handler.py | 3 + .../session_handler.py | 164 ++- .../streaming_iterator.py | 78 +- .../tool_search.py | 18 +- .../transformation.py | 265 ++++- litellm/responses/main.py | 14 +- litellm/responses/streaming_iterator.py | 1 + litellm/router.py | 1 + litellm/types/litellm_params.py | 3 + litellm/types/llms/anthropic.py | 6 + litellm/types/llms/openai.py | 2 + .../test_responses_compaction_benchmarks.py | 35 + .../test_responses_compaction_items_wire.py | 1057 +++++++++++++++++ .../test_redact_messages.py | 22 + ...test_openai_responses_guardrail_handler.py | 180 +++ .../proxy/guardrails/test_content_utils.py | 41 + .../hooks/test_proxy_track_cost_callback.py | 68 ++ tests/unit/proxy/hooks/test_tpm_concurrent.py | 50 +- .../response_api_endpoints/test_endpoints.py | 122 ++ .../proxy/spend_tracking/test_input_tokens.py | 57 + .../test_spend_tracking_utils.py | 233 +++- .../proxy/test_common_request_processing.py | 67 ++ .../unit/proxy/test_litellm_pre_call_utils.py | 2 + ..._utils_model_creation_and_error_logging.py | 135 ++- .../test_anthropic_responses_bridge.py | 639 +++++++++- .../test_compaction_items.py | 87 ++ .../test_count_request.py | 75 ++ .../test_litellm_completion_responses.py | 402 +++++++ .../test_session_handler.py | 367 +++++- .../test_streaming_iterator_transformation.py | 34 + .../test_responses_api_request_body.py | 56 + .../unit/responses/test_streaming_iterator.py | 21 + tests/unit/test_router/test_router.py | 23 + tests/unit/types/test_litellm_params.py | 1 + ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 52 files changed, 4879 insertions(+), 230 deletions(-) create mode 100644 litellm/responses/litellm_completion_transformation/compaction_items.py create mode 100644 tests/benchmarks/test_responses_compaction_benchmarks.py create mode 100644 tests/integration/providers/test_responses_compaction_items_wire.py create mode 100644 tests/unit/responses/litellm_completion_transformation/test_compaction_items.py diff --git a/litellm/constants.py b/litellm/constants.py index 6a503463a8c..c0fe0491775 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -611,6 +611,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/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 0168cdad7c0..7f8a38cc8ec 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -24,6 +24,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, @@ -41,6 +42,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObject = Any +MESSAGE_REDACTION_ENABLE_HEADERS: Final = ("litellm-enable-message-redaction", "x-litellm-enable-message-redaction") + def _opted_out_of_message_logging(custom_logger: CustomLogger) -> bool: return hasattr(custom_logger, "message_logging") and custom_logger.message_logging is not True @@ -195,33 +198,45 @@ 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", + "compaction_delta", + ) + 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 @@ -262,6 +277,11 @@ def _redacted_responses_api_output_item(item: object, redacted_str: str) -> obje if item_mapping.get("type") == "reasoning" and isinstance(summary, list) else {} ), + **( + {"encrypted_content": redacted_str} + if item_mapping.get("type") == "compaction" and "encrypted_content" in item_mapping + else {} + ), **( {"arguments": redacted_str} if item_mapping.get("type") == "function_call" and "arguments" in item_mapping @@ -556,13 +576,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/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index f0143f64be8..13ceea3e31b 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -856,11 +856,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 and isinstance(compaction_delta, str): + 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/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 84ff790ae48..3ea126c5a2a 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -65,6 +65,11 @@ 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, + decode_compaction_item, + 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 ( @@ -110,6 +115,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 @@ -266,7 +281,17 @@ def _item_rewrite_field(item: Mapping[str, object]) -> str | None: return _PATCHABLE_ITEM_FIELDS.get(item_type) +def _rewritten_compaction_item(item: Mapping[str, object], rewritten: object) -> Mapping[str, object] | None: + encrypted_content: Final = item.get("encrypted_content") + summary: Final = rewritten.get("content") if isinstance(rewritten, Mapping) else None + if not isinstance(encrypted_content, str) or not isinstance(summary, str) or not summary: + return None + return {**item, "encrypted_content": with_compaction_summary(encrypted_content, summary)} + + def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapping[str, object] | None: + if compaction_summary(item) is not None: + return _rewritten_compaction_item(item, rewritten) field: Final = _item_rewrite_field(item) if field is None or not isinstance(rewritten, Mapping): return None @@ -804,6 +829,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 @@ -849,6 +879,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 @@ -1116,13 +1152,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( @@ -1214,6 +1254,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], @@ -1256,11 +1304,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) @@ -1481,13 +1531,13 @@ 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 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( - 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 @@ -1610,6 +1660,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): @@ -1677,6 +1733,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/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 545b317055a..4572e59944b 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -28240,6 +28240,9 @@ }, { "$ref": "#/components/schemas/OutputToolSearchCall" + }, + { + "$ref": "#/components/schemas/ResponseCompactionItem" } ] }, diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6fda5780f63..23faa964d5f 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -1985,6 +1985,7 @@ class ProxyBaseLLMRequestProcessing: self.data = data self.served_model = served_model self._tags_before_guardrails: frozenset[str] | None = None + self.route_type: str | None = None @property def litellm_call_id(self) -> str | None: @@ -2434,6 +2435,7 @@ class ProxyBaseLLMRequestProcessing: ## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call ## IMPORTANT Note: - initialize this before running pre-call checks. Ensures we log rejected requests to langfuse. + self.route_type = route_type logging_obj, self.data = litellm.utils.function_setup( original_function=route_type, rules_obj=litellm.utils.Rules(), @@ -2800,6 +2802,7 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict=user_api_key_dict, original_exception=exc, request_data=self.data, + dispatched_call_type=self.route_type, ) return await open_sse_before_first_byte( @@ -2864,6 +2867,7 @@ class ProxyBaseLLMRequestProcessing: "skip_pre_call_logic=True requires litellm_logging_obj to be set in data. " "Ensure common_processing_pre_call_logic was called before using this parameter." ) + self.route_type = route_type else: self.data, logging_obj = await self._pre_call_with_fallbacks( request=request, @@ -3473,6 +3477,7 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict=user_api_key_dict, original_exception=exc, request_data=request_data, + dispatched_call_type=self.route_type, ) except HTTPException: return @@ -4142,6 +4147,7 @@ class ProxyBaseLLMRequestProcessing: user_api_key_dict=user_api_key_dict, original_exception=e, request_data=self.data, + dispatched_call_type=self.route_type, ) # Use transformed exception if callback returned one, otherwise use original if transformed_exception is not None: diff --git a/litellm/proxy/guardrails/_content_utils.py b/litellm/proxy/guardrails/_content_utils.py index 64c935ce0d0..17bea031794 100644 --- a/litellm/proxy/guardrails/_content_utils.py +++ b/litellm/proxy/guardrails/_content_utils.py @@ -11,6 +11,12 @@ every text fragment. from collections.abc import Callable, Iterator, Mapping, Sequence from typing import Any, Final +from litellm.responses.litellm_completion_transformation.compaction_items import ( + compaction_summary, + input_items, + with_visited_compaction_summary, +) + # Call types whose body carries free-form chat / prompt text that # text-content guardrails (banned keywords, content moderation, secret # detection, …) should inspect. The proxy ingress passes ``route_type`` @@ -95,12 +101,12 @@ def _coerce_input_to_messages(input_value: object) -> list[dict[str, object]]: """Coerce a Responses-API ``data["input"]`` value into chat-style messages.""" if isinstance(input_value, str): return [{"role": "user", "content": input_value}] - if not isinstance(input_value, list): - return [] messages: Final[list[dict[str, object]]] = [] - for item in input_value: + for item in input_items(input_value): if isinstance(item, str): messages.append({"role": "user", "content": item}) + elif (summary := compaction_summary(item)) is not None: + messages.append({"role": "user", "content": summary}) elif isinstance(item, dict): if _part_text(item) is not None: messages.append({"role": item.get("role") or "user", "content": [item]}) @@ -189,11 +195,14 @@ def walk_user_text(data: dict[str, object], visit: Callable[[str], str]) -> int: data["input"] = visit(input_value) return visited if isinstance(input_value, list): - for idx, item in enumerate(input_value): + for idx, item in enumerate(input_items(input_value)): if isinstance(item, str): if item: visited += 1 input_value[idx] = visit(item) + elif (compaction := with_visited_compaction_summary(item, visit)) is not None: + visited += 1 + input_value[idx] = compaction elif isinstance(item, dict): if _part_text(item) is not None: visited += 1 diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index a8ba8d65785..a2bb691b401 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -73,6 +73,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 billable_input_since_newest_compaction from litellm.router_utils.add_retry_fallback_headers import ( ensure_response_additional_headers, response_has_hidden_params, @@ -1247,7 +1248,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 = ( + billable_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)) @@ -3835,6 +3839,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/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 994310c6619..6801c8521e5 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -22,6 +22,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.llms.azure.ptu_capacity import is_str_keyed_mapping from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( @@ -127,6 +128,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: @@ -183,6 +187,26 @@ def _failure_snapshot_with_guardrails_recorded_after_it( return {**failure_snapshot, "guardrail_information": recorded_guardrail_information} +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") + 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: from litellm.proxy.proxy_server import proxy_logging_obj @@ -417,13 +441,17 @@ class _ProxyDBLogger(CustomLogger): existing_metadata.update({**_routed_metadata_fallback(request_fields), **_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 + litellm_metadata: Final[ObjectMapping] = ( + _OBJECT_MAPPING.validate_python(litellm_metadata_bucket) + if isinstance(litellm_metadata_bucket, dict) + else {} + ) + existing_metadata.update( + _guardrail_info_lifted_from_litellm_metadata( + litellm_metadata, + guardrail_info_already_set="standard_logging_guardrail_information" in existing_metadata, + ) + ) if "litellm_params" not in request_data: request_data["litellm_params"] = {} @@ -438,7 +466,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/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d34787cd5e9..8a45f658f36 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -353,6 +353,7 @@ LITELLM_TRACE_CONTROL_METADATA_FIELDS: Final = frozenset( _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = ( "weights", "_router_weights", + "_encrypted_content_affinity_pinned", "fallback_depth", "_target_order", "attempted_targets", diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b48b5ab1c1e..f632317d9a8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10384,6 +10384,7 @@ async def async_data_generator( request: Request | None = None, *, responses_stream_errors: bool = False, + dispatched_call_type: str | None = None, ): verbose_proxy_logger.debug("inside generator") stream_completed = False @@ -10599,6 +10600,7 @@ async def async_data_generator( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=request_data, + dispatched_call_type=dispatched_call_type, ) client_error: Final[Exception] = e if replacement is None else replacement verbose_proxy_logger.debug( diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index c918ea3f1d1..312c5124ca5 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -37,6 +37,7 @@ from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADAT from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, first_validation_error from litellm.llms.openai.responses.guardrail_translation.handler import ( @@ -79,7 +80,8 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attributio request_tags_from_request_data, ) from litellm.proxy.route_llm_request import raise_if_required_body_param_missing -from litellm.proxy.utils import handle_exception_on_proxy +from litellm.proxy.utils import count_request_input_tokens, handle_exception_on_proxy +from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction from litellm.responses.litellm_completion_transformation.count_request import MODEL_FILE_ID_MAPPING_FIELD from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard, spawn_logging_task from litellm.router import Router @@ -94,7 +96,6 @@ from litellm.types.token_count import TokenCountSourceFields from litellm.types.utils import ADDRESSED_RESPONSE_ID_FIELD, TokenCountResponse if TYPE_CHECKING: - from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy.utils import ProxyLogging router: Final = APIRouter() @@ -532,6 +533,14 @@ async def cursor_model_list( return await model_list(user_api_key_dict=user_api_key_dict) +def _dispatched_prompt_token_counter(logging_obj: LiteLLMLoggingObj) -> Callable[[], int]: + return lambda: count_request_input_tokens( + model=str(logging_obj.model_call_details.get("model") or ""), + request_input=messages_since_newest_compaction(logging_obj.model_call_details.get("messages")), + system_input=None, + ) + + @router.post( "/cursor/chat/completions", dependencies=[ @@ -656,6 +665,11 @@ async def cursor_chat_completions( model=request_data.get("model", ""), custom_llm_provider=None, logging_obj=logging_obj, + count_prompt_tokens=( + _dispatched_prompt_token_counter(logging_obj) + if isinstance(logging_obj, LiteLLMLoggingObj) + else None + ), ) # Use async_data_generator to format as SSE return async_data_generator( @@ -663,6 +677,7 @@ async def cursor_chat_completions( user_api_key_dict=user_api_key_dict, request_data=request_data, request=request, + dispatched_call_type=processor.route_type, ) # Otherwise, use the default generator return async_data_generator( @@ -670,6 +685,7 @@ async def cursor_chat_completions( user_api_key_dict=user_api_key_dict, request_data=request_data, request=request, + dispatched_call_type=processor.route_type, ) try: @@ -1585,6 +1601,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/proxy/spend_tracking/input_tokens.py b/litellm/proxy/spend_tracking/input_tokens.py index 6c7083fb6db..927152296e9 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 ( + billable_input_since_newest_compaction, + holds_compaction_summary, +) 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,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=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/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 62aee788ed0..94b079163e9 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -27,6 +27,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, @@ -70,6 +71,11 @@ 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 ( + compaction_summary, + summary_only_compaction_content, + with_compaction_summary, +) from litellm.types.passthrough_endpoints.pass_through_endpoints import ( SETTLES_BACKGROUND_RESPONSE_KWARG, BackgroundResponseSettlement, @@ -1227,11 +1233,42 @@ def _is_spend_log_content(parent: Mapping[str, object], key: str) -> bool: ) +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 ( + 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[-min(end_chars, max_length - start_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) + ) + + +def _bounded_compaction_summary( + parent: Mapping[str, object], key: str, value: object, keep_compaction_blocks: bool +) -> str | None: + if not keep_compaction_blocks or not isinstance(value, str) or parent.get("type") != "compaction": + return None + if key == "content": + return _compaction_summary_for_db(value) + summary: Final = compaction_summary(parent) if key == "encrypted_content" else None + return None if summary is None else with_compaction_summary(value, _compaction_summary_for_db(summary)) + + 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. @@ -1242,11 +1279,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: @@ -1261,51 +1293,23 @@ 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] 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 _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 - } + def _stored_value(key: str, value: object) -> object: + if mask_credentials and _is_request_body_credential(key, value): + return REDACTED_BY_LITELM_STRING + bounded_summary: Final = _bounded_compaction_summary(request_body, key, value, keep_compaction_blocks) + if bounded_summary is not None: + return bounded_summary + return _sanitize_value(value, mask_credentials and not _is_spend_log_content(request_body, key)) + + return {k: _stored_value(k, v) for k, v in request_body.items() if k not in _SENSITIVE_REQUEST_BODY_KEYS} # Quoted-key form: ``"input"`` / ``'messages'`` / ``"prompt"`` followed by @@ -1820,6 +1824,33 @@ def _stored_request_without_requester_ip_address(request_body: Mapping[object, o } +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_compaction_content(_compaction_summary_for_db(summary)), + } + + +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_bounded_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], @@ -1864,8 +1895,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 @@ -1939,7 +1969,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/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2a328971817..52b3866d04d 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -89,6 +89,7 @@ from litellm.proxy.common_utils.openai_error_payload import ( with_litellm_call_id, ) from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error +from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction from litellm.types.guardrails import GuardrailEventHooks from litellm.types.proxy.model_listing import ModelInfoResponse from litellm.types.utils import MCP_GUARDRAIL_CALL_TYPES, CallTypes, CallTypesLiteral, ModelInfo, Usage @@ -937,7 +938,7 @@ def _system_prompt_text(system_input: object) -> str: return "".join(_prompt_block_text(block) for block in system_input) -def _count_request_input_tokens(model: str, request_input: object, system_input: object) -> int: +def count_request_input_tokens(model: str, request_input: object, system_input: object) -> int: system_text: Final = _system_prompt_text(system_input) system_tokens: Final = litellm.token_counter(model=model, text=system_text) if system_text else 0 if isinstance(request_input, str): @@ -958,7 +959,7 @@ def _estimate_dispatched_failure_usage(model: str, request_input: object, system the same tokenizer fallback interrupted streams use, so the spend log's failure row records what was sent instead of zero.""" try: - input_tokens: Final = _count_request_input_tokens( + input_tokens: Final = count_request_input_tokens( model=model, request_input=request_input, system_input=system_input ) except Exception: @@ -997,6 +998,8 @@ _INPUT_ESTIMABLE_CALL_TYPES: Final = frozenset( ) +_RESPONSES_CALL_TYPES: Final = frozenset((CallTypes.responses.value, CallTypes.aresponses.value)) + _FAILED_ATTEMPTS_USAGE: Final = TypeAdapter(tuple[tuple[Usage, float], ...]) @@ -1022,6 +1025,7 @@ def _current_failure_usage_to_lift( provider_billed: bool = False, deployment_id: str | None = None, base_model: str | None = None, + dispatched_call_type: str | None = None, ) -> tuple[object, object] | None: """A stream that broke mid-flight still billed the provider for the chunks already delivered; the streaming handler stashes that recovered usage and @@ -1045,7 +1049,7 @@ def _current_failure_usage_to_lift( # The proxy-only error handler stamps this flag on a client disconnect too, even though the # request had already reached the provider, so a billed failure keeps its estimate. return None - estimated_usage: Final = _estimated_prompt_usage(model_call_details, request_body) + estimated_usage: Final = _estimated_prompt_usage(model_call_details, request_body, dispatched_call_type) if estimated_usage is None: return None if not provider_billed: @@ -1056,13 +1060,14 @@ def _current_failure_usage_to_lift( def _estimated_prompt_usage( - model_call_details: Mapping[str, object], request_body: Mapping[str, object] + model_call_details: Mapping[str, object], request_body: Mapping[str, object], dispatched_call_type: str | None ) -> Usage | None: """The input-only Usage estimate for the prompt this request sent, or None when the call type has no estimable input. The raw request body backfills the system prompt when the SDK bridges an endpoint (e.g. /v1/messages on a chat-completions provider) without filling optional_params.""" if str(model_call_details.get("call_type")) not in _INPUT_ESTIMABLE_CALL_TYPES: return None + dispatched_input: Final = model_call_details.get("messages") optional_params: Final = model_call_details.get("optional_params") dispatched_system: Final = ( (optional_params.get("system") or optional_params.get("instructions")) @@ -1079,7 +1084,11 @@ def _estimated_prompt_usage( ) return _estimate_dispatched_failure_usage( model=str(model_call_details.get("model") or ""), - request_input=model_call_details.get("messages"), + request_input=( + messages_since_newest_compaction(dispatched_input) + if dispatched_call_type in _RESPONSES_CALL_TYPES + else dispatched_input + ), system_input=system_input, ) @@ -1444,7 +1453,9 @@ def _provider_billed_the_failure(exception: Exception) -> bool: return _timed_out_reading_from_provider(exception) or _client_disconnected(exception) -def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed: bool = False) -> Mapping[str, object]: +def _failure_fields_to_lift( + request_data: Mapping[str, object], provider_billed: bool = False, dispatched_call_type: str | None = None +) -> Mapping[str, object]: """Failure-path callbacks run after ``litellm_logging_obj`` is popped from request_data (it is not serialisable), so the caller merges these fields onto request_data first: the request start and first-handoff instants for @@ -1468,9 +1479,10 @@ def _failure_fields_to_lift(request_data: Mapping[str, object], provider_billed: provider_billed=_billed, deployment_id=_deployment_id, base_model=_base_model, + dispatched_call_type=dispatched_call_type, ) _prompt_estimate: Final = ( - _estimated_prompt_usage(_model_call_details, request_data) + _estimated_prompt_usage(_model_call_details, request_data, dispatched_call_type) if _billed and _first_handoff is not None and not _model_call_details.get("failed_attempts_usage") else None ) @@ -3634,6 +3646,7 @@ class ProxyLogging: error_type: ProxyErrorTypes | None = None, route: str | None = None, traceback_str: str | None = None, + dispatched_call_type: str | None = None, ) -> HTTPException | None: """ Allows users to raise custom exceptions/log when a call fails, without having to deal with parsing Request body. @@ -3704,7 +3717,9 @@ class ProxyLogging: request_data.update( await offload_token_count(_failure_fields_to_lift)( - request_data, provider_billed=_provider_billed_the_failure(original_exception) + request_data, + provider_billed=_provider_billed_the_failure(original_exception), + dispatched_call_type=dispatched_call_type or _call_type_for_route(user_api_key_dict.request_route), ) ) 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..8bd82c49090 --- /dev/null +++ b/litellm/responses/litellm_completion_transformation/compaction_items.py @@ -0,0 +1,329 @@ +import json +import uuid +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from itertools import chain, groupby +from types import MappingProxyType +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 pydantic import BaseModel, ConfigDict, 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")) + + +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) +class CompactionBlock: + summary: str + wire_block: Mapping[str, object] + + +def _compaction_block(value: object) -> CompactionBlock | None: + 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: + 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 _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 + + +class _AssistantTurnFields(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + content: object = None + tool_calls: object = None + provider_specific_fields: object = None + + +def carries_only_compaction(message: object) -> bool: + fields: Final = _AssistantTurnFields.model_validate(message) + return ( + not fields.content + and not fields.tool_calls + and provider_compaction_summary(fields.provider_specific_fields) is not None + ) + + +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) + 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(fields.get("choices")) + except ValidationError: + 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( + id=f"cmp_{uuid.uuid4()}", + encrypted_content=json.dumps(dict(block.wire_block), separators=(",", ":")), + type="compaction", + ) + + +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 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 + + +def decode_compaction_item(encrypted_content: object) -> CompactionBlock | None: + if not isinstance(encrypted_content, str): + return None + try: + 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 summary_only_compaction_content(summary: str) -> str: + 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) + if block is None: + return encrypted_content + rewritten: Final = json.dumps({**block.wire_block, "content": summary}, separators=(",", ":")) + return encrypted_content.removesuffix(original) + rewritten + + +def with_visited_compaction_summary(item: object, visit: Callable[[str], str]) -> Mapping[str, object] | None: + fields: Final = _json_object(item) + if fields is None: + return None + summary: Final = compaction_summary(fields) + encrypted_content: Final = fields.get("encrypted_content") + if summary is None or not isinstance(encrypted_content, str): + return None + return {**fields, "encrypted_content": with_compaction_summary(encrypted_content, visit(summary))} + + +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) + 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, ...]: + try: + return tuple(_OBJECT_LIST.validate_python(input)) + except ValidationError: + return () + + +def instruction_items(input: object) -> tuple[object, ...]: + return tuple(item for item in input_items(input) if _is_instruction_item(item)) + + +def _call_id(item: object) -> str | None: + fields: Final = _json_object(item) + call_id: Final = None if fields is None else fields.get("call_id") + return call_id if isinstance(call_id, str) else None + + +def _is_turn_item(item: object) -> bool: + fields: Final = _json_object(item) + item_type: Final = None if fields is None else fields.get("type", "message") + if fields is None or not isinstance(item_type, str): + return False + if item_type == "message": + return fields.get("role") == "assistant" + return item_type == "reasoning" or (_call_id(fields) is not None and not item_type.endswith("_output")) + + +def _turn_indices_answered_later( + turn: tuple[tuple[int, object], ...], calls_answered_later: frozenset[int] +) -> tuple[int, ...]: + if not any(index in calls_answered_later for index, _ in turn): + return () + return tuple(index for index, item in turn if _call_id(item) is None or index in calls_answered_later) + + +def _is_turn_entry(entry: tuple[int, object]) -> bool: + return _is_turn_item(entry[1]) + + +def _calls_answered_only_later(earlier: Sequence[object], answered: frozenset[str]) -> frozenset[int]: + last_index_by_call_id: Final = MappingProxyType( + {call_id: index for index, item in enumerate(earlier) if (call_id := _call_id(item)) is not None} + ) + return frozenset(index for call_id, index in last_index_by_call_id.items() if call_id in answered) + + +def _earlier_turn_indices_answered_later(earlier: Sequence[object], answered: frozenset[str]) -> tuple[int, ...]: + calls_answered_later: Final = _calls_answered_only_later(earlier, answered) + turns: Final = (tuple(turn) for is_turn, turn in groupby(enumerate(earlier), key=_is_turn_entry) if is_turn) + return tuple(chain.from_iterable(_turn_indices_answered_later(turn, calls_answered_later) for turn in turns)) + + +def _call_ids(items: Sequence[object]) -> frozenset[str]: + return frozenset(call_id for item in items if (call_id := _call_id(item)) is not None) + + +def _keeps_every_item(item: object) -> bool: + return True + + +def _with_answered_turns_after_newest_compaction( + input: Sequence[_Item], keeps_earlier_item: Callable[[object], bool] +) -> tuple[_Item, ...]: + newest: Final = next( + (index for index in range(len(input) - 1, 0, -1) if compaction_summary(input[index]) is not None), 0 + ) + later: Final = input[newest:] + moved: Final = _earlier_turn_indices_answered_later(input[:newest], _call_ids(later)) + moved_indices: Final = frozenset(moved) + moved_call_ids: Final = _call_ids([input[index] for index in moved]) + first_answer: Final = next( + (index for index, item in enumerate(later) if _call_id(item) in moved_call_ids), len(later) + ) + return ( + *(item for index, item in enumerate(input[:newest]) if index not in moved_indices and keeps_earlier_item(item)), + *later[:first_answer], + *(input[index] for index in moved), + *later[first_answer:], + ) + + +def _items_from_newest_compaction(input: Sequence[_Item]) -> tuple[_Item, ...]: + return _with_answered_turns_after_newest_compaction(input, _is_instruction_item) + + +def with_answered_turns_after_newest_compaction(input: str | ResponseInputParam) -> str | ResponseInputParam: + return ( + list(_with_answered_turns_after_newest_compaction(input, _keeps_every_item)) + if isinstance(input, list) + else input + ) + + +def without_items_before_compaction(input: str | ResponseInputParam) -> str | ResponseInputParam: + return list(_items_from_newest_compaction(input)) if isinstance(input, list) else input + + +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 [_summary_or_item(item) for item in _items_from_newest_compaction(items)] + + +def _as_summary_message(item: _Item) -> _Item | EasyInputMessageParam: + 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)] + + +def messages_since_newest_compaction(input: object) -> object: + try: + items: Final = _OBJECT_LIST.validate_python(input) + except ValidationError: + return input + return [_as_summary_message(item) for item in _items_from_newest_compaction(items)] diff --git a/litellm/responses/litellm_completion_transformation/count_request.py b/litellm/responses/litellm_completion_transformation/count_request.py index 56a0b9024c5..7679b4ede71 100644 --- a/litellm/responses/litellm_completion_transformation/count_request.py +++ b/litellm/responses/litellm_completion_transformation/count_request.py @@ -1,4 +1,4 @@ -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass from typing import Final, cast # noqa: TID251 # input items are forwarded as sent, as inference does @@ -14,7 +14,11 @@ from litellm.litellm_core_utils.tool_fields import DropPolicy from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body from litellm.responses.additional_tools import hoist_additional_tools from litellm.responses.file_search.emulated_handler import replace_file_search_tools -from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig +from litellm.responses.litellm_completion_transformation.compaction_items import holds_compaction_summary +from litellm.responses.litellm_completion_transformation.transformation import ( + ChatCompletionSession, + LiteLLMCompletionResponsesConfig, +) from litellm.responses.main import previous_response_client_ids from litellm.types.llms.openai import ResponseInputParam from litellm.utils import ProviderConfigManager @@ -132,12 +136,17 @@ def _with_emulated_file_search(responses_request: Mapping[str, object]) -> Mappi return {**responses_request, "tools": replace_file_search_tools(_ITEMS.validate_python(tools))[0]} +def _with_compaction_as_summary_text(responses_request: Mapping[str, object]) -> Mapping[str, object]: + return {name: value for name, value in responses_request.items() if name != "context_management"} + + async def _with_previous_turns( chat_body: Mapping[str, object], model: str, custom_llm_provider: str, responses_request: Mapping[str, object], policy: DropPolicy, + session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None, ) -> Mapping[str, object]: previous_response_id: Final = responses_request.get("previous_response_id") if not isinstance(previous_response_id, str) or not previous_response_id: @@ -147,8 +156,10 @@ async def _with_previous_turns( previous_response_id=previous_response_id, litellm_completion_request={**chat_body, "model": model, "custom_llm_provider": custom_llm_provider}, instructions=instructions if isinstance(instructions, str) else None, + session_lookup=session_lookup, request_kwargs={"drop_params": policy.drop_params}, previous_response_client_ids=previous_response_client_ids(previous_response_id, responses_request), + input_replaces_history=holds_compaction_summary(responses_request.get("input")), ) @@ -171,6 +182,7 @@ async def chat_count_request( api_base: str | None, responses_request: Mapping[str, object], policy: DropPolicy, + session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None = None, ) -> ChatCountRequest | ChatCountBridgeError | None: """The chat request Responses inference sends a provider with no native Responses API, or None when inference sends this provider the Responses body itself.""" @@ -184,7 +196,9 @@ async def chat_count_request( try: chat_body: Final = await _with_previous_turns( responses_batch_body_to_chat_body( - _with_emulated_file_search(_with_hoisted_additional_tools(responses_request)), + _with_compaction_as_summary_text( + _with_emulated_file_search(_with_hoisted_additional_tools(responses_request)) + ), custom_llm_provider, policy.drop_params, ), @@ -192,6 +206,7 @@ async def chat_count_request( custom_llm_provider, responses_request, policy, + session_lookup, ) tools: Final = ( () if "tools" in policy.additional_drop_params else tuple(map(_unmarked, _mappings(chat_body.get("tools")))) diff --git a/litellm/responses/litellm_completion_transformation/handler.py b/litellm/responses/litellm_completion_transformation/handler.py index 5489d2223ac..c9d800421f7 100644 --- a/litellm/responses/litellm_completion_transformation/handler.py +++ b/litellm/responses/litellm_completion_transformation/handler.py @@ -9,11 +9,13 @@ 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, ) 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 ( @@ -120,6 +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(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 35160830fdc..d50e9a61edc 100644 --- a/litellm/responses/litellm_completion_transformation/session_handler.py +++ b/litellm/responses/litellm_completion_transformation/session_handler.py @@ -1,12 +1,20 @@ 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 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, + input_items, + instruction_items, + response_holds_compaction_block, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( AllMessageValues, @@ -20,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 @@ -66,6 +75,95 @@ def _stored_responses_api_output(stored_response: object) -> ResponseInputParam return None +@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( + ( + 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 + ), + _ReplayStart(index=0, at_output=False), + ) + + +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) + 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, + start: _ReplayStart, + earlier_instructions: tuple[object, ...], +) -> Mapping[str, object] | None: + if proxy_server_request is None or index != start.index: + return proxy_server_request + 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: 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 @@ -91,6 +189,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") @@ -107,26 +206,34 @@ 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 ] - for spend_log, proxy_server_request_dict in zip(all_spend_logs, proxy_server_requests): - 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, - chat_completion_message_history=chat_completion_message_history, - item_cache_control=item_cache_control, - ) - ) + 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]) + 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 : replayed_turns], + [ + _replayed_request(index, proxy_server_requests[index], start, earlier_instructions) + for index in range(start.index, replayed_turns) + ], + 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", @@ -138,6 +245,24 @@ class ResponsesSessionHandler: instructions=_latest_replayable_instructions(proxy_server_requests), ) + @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", @@ -182,6 +307,7 @@ class ResponsesSessionHandler: responses_api_request={}, replay_reasoning=True, item_cache_control=item_cache_control, + only_since_compaction=True, ) ) @@ -196,6 +322,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 != {}: @@ -342,12 +469,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/streaming_iterator.py b/litellm/responses/litellm_completion_transformation/streaming_iterator.py index 4ef68aa072a..c25dd4a405e 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -8,6 +8,7 @@ import litellm from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params, served_from_cache from litellm.litellm_core_utils.litellm_logging import BilledRoundsSpend 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, @@ -83,6 +84,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) @@ -146,6 +152,9 @@ 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._sent_compaction_item_done: bool = False self._cached_response_id: str | None = None self._buffered_chunk: ModelResponseStream | None = None self._upstream_exhausted: bool = False @@ -186,7 +195,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, @@ -194,7 +203,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: @@ -248,6 +260,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)) @@ -256,6 +270,45 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): index: Final = self._streamed_output_index(item) return (index is None, index or 0) + 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 + streamed_item: Final = self._streamed_compaction_item() + if streamed_item is None: + return + self._compaction_output_index = self._allocate_output_index() + 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: + return + streamed_item: Final = self._streamed_compaction_item() + if streamed_item is None: + return + self._sent_compaction_item_done = True + self._sequence_number += 1 + 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, + ) + ) + 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: @@ -592,6 +645,12 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): } ) + def _ordered_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: + ordered: Final = self._reasoning_part_first(chunk) + if not _streams_compaction(ordered): + self._queue_compaction_item_added() + return ordered + def _with_encoded_response_id(self, response: ResponsesAPIResponse) -> ResponsesAPIResponse: return ResponsesAPIRequestUtils.update_responses_api_response_id_with_model_id( responses_api_response=response, @@ -1004,6 +1063,10 @@ 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_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. if isinstance(self.litellm_model_response, ModelResponse): self._queue_final_tool_call_done_events(self.litellm_model_response) @@ -1091,6 +1154,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): self._reasoning_item_id or self._cached_reasoning_item_id or mint_reasoning_item_id() ) reasoning_item_id: Final = self._cached_reasoning_item_id + self._queue_compaction_item_done() self._sequence_number += 1 text_done_event: Final = self.create_reasoning_summary_text_done_event( reasoning_item_id=reasoning_item_id, @@ -1139,7 +1203,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): raise StopAsyncIteration chunk = await self.litellm_custom_stream_wrapper.__anext__() if chunk is not None: - chunk = self._reasoning_part_first(cast(ModelResponseStream, chunk)) + chunk = self._ordered_chunk(cast(ModelResponseStream, chunk)) for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -1195,11 +1259,11 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): try: buffered_chunk = self._take_buffered_chunk() if buffered_chunk is not None: - chunk = self._reasoning_part_first(buffered_chunk) + chunk = self._ordered_chunk(buffered_chunk) elif self._upstream_exhausted: raise StopIteration else: - chunk = self._reasoning_part_first(self.litellm_custom_stream_wrapper.__next__()) + chunk = self._ordered_chunk(self.litellm_custom_stream_wrapper.__next__()) for src in ( getattr(chunk, "provider_specific_fields", None), getattr( @@ -1356,7 +1420,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: @@ -1388,6 +1453,7 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): messages=LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( # pyright: ignore[reportUnknownMemberType] # the transformer's signature is partially untyped input=self.request_input, responses_api_request=dict(self.responses_api_request), + only_since_compaction=True, ), logging_obj=self.litellm_logging_obj, ) 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 d9ebf50cc1a..02d9566d396 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -8,7 +8,7 @@ import uuid from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence from dataclasses import dataclass from functools import partial -from itertools import accumulate, chain, islice +from itertools import accumulate, chain, groupby, islice from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -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 @@ -120,6 +121,17 @@ from litellm.types.utils import ( Usage, ) +from .compaction_items import ( + carries_only_compaction, + compaction_output_item, + decode_compaction_item, + holds_compaction_summary, + provider_compaction_summary, + provider_takes_compaction_blocks, + requests_compaction, + with_answered_turns_after_newest_compaction, + without_items_before_compaction, +) from .custom_tools import ( convert_custom_tool_to_function_tool, extract_custom_tool_names, @@ -224,6 +236,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) @@ -306,6 +319,126 @@ 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 _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 _call_ids_at_turn(indexed_message: tuple[int, object]) -> tuple[tuple[str | None, int], ...]: + index, message = indexed_message + return tuple((_str_field(call, "id"), index) for call in _assistant_tool_calls(message)) + + +def _newest_stored_calls_the_input_answers( + session_history: Sequence[SessionMessage], answered: frozenset[str] +) -> Iterator[SessionMessage | None]: + newest_turn_by_call_id: Final = MappingProxyType( + dict(chain.from_iterable(map(_call_ids_at_turn, enumerate(session_history)))) + ) + answered_turns: Final = tuple( + sorted((newest_turn_by_call_id[call_id], call_id) for call_id in answered if call_id in newest_turn_by_call_id) + ) + answered_by_turn: Final = MappingProxyType( + { + turn: frozenset(call_id for _, call_id in calls) + for turn, calls in groupby(answered_turns, key=lambda turn_and_call: turn_and_call[0]) + } + ) + return ( + _stored_turn_with_calls(message, answered_by_turn.get(index, frozenset())) + for index, message in enumerate(session_history) + ) + + +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 _newest_stored_calls_the_input_answers(session_history, answered) 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,) + 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,) + + def _input_item_cache_control(input_item: Mapping[str, object]) -> ChatCompletionCachedContent | None: cache_control: Final = input_item.get("cache_control") if not cache_control: @@ -340,6 +473,32 @@ def without_item_cache_control(input: str | ResponseInputParam) -> str | Respons ) +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 _input_to_read( + input: str | ResponseInputParam, only_since_compaction: bool, native_compaction: bool +) -> str | ResponseInputParam: + if only_since_compaction: + return without_items_before_compaction(input) + if native_compaction: + return with_answered_turns_after_newest_compaction(input) + return input + + +def _replays_compaction_block(input: str | ResponseInputParam, context_management: object) -> bool: + return requests_compaction(context_management) and holds_compaction_summary(input) + + _PROVIDERS_KEEPING_MESSAGE_CACHE_CONTROL: Final = frozenset(("anthropic", "dashscope", "minimax", "sap", "zai")) @@ -426,9 +585,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 @@ -870,14 +1026,21 @@ 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")) ) + request_input: Final = with_compaction_items_as_params(input) + native_compaction: Final = _replays_compaction_block( + request_input, context_management + ) and provider_takes_compaction_blocks(model, custom_llm_provider) bridged_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=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 = { @@ -909,7 +1072,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, @@ -939,6 +1102,8 @@ class LiteLLMCompletionResponsesConfig: responses_api_request: ResponsesAPIOptionalRequestParams | dict, replay_reasoning: bool = False, item_cache_control: bool = False, + native_compaction: bool = False, + only_since_compaction: bool = False, ) -> list[ AllMessageValues | GenericChatCompletionMessage @@ -960,6 +1125,19 @@ 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 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 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 @@ -975,10 +1153,13 @@ class LiteLLMCompletionResponsesConfig: ) ) + dumped_input: Final = with_compaction_items_as_params(input) + read_input: Final = _input_to_read(dumped_input, only_since_compaction, native_compaction) 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, ) ) @@ -997,6 +1178,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, @@ -1005,6 +1198,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[str, object]: """ Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages @@ -1027,11 +1221,15 @@ 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) + session_history: Final = chat_completion_session.get("messages") or [] + stored_messages: Final = LiteLLMCompletionResponsesConfig._with_compaction_blocks_the_route_takes( + _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, 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 ) @@ -1062,11 +1260,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 ) @@ -1115,6 +1318,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 ]: @@ -1140,6 +1344,7 @@ class LiteLLMCompletionResponsesConfig: LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( input_item=_input, replay_reasoning=replay_reasoning, + native_compaction=native_compaction, ) ) @@ -1523,7 +1728,7 @@ class LiteLLMCompletionResponsesConfig: """Find the index of the previous assistant message.""" for j in range(current_idx - 1, -1, -1): if messages[j].get("role") == "assistant": - return j + return None if carries_only_compaction(messages[j]) else j return None @staticmethod @@ -1839,6 +2044,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 @@ -1877,6 +2083,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 @@ -1928,6 +2140,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[ChatCompletionUserMessage | 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 (ChatCompletionUserMessage(role="user", content=compaction.summary),) + @staticmethod def _transform_responses_api_message_item( input_item: Mapping[str, object], content: object @@ -3038,6 +3267,7 @@ class LiteLLMCompletionResponsesConfig: | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | ResponseCompactionItem ]: responses_output: list[ GenericResponseOutputItem @@ -3048,8 +3278,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) ) @@ -3115,6 +3347,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/responses/main.py b/litellm/responses/main.py index 3bdb1c9c375..09eea7cb323 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -40,9 +40,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.mcp.round_billing import ( billed_for_every_round, @@ -1669,7 +1675,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: @@ -2516,7 +2524,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 bc30da72df9..0d59c93786c 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1707,6 +1707,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 fe418b5830b..da04e59f18a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -14069,6 +14069,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/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 0ed156524be..089f21a33c1 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -386,6 +386,9 @@ class AgenticLoopState: @dataclass(frozen=True, slots=True, kw_only=True) class RouterState: weights: "RouterWeights | None" = field(default=None, metadata=wire("_router_weights")) + encrypted_content_affinity_pinned: bool | None = field( + default=None, metadata=wire("_encrypted_content_affinity_pinned") + ) fallback_depth: int | None = None max_fallbacks: int | None = None attempted_targets: "AttemptedFallbackTargets | None" = None diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 6aeead7679f..ba422a0a901 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -481,6 +481,11 @@ class ContentThinkingSignatureBlockDelta(TypedDict): signature: str +class ContentCompactionBlockDelta(TypedDict): + type: ReadOnly[Literal["compaction_delta"]] + content: ReadOnly[str | None] + + StreamingContentBlockDeltaType = Literal["text_delta", "input_json_delta", "thinking_delta", "signature_delta"] @@ -493,6 +498,7 @@ class ContentBlockDelta(TypedDict): | ContentCitationsBlockDelta | ContentThinkingBlockDelta | ContentThinkingSignatureBlockDelta + | ContentCompactionBlockDelta ) diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 8dbc0d97d54..62dc152ca84 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, @@ -1488,6 +1489,7 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | ResponseCompactionItem ] ) parallel_tool_calls: bool | None = None diff --git a/tests/benchmarks/test_responses_compaction_benchmarks.py b/tests/benchmarks/test_responses_compaction_benchmarks.py new file mode 100644 index 00000000000..62a5e95255e --- /dev/null +++ b/tests/benchmarks/test_responses_compaction_benchmarks.py @@ -0,0 +1,35 @@ +import json +from collections.abc import Mapping +from typing import Final + +import pytest + +from litellm.responses.litellm_completion_transformation.compaction_items import ( + with_answered_turns_after_newest_compaction, +) + +COMPACTION_ITEM: Final[Mapping[str, object]] = { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Turn one planned Tokyo."}), +} +LONG_TURN_OF_ASSISTANT_TEXT: Final[tuple[Mapping[str, object], ...]] = tuple( + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": f"Step {index}."}]} + for index in range(5_000) +) +CALL: Final[Mapping[str, object]] = { + "type": "function_call", + "call_id": "call_osaka", + "name": "get_weather", + "arguments": "{}", +} +OUTPUT: Final[Mapping[str, object]] = {"type": "function_call_output", "call_id": "call_osaka", "output": "Sunny."} + + +@pytest.mark.benchmark +def test_with_answered_turns_after_newest_compaction_long_turn_split_from_its_output() -> None: + replayed: Final = with_answered_turns_after_newest_compaction( + [*LONG_TURN_OF_ASSISTANT_TEXT, CALL, COMPACTION_ITEM, OUTPUT] + ) + + assert replayed == [COMPACTION_ITEM, *LONG_TURN_OF_ASSISTANT_TEXT, CALL, OUTPUT] 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..5b8e2b534d7 --- /dev/null +++ b/tests/integration/providers/test_responses_compaction_items_wire.py @@ -0,0 +1,1057 @@ +import json +import os +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 + +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), pytest.mark.shares_worker_state] + +_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"} +_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"] + + +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 _fail(marker: str) -> str: + return f"fail-now-{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}" + + +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}) + 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, ...]: + 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( + chain.from_iterable(_anthropic_block_frames(index, block) for index, block in enumerate(content)) + ) + 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) + 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)}) + 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(chain.from_iterable(map(_message_texts, messages))) + + +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"])) + 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(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(chain.from_iterable(map(_message_texts, items))) + + +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], 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) + + +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( + 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") + 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_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: + 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 = _streamed_compaction_blocks(chunks) + 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, **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]] = { + **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, + }, + "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"), + _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 + + +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, **litellm_settings) + with owned_proxy(gateway, directory, {}, config=config, workers=2) as owned: + yield _StoredRig(owned, provider, guardrail) + + +@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"' + " WHERE request_id=%s", + (request_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + 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"), + 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) + + +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) + + +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] diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index a91888fea79..6a734ca3b70 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -57,6 +57,8 @@ def _provider_fields() -> dict[str, object]: "reasoningContentBlocks": [{"reasoningText": {"text": SECRET}}], "citations": [{"cited_text": SECRET}], "thought_signatures": [SIGNATURE], + "compaction_blocks": [{"type": "compaction", "content": SECRET}], + "compaction_delta": {"type": "compaction_delta", "content": SECRET}, } @@ -165,6 +167,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.mark.parametrize("surface", ("typed", "dict", "standard", "callback", "helper")) def test_responses_redaction_removes_instructions_without_changing_the_response(surface: str) -> None: response: Final = litellm.ResponsesAPIResponse.model_validate( 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..2f10fd6fd0d 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 @@ -34,6 +34,7 @@ from litellm.llms.openai.responses.guardrail_translation.handler import ( from litellm.llms.openai.responses.guardrail_translation.tool_merge import merge_guardrailed_tools from litellm.proxy.guardrails.guardrail_hooks.generic_guardrail_api import GenericGuardrailAPI from litellm.types.llms.openai import ChatCompletionToolCallChunk +from litellm.responses.litellm_completion_transformation.compaction_items import compaction_summary from litellm.responses.litellm_completion_transformation.transformation import ( LiteLLMCompletionResponsesConfig, ) @@ -298,6 +299,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""" @@ -3072,6 +3099,21 @@ class DictStructuredMessagesGuardrail(CustomGuardrail): return {**inputs, "structured_messages": {"evaluation": "allowed", "messages": []}} +class NameMaskingStructuredGuardrail(CustomGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + rewritten: Final = [ + {**m, "content": m["content"].replace("Biscuit", "[NAME]")} if isinstance(m.get("content"), str) else m + for m in inputs.get("structured_messages") or [] + ] + return {**inputs, "structured_messages": rewritten} + + def _parallel_tool_call_input() -> list: return [ {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read_a", "arguments": "{}"}, @@ -3298,6 +3340,30 @@ class TestProvenancePatching: assistant_items = [item for item in result["input"] if item.get("role") == "assistant"] assert [_texts(item) for item in assistant_items] == [[COMPRESSED_MARKER]] + @pytest.mark.asyncio + async def test_masked_compaction_summary_stays_a_compaction_item_that_drops_older_history(self) -> None: + data: Final = { + "model": "claude-sonnet-5-5", + "input": [ + {"role": "user", "content": "My cat is Biscuit. " * 200}, + {"role": "assistant", "content": "Noted."}, + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "Cat is Biscuit."}), + }, + {"role": "user", "content": "Cat name?"}, + ], + } + + result: Final = await OpenAIResponsesHandler().process_input_messages(data, NameMaskingStructuredGuardrail()) + + assert compaction_summary(result["input"][2]) == "Cat is [NAME]." + provider_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=result["input"], responses_api_request={}, only_since_compaction=True + ) + assert [message["content"] for message in provider_messages] == ["Cat is [NAME].", "Cat name?"] + def test_provenance_bails_on_non_mapping_item(self): from litellm.llms.openai.responses.guardrail_translation.handler import _input_item_provenance @@ -4174,3 +4240,117 @@ 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 + + @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." + + @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/guardrails/test_content_utils.py b/tests/unit/proxy/guardrails/test_content_utils.py index 920ffc77095..c14375135dc 100644 --- a/tests/unit/proxy/guardrails/test_content_utils.py +++ b/tests/unit/proxy/guardrails/test_content_utils.py @@ -1,5 +1,7 @@ """Tests for the shared guardrail content extraction helpers.""" +import json + from litellm.proxy.guardrails._content_utils import ( apply_redacted_messages_back, build_inspection_messages, @@ -167,6 +169,21 @@ def test_iter_message_text_inspects_reasoning_content_and_summary(): assert list(iter_message_text(data)) == ["content secret", "summary secret"] + +def test_iter_message_text_inspects_compaction_summary(): + data = { + "input": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": "banned word"}), + }, + {"role": "user", "content": "next turn"}, + ] + } + assert list(iter_message_text(data)) == ["banned word", "next turn"] + + # ── walk_user_text ──────────────────────────────────────────────────────────── @@ -347,6 +364,30 @@ def test_walk_user_text_redacts_reasoning_content_and_summary(): assert item["id"] == "rs_1" + +def test_walk_user_text_redacts_compaction_summary_and_keeps_the_block(): + data = { + "input": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps( + {"type": "compaction", "content": "AKIAEXAMPLE summary", "signature": "sig"} + ), + } + ] + } + visited = walk_user_text(data, lambda s: s.replace("AKIAEXAMPLE", "[REDACTED]")) + item = data["input"][0] + assert visited == 1 + assert (item["type"], item["id"]) == ("compaction", "cmp_1") + assert json.loads(item["encrypted_content"]) == { + "type": "compaction", + "content": "[REDACTED] summary", + "signature": "sig", + } + + # ── build_inspection_messages ───────────────────────────────────────────────── 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 bd997f10322..76ad1e08e28 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -415,6 +415,74 @@ 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, 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], + 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": 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, + "metadata": body_metadata, + "litellm_metadata": {"headers": litellm_metadata_headers}, + }, + }, + }, + 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 + assert stored_request["metadata"].get("headers") == caller_metadata.get("headers") + assert stored_request["litellm_metadata"]["headers"] == request_headers + + @pytest.mark.asyncio @pytest.mark.parametrize( "used_client_oauth_token, custom_llm_provider, expected", diff --git a/tests/unit/proxy/hooks/test_tpm_concurrent.py b/tests/unit/proxy/hooks/test_tpm_concurrent.py index bb2a89b1d98..3d1e494e36c 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,50 @@ 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" + ) + + +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/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index abc1fe88c7a..985aee061bc 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -176,6 +176,105 @@ async def test_streaming_upstream_errors_keep_the_client_protocol( assert "error" in events[-1] +_COMPACTION_ITEM: Final = { + "type": "compaction", + "encrypted_content": json.dumps({"type": "compaction", "content": "the cat is named Biscuit"}), +} +_HISTORY_BEFORE_COMPACTION: Final = ( + {"role": "user", "content": "a long turn " * 200}, + {"role": "assistant", "content": "a long reply " * 200}, +) +_TURN_AFTER_COMPACTION: Final = {"role": "user", "content": "what is the cat's name?"} + + +async def _cursor_stream_failure_prompt_tokens( + monkeypatch: pytest.MonkeyPatch, request_input: list[Mapping[str, str]], fails_after_output: bool +) -> int: + import litellm.proxy.proxy_server as ps + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + class _FailureRowRecorder(CustomLogger): + prompt_tokens: int | None = None + + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: object, + traceback_str: str | None = None, + ) -> None: + usage: Final = request_data.get("combined_usage_object") + self.prompt_tokens = usage.prompt_tokens if isinstance(usage, litellm.Usage) else None + + model: Final = "gpt-6-astra" + upstream_url: Final = "https://streaming.example/v1" + created: Final = { + "type": "response.created", + "sequence_number": 0, + "response": { + "id": "resp_upstream", + "object": "response", + "created_at": 1, + "status": "in_progress", + "model": model, + "output": [], + }, + } + output_delta: Final = { + "type": "response.output_text.delta", + "sequence_number": 1, + "item_id": "msg_partial", + "output_index": 0, + "content_index": 0, + "delta": "Biscuit", + } + failed: Final = {"type": "error", "error": {"message": "upstream broke", "code": "server_error", "type": None}} + events: Final = (created, output_delta, failed) if fails_after_output else (created, failed) + wire: Final = "".join("data: " + json.dumps(event) + "\n\n" for event in events) + recorder: Final = _FailureRowRecorder() + router: Final = litellm.Router( + model_list=[ + { + "model_name": model, + "litellm_params": {"model": "openai/" + model, "api_base": upstream_url, "api_key": "fixture-key"}, + } + ], + num_retries=0, + ) + monkeypatch.setattr(ps, "llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(app.dependency_overrides, user_api_key_auth, _auth_override) + with respx.mock as transport: + transport.post(upstream_url + "/responses").respond( + 200, content=wire, headers={"Content-Type": "text/event-stream"} + ) + async with httpx.AsyncClient(transport=httpx.ASGITransport(app), base_url="http://testserver") as client: + result: Final = await client.post( + "/cursor/chat/completions", json={"model": model, "stream": True, "input": request_input} + ) + assert recorder.prompt_tokens is not None, result.text + return recorder.prompt_tokens + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fails_after_output", [False, True]) +async def test_a_failed_cursor_stream_counts_its_input_from_the_newest_compaction_item( + monkeypatch: pytest.MonkeyPatch, fails_after_output: bool +) -> None: + """Cursor's route maps to no call type, yet it dispatches a Responses call, which sends + the provider only the newest compaction item onward, so the failure row counts that much.""" + resent: Final = await _cursor_stream_failure_prompt_tokens( + monkeypatch, [*_HISTORY_BEFORE_COMPACTION, _COMPACTION_ITEM, _TURN_AFTER_COMPACTION], fails_after_output + ) + dispatched_after_trim: Final = await _cursor_stream_failure_prompt_tokens( + monkeypatch, [_COMPACTION_ITEM, _TURN_AFTER_COMPACTION], fails_after_output + ) + + assert resent == dispatched_after_trim + + @pytest.mark.asyncio async def test_responses_api_background_polling_rejects_missing_input(): from fastapi import Response as FastAPIResponse @@ -4832,3 +4931,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/proxy/spend_tracking/test_input_tokens.py b/tests/unit/proxy/spend_tracking/test_input_tokens.py index a1da6c79cd7..71f4dec2345 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,59 @@ 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( + "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, earlier_item: dict[str, object]) -> None: + compacted_body: Final = {"model": model, "input": [_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,) + ) + 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] + + +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/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 55358a3b0de..78f5e1ace0c 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,23 @@ 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 KUBERNETES_POD_ROUTING_KEY, 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 ( + KUBERNETES_POD_ROUTING_KEY, + 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 @@ -366,7 +381,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, @@ -380,7 +398,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( @@ -394,7 +416,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"), @@ -732,6 +757,19 @@ def test_large_request_no_truncation_threshold(): assert "skipped" in truncated_content +@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(f") ...{stored_tail}") + + def test_small_request_no_truncation(): """ Test that small strings are not truncated by MAX_STRING_LENGTH_PROMPT_IN_DB @@ -750,9 +788,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): @@ -1144,6 +1180,187 @@ def test_get_response_for_spend_logs_payload_truncates_large_base64(mock_should_ assert parsed["data"][0]["other_field"] == "value" +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 + + 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}" + 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_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 + + 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}) + 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}}, + } + + 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 = 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"] + + +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) + + +@pytest.mark.parametrize("summary_past_the_summary_cap", [False, True]) +def test_responses_shaped_spend_row_keeps_a_compaction_output_item_summary_bounded_like_a_chat_block( + monkeypatch: pytest.MonkeyPatch, summary_past_the_summary_cap: bool +): + from litellm.constants import MAX_STRING_LENGTH_PROMPT_IN_DB + from litellm.responses.litellm_completion_transformation.compaction_items import compaction_summary + from litellm.responses.utils import ResponsesAPIRequestUtils + + monkeypatch.setenv("STORE_PROMPTS_IN_SPEND_LOGS", "true") + filler: Final = "detail " * (MAX_STRING_LENGTH_PROMPT_IN_DB // 7) + summary: Final = ( + _compaction_summary_past_the_summary_cap() + if summary_past_the_summary_cap + else f"{filler}The cat is named Biscuit. {filler}" + ) + payload: Final = cast( + StandardLoggingPayload, + { + "response": { + "object": "response", + "output": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": ResponsesAPIRequestUtils.wrap_encrypted_content_with_model_id( + json.dumps({"type": "compaction", "content": summary}), "deployment-1" + ), + }, + {"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "OK"}]}, + ], + } + }, + ) + + stored_item: Final = json.loads(_get_response_for_spend_logs_payload(payload))["output"][0] + model_id, _ = ResponsesAPIRequestUtils.unwrap_encrypted_content_with_model_id(stored_item["encrypted_content"]) + stored_summary: Final = compaction_summary(stored_item) + assert model_id == "deployment-1" + assert stored_summary is not None + if summary_past_the_summary_cap: + _assert_bounded_like_a_db_string(stored_summary, summary) + return + assert stored_summary == 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/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index a226ce8a4a9..92bf2b45fab 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -4066,6 +4066,72 @@ class TestHasAttributeErrorInChain: assert _has_attribute_error_in_chain(exc_a) is False +_COMPACTION_ITEM: Final = { + "type": "compaction", + "encrypted_content": json.dumps({"type": "compaction", "content": "the cat is named Biscuit"}), +} +_HISTORY_BEFORE_COMPACTION: Final = ( + {"role": "user", "content": "a long turn " * 200}, + {"role": "assistant", "content": "a long reply " * 200}, +) +_TURN_AFTER_COMPACTION: Final = {"role": "user", "content": "what is the cat's name?"} + + +class TestHandleLLMApiExceptionPricesWhatTheServerDispatched: + async def _failed_prompt_tokens( + self, + route_type: Literal["acompletion", "aresponses"], + logged_call_type: str, + request_input: list[dict[str, str]], + ) -> int: + processing_obj: Final = ProxyBaseLLMRequestProcessing(data={"model": "gpt-4o", "input": request_input}) + cursor_caller: Final = ProxyUserAPIKeyAuth(request_route="/cursor/chat/completions") + mock_request: Final = MagicMock(spec=Request) + mock_request.headers = {} + proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging.alert_types = [] + _, logging_obj = await processing_obj.common_processing_pre_call_logic( + request=mock_request, + general_settings={}, + user_api_key_dict=cursor_caller, + proxy_logging_obj=proxy_logging, + proxy_config=MagicMock(spec=ProxyConfig), + route_type=route_type, + ) + logging_obj.model_call_details.update( + first_api_call_start_time=datetime.datetime.now(), + messages=request_input, + call_type=logged_call_type, + ) + with pytest.raises(ProxyException): + await processing_obj.handle_llm_api_exception( + e=Exception("boom"), + user_api_key_dict=cursor_caller, + proxy_logging_obj=proxy_logging, + ) + return processing_obj.data["combined_usage_object"].prompt_tokens + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "route_type, logged_call_type, trims_at_the_compaction", + [("aresponses", "aresponses", True), ("acompletion", "responses", False)], + ) + async def test_a_failed_call_counts_from_the_newest_compaction_only_when_dispatched_as_responses( + self, + route_type: Literal["acompletion", "aresponses"], + logged_call_type: str, + trims_at_the_compaction: bool, + ) -> None: + async def prompt_tokens(request_input: list[dict[str, str]]) -> int: + return await self._failed_prompt_tokens(route_type, logged_call_type, request_input) + + resent: Final = await prompt_tokens([*_HISTORY_BEFORE_COMPACTION, _COMPACTION_ITEM, _TURN_AFTER_COMPACTION]) + dispatched_after_trim: Final = await prompt_tokens([_COMPACTION_ITEM, _TURN_AFTER_COMPACTION]) + + assert (resent == dispatched_after_trim) is trims_at_the_compaction + assert resent >= dispatched_after_trim + + @pytest.mark.asyncio class TestHandleLLMApiExceptionDictDetail: """ @@ -10817,6 +10883,7 @@ class TestDetachedStreamFailureHook: "user_api_key_dict": user_api_key_dict, "original_exception": failure, "request_data": processing_obj.data, + "dispatched_call_type": "anthropic_messages", } ] diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 3f8e3bc93d1..d9def23bde8 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -1229,6 +1229,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "litellm_metadata": copy.deepcopy(malicious_metadata), "weights": {"gpt-3.5-turbo": {"forged-deployment-id": 100}}, "_router_weights": {"gpt-3.5-turbo": {"forged-deployment-id": 100}}, + "_encrypted_content_affinity_pinned": True, } updated = await add_litellm_data_to_request( @@ -1251,6 +1252,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): assert "_router_weights" not in updated assert "weights" not in updated["proxy_server_request"]["body"] assert "_router_weights" not in updated["proxy_server_request"]["body"] + assert "_encrypted_content_affinity_pinned" not in updated stripped_keys = { "disable_global_guardrails", diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index ab6a874a010..b5940e4c212 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -1,5 +1,6 @@ import asyncio import datetime as real_datetime +import json import smtplib import traceback from collections.abc import Mapping @@ -745,7 +746,9 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: post-dispatch failures keep zero spend. """ - async def _run(self, request_data, original_exception: Exception | None = None): + async def _run( + self, request_data, original_exception: Exception | None = None, request_route=None, dispatched_call_type=None + ): from unittest.mock import AsyncMock, patch from litellm.proxy._types import UserAPIKeyAuth @@ -756,7 +759,8 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: await proxy_logging_obj.post_call_failure_hook( request_data=request_data, original_exception=original_exception or Exception("boom"), - user_api_key_dict=UserAPIKeyAuth(), + user_api_key_dict=UserAPIKeyAuth(request_route=request_route), + dispatched_call_type=dispatched_call_type, ) def _timeout(self, model: str, cause: Exception) -> Exception: @@ -773,6 +777,38 @@ class TestPostCallFailureHookEstimatesDispatchedInputTokens: logging_obj.model_call_details = model_call_details return logging_obj + @pytest.mark.asyncio + @pytest.mark.parametrize( + "request_route, dispatched_call_type", + [("/v1/responses", None), ("/cursor/chat/completions", "aresponses")], + ) + async def test_a_failed_responses_request_counts_its_input_from_the_newest_compaction_item( + self, request_route, dispatched_call_type + ): + from datetime import datetime + + async def prompt_tokens(request_input): + request_data = { + "litellm_logging_obj": self._logging_obj( + { + "first_api_call_start_time": datetime.now(), + "model": FAILURE_USAGE_MODEL, + "messages": request_input, + "call_type": "aresponses", + } + ), + "metadata": {}, + } + await self._run(request_data, request_route=request_route, dispatched_call_type=dispatched_call_type) + return request_data["combined_usage_object"].prompt_tokens + + resent = await prompt_tokens( + [COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION] + ) + + assert resent == await prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) + assert resent < await prompt_tokens([COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, TURN_AFTER_COMPACTION]) + @pytest.mark.asyncio async def test_timed_out_request_estimates_input_tokens_priced_at_input_rate(self): from datetime import datetime @@ -2164,29 +2200,29 @@ def test_a_system_prompt_reads_the_same_whatever_shape_it_arrived_in(system_inpu def test_a_system_prompt_is_counted_on_top_of_the_request(): - from litellm.proxy.utils import _count_request_input_tokens + from litellm.proxy.utils import count_request_input_tokens - without = _count_request_input_tokens(FAILURE_USAGE_MODEL, "hello world", None) - with_system = _count_request_input_tokens(FAILURE_USAGE_MODEL, "hello world", "be brief") + without = count_request_input_tokens(FAILURE_USAGE_MODEL, "hello world", None) + with_system = count_request_input_tokens(FAILURE_USAGE_MODEL, "hello world", "be brief") assert without > 0 assert with_system > without def test_a_request_with_nothing_in_it_counts_zero(): - from litellm.proxy.utils import _count_request_input_tokens + from litellm.proxy.utils import count_request_input_tokens - assert _count_request_input_tokens(FAILURE_USAGE_MODEL, [], None) == 0 - assert _count_request_input_tokens(FAILURE_USAGE_MODEL, None, None) == 0 + assert count_request_input_tokens(FAILURE_USAGE_MODEL, [], None) == 0 + assert count_request_input_tokens(FAILURE_USAGE_MODEL, None, None) == 0 def test_a_failed_dispatch_is_estimated_as_input_only(): - from litellm.proxy.utils import _count_request_input_tokens, _estimate_dispatched_failure_usage + from litellm.proxy.utils import count_request_input_tokens, _estimate_dispatched_failure_usage usage = _estimate_dispatched_failure_usage(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) assert usage is not None - assert usage.prompt_tokens == _count_request_input_tokens(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) + assert usage.prompt_tokens == count_request_input_tokens(FAILURE_USAGE_MODEL, ONE_USER_MESSAGE, None) assert usage.completion_tokens == 0 assert usage.total_tokens == usage.prompt_tokens @@ -2450,7 +2486,7 @@ def test_earlier_attempts_are_priced_from_the_prompt_estimate_not_the_final_stre lifted = _failure_fields_to_lift({"litellm_logging_obj": logging_obj}, provider_billed=True) - estimate = _estimated_prompt_usage(logging_obj.model_call_details, {}) + estimate = _estimated_prompt_usage(logging_obj.model_call_details, {}, None) assert lifted["combined_usage_object"] is recovered assert [(a.usage.prompt_tokens, a.usage.completion_tokens) for a in lifted["billed_failed_attempts"]] == [ (estimate.prompt_tokens, 0) @@ -3336,6 +3372,83 @@ def test_a_request_that_timed_out_at_a_provider_bills_its_input_at_the_input_rat assert response_cost == pytest.approx(expected_input_cost) + +COMPACTION_INSTRUCTIONS = {"role": "system", "content": "answer briefly"} +COMPACTED_HISTORY = [ + {"role": "user", "content": "a long turn " * 200}, + {"role": "assistant", "content": "a long reply " * 200}, +] +COMPACTION_ITEM = { + "type": "compaction", + "encrypted_content": json.dumps({"type": "compaction", "content": "the cat is named Biscuit"}), +} +TURN_AFTER_COMPACTION = {"role": "user", "content": "what is the cat's name?"} + + +def _failed_call_prompt_tokens(call_type, dispatched_call_type, request_input): + from litellm.proxy.utils import _failure_usage_to_lift + + lifted = _failure_usage_to_lift( + model_call_details={"call_type": call_type, "model": FAILURE_USAGE_MODEL, "messages": request_input}, + request_body={}, + dispatched=True, + dispatched_call_type=dispatched_call_type, + ) + assert lifted is not None + return lifted[0].prompt_tokens + + +@pytest.mark.parametrize( + "call_type, dispatched_call_type", + [ + ("acompletion", "acompletion"), + ("aanthropic_messages", "anthropic_messages"), + ("responses", "acompletion"), + ], +) +def test_a_failed_call_off_the_responses_route_counts_the_history_before_a_compaction_typed_message( + call_type, dispatched_call_type +): + def prompt_tokens(request_input): + return _failed_call_prompt_tokens(call_type, dispatched_call_type, request_input) + + sent = prompt_tokens([COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) + + assert sent > prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) + + +def test_an_earlier_timed_out_responses_attempt_counts_its_input_from_the_newest_compaction_item(): + from litellm.litellm_core_utils.get_litellm_params import get_litellm_params + from litellm.proxy.utils import _failure_fields_to_lift + + def earlier_attempt_prompt_tokens(request_input): + logging_obj = _dispatched_failure_logging_obj( + FAILURE_USAGE_MODEL, + get_litellm_params( + metadata={ + "attempted_retries": 1, + "model_info": {"id": "dep-compaction-final"}, + "previous_models": ( + {"deployment_id": "dep-compaction-final", "exception_type": "Timeout", "attempted_retries": 0}, + ), + } + ), + ) + logging_obj.model_call_details["call_type"] = "aresponses" + logging_obj.model_call_details["messages"] = request_input + lifted = _failure_fields_to_lift( + {"litellm_logging_obj": logging_obj}, provider_billed=True, dispatched_call_type="aresponses" + ) + (earlier_attempt,) = lifted["billed_failed_attempts"] + return earlier_attempt.usage.prompt_tokens + + sent = earlier_attempt_prompt_tokens( + [COMPACTION_INSTRUCTIONS, *COMPACTED_HISTORY, COMPACTION_ITEM, TURN_AFTER_COMPACTION] + ) + + assert sent == earlier_attempt_prompt_tokens([COMPACTION_INSTRUCTIONS, COMPACTION_ITEM, TURN_AFTER_COMPACTION]) + + @pytest.mark.parametrize( "current_attempt_api_call_start_time", [None, 1700000001.0], 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 94efef050a6..36df738c0bb 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.constants import REDACTED_BY_LITELLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.responses.litellm_completion_transformation.handler import ( @@ -52,6 +54,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 @@ -561,3 +564,637 @@ 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?"},) +_COMPACT_AT_50K: Final = ({"type": "compaction", "compact_threshold": 50000},) + + +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), 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."}] + + +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), + 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), 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"), + (ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, 1, "message"), + (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) + ] + assert completed.output[0].id == item_events[2].item.id + assert _sent_assistant_content(route) == [_COMPACTION_BLOCK, {"type": "text", "text": "The."}] + + +def test_streamed_responses_over_anthropic_answer_after_a_failed_compaction(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_failed_compaction", + "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": None}}, + {"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"}, + ), + ) + ) + + events: Final = tuple( + litellm.responses( + model=_MODEL, + input=list(_FIRST_INPUT), + context_management=list(_COMPACT_AT_50K), + stream=True, + api_key="mock_api_key", + ) + ) + + assert events[-1].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED + assert [item.type for item in events[-1].response.output] == ["message"] + assert events[-1].response.output[0].content[0].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( + ( + { + "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."}] + assert [key for key in json.loads(route.calls.last.request.content) if key.startswith("_")] == [] + + +_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"]) +@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() + + 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 + + +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, + ] + + +_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 + + +_OSAKA_CALL: Final = {**_WEATHER_CALL, "function": {"name": "get_weather", "arguments": '{"city": "Osaka"}'}} +_TOOL_TURN_REUSING_THE_CALL_ID: Final = { + **_TOOL_TURN_BEFORE_COMPACTION, + "request_id": "chatcmpl-osaka-turn", + "proxy_server_request": { + "model": _MODEL, + "input": [{"type": "function_call_output", "call_id": "toolu_weather", "output": "18C and sunny"}], + }, + "response": { + "id": "chatcmpl-osaka-turn", + "object": "chat.completion", + "created": 0, + "model": _MODEL, + "choices": [ + { + "index": 0, + "finish_reason": "tool_calls", + "message": {"role": "assistant", "content": None, "tool_calls": [_OSAKA_CALL]}, + } + ], + }, +} + + +@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_restores_only_the_newest_stored_call_for_a_reused_call_id( + 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": "12C and rain"} + spend_logs: Final = [_TOOL_TURN_BEFORE_COMPACTION, _TOOL_TURN_REUSING_THE_CALL_ID] + + with patch("litellm.proxy.proxy_server.prisma_client", _FakePrismaClient(spend_logs)): + await litellm.aresponses( + model=_MODEL, + previous_response_id="chatcmpl-osaka-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"] + blocks: Final = list( + chain.from_iterable( + ((index, block) for block in _content_blocks(message)) for index, message in enumerate(messages) + ) + ) + tool_uses: Final = [(index, block["id"], block["input"]) for index, block in blocks if block["type"] == "tool_use"] + tool_results: Final = [(index, block["tool_use_id"]) for index, block in blocks if block["type"] == "tool_result"] + assert [(tool_id, tool_input) for _, tool_id, tool_input in tool_uses] == [("toolu_weather", {"city": "Osaka"})] + assert tool_results == [(tool_uses[0][0] + 1, "toolu_weather")] + assert "12C and rain" in json.dumps(messages[tool_uses[0][0] + 1]) + + +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=[*_FIRST_INPUT, {"role": "assistant", "content": "The."}, compaction_item, *_NEXT_INPUT], + api_key="mock_api_key", + ) + + 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_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_count_request.py b/tests/unit/responses/litellm_completion_transformation/test_count_request.py index 464bdb1a367..0477ae15b8e 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_count_request.py +++ b/tests/unit/responses/litellm_completion_transformation/test_count_request.py @@ -10,6 +10,7 @@ from litellm.responses.litellm_completion_transformation.count_request import ( ChatCountRequest, chat_count_request, ) +from litellm.responses.litellm_completion_transformation.transformation import ChatCompletionSession _MARK: Final[dict[str, JsonValue]] = {"type": "ephemeral"} _MARKED_RESPONSES_REQUEST: Final[dict[str, JsonValue]] = { @@ -125,3 +126,77 @@ async def test_claude_count_request_leaves_out_tools_the_deployment_drops() -> N assert isinstance(dropped, ChatCountRequest) assert dropped.tools is None assert dropped == await _claude_count({"input": [_QUESTION]}, policy) + + +_SUMMARY: Final = "Earlier turns covered the Tokyo trip." +_COMPACTED_INPUT: Final[list[JsonValue]] = [ + {"type": "compaction", "id": "cmp_1", "encrypted_content": json.dumps({"type": "compaction", "content": _SUMMARY})}, + {"type": "message", "role": "user", "content": "And Osaka?"}, +] +_COMPACTION_REQUESTS: Final = ( + pytest.param({}, id="summary_route"), + pytest.param({"context_management": [{"type": "compaction", "compact_threshold": 50000}]}, id="native_route"), +) + + +async def _stored_turns_the_compaction_replaces(previous_response_id: str) -> ChatCompletionSession: + return ChatCompletionSession( + messages=[ + {"role": "user", "content": "Weather in SF?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_0", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + ], + }, + ], + litellm_session_id=None, + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("compaction_request", _COMPACTION_REQUESTS) +async def test_claude_count_request_counts_a_compaction_item_as_its_summary_text( + compaction_request: dict[str, JsonValue], +) -> None: + request: Final = await chat_count_request( + model="claude-opus-5-5", + custom_llm_provider="anthropic", + api_base=None, + responses_request={"model": "claude-opus-5-5", "input": _COMPACTED_INPUT, **compaction_request}, + policy=DropPolicy(drop_params=False, additional_drop_params=frozenset()), + ) + + assert isinstance(request, ChatCountRequest) + assert [(message.get("role"), message.get("content")) for message in request.messages] == [ + ("user", _SUMMARY), + ("user", "And Osaka?"), + ] + assert "compaction_blocks" not in json.dumps(request.messages, default=str) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("compaction_request", _COMPACTION_REQUESTS) +async def test_claude_count_request_leaves_out_the_stored_turns_a_compaction_item_replaces( + compaction_request: dict[str, JsonValue], +) -> None: + request: Final = await chat_count_request( + model="claude-opus-5-5", + custom_llm_provider="anthropic", + api_base=None, + responses_request={ + "model": "claude-opus-5-5", + "previous_response_id": "resp_before_the_compaction", + "input": _COMPACTED_INPUT, + **compaction_request, + }, + policy=DropPolicy(drop_params=False, additional_drop_params=frozenset()), + session_lookup=_stored_turns_the_compaction_replaces, + ) + + assert isinstance(request, ChatCountRequest) + counted: Final = json.dumps(request.messages, default=str) + assert _SUMMARY in counted and "And Osaka?" in counted + assert "Weather in SF?" not in counted + assert "call_0" not in counted 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 f34ca7d0ff8..30e13442c92 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 @@ -24,6 +24,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( _bedrock_converse_messages_pt, anthropic_messages_pt, ) +from litellm.responses.litellm_completion_transformation.reasoning_items import encode_thinking_blocks from litellm.responses.litellm_completion_transformation.transformation import ( TOOL_CALLS_CACHE, ChatCompletionSession, @@ -6502,6 +6503,363 @@ 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."} + + +_COMPACTION_REQUESTED: Final = ResponsesAPIOptionalRequestParams( + context_management=[{"type": "compaction", "compact_threshold": 50000}] +) +_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( + ("model", "responses_api_request", "replayed_compaction"), + [ + ( + "anthropic/claude-sonnet-5-5", + _COMPACTION_REQUESTED, + { + "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), + ("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( + 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, + 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=responses_api_request, + ) + + assert [dict(message) for message in result["messages"]] == [ + replayed_compaction, + {"role": "user", "content": "And Osaka?"}, + ] + + +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]}, + ), + {"role": "user", "content": "Day two?"}, + Message(role="assistant", content="Ueno."), + ], + 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"] + ] + + +_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), + *_TURNS_AFTER_THE_COMPACTED_ANSWER, +] + + +@pytest.mark.parametrize( + ("model", "context_management", "replayed"), + [ + ( + "anthropic/claude-sonnet-5-5", + _COMPACTION_REQUESTED["context_management"], + [("assistant", "Day one, Asakusa.", [_COMPACTION_BLOCK]), *_TURNS_AFTER_THE_COMPACTED_ANSWER], + ), + ("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", + 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?"}, + ] + + + +@pytest.mark.parametrize( + ("call", "output"), + [ + pytest.param( + {"type": "function_call", "call_id": "call_weather", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_weather", "output": "Rain."}, + id="function_call", + ), + pytest.param( + {"type": "custom_tool_call", "call_id": "call_patch", "name": "apply_patch", "input": "*** Begin Patch"}, + {"type": "custom_tool_call_output", "call_id": "call_patch", "output": "Done."}, + id="custom_tool_call", + ), + ], +) +def test_summary_replay_keeps_a_call_its_compaction_item_replaces_when_a_later_output_answers_it( + call: Mapping[str, str], output: Mapping[str, str] +) -> None: + messages: 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."}, + call, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + output, + ], + ), + responses_api_request=_COMPACTION_REQUESTED, + )["messages"] + + assert [message["role"] for message in messages] == ["system", "user", "assistant", "tool"] + assert dict(messages[1]) == _SUMMARY_AS_TEXT + assert [tool_call["id"] for tool_call in messages[2]["tool_calls"]] == [output["call_id"]] + assert messages[3]["tool_call_id"] == output["call_id"] + + +_ANSWERED_THINKING: Final = {"type": "thinking", "thinking": "Need Osaka weather.", "signature": "sig_osaka"} +_CHECKING_OSAKA: Final = { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "Checking Osaka."}], +} + + +def _reasoning_item(item_id: str, thinking: Mapping[str, str]) -> Mapping[str, object]: + return {"type": "reasoning", "id": item_id, "summary": [], "encrypted_content": encode_thinking_blocks([thinking])} + + +def test_summary_replay_keeps_only_the_reasoning_of_the_turn_whose_call_a_later_output_answers() -> None: + tokyo_thinking: Final = {"type": "thinking", "thinking": "Need Tokyo weather.", "signature": "sig_tokyo"} + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( + ResponseInputParam, + [ + {"role": "user", "content": "Plan a Tokyo trip."}, + _reasoning_item("rs_tokyo", tokyo_thinking), + {"type": "function_call", "call_id": "call_tokyo", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_tokyo", "output": "Sun."}, + _reasoning_item("rs_osaka", _ANSWERED_THINKING), + _CHECKING_OSAKA, + {"type": "function_call", "call_id": "call_osaka", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call", "call_id": "call_kyoto", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_kyoto", "output": "Clouds."}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"type": "function_call_output", "call_id": "call_osaka", "output": "Rain."}, + ], + ), + responses_api_request=ResponsesAPIOptionalRequestParams(), + )["messages"] + + assert [message["role"] for message in messages] == ["user", "assistant", "tool"] + assert dict(messages[0]) == _SUMMARY_AS_TEXT + assert messages[1]["thinking_blocks"] == [_ANSWERED_THINKING] + assert [tool_call["id"] for tool_call in messages[1]["tool_calls"]] == ["call_osaka"] + + +def test_native_compaction_replay_moves_a_call_its_compaction_item_splits_from_its_output_after_the_block() -> None: + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( + ResponseInputParam, + [ + {"role": "user", "content": "Plan a Tokyo trip."}, + _reasoning_item("rs_osaka", _ANSWERED_THINKING), + _CHECKING_OSAKA, + {"type": "function_call", "call_id": "call_osaka", "name": "get_weather", "arguments": "{}"}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"type": "function_call_output", "call_id": "call_osaka", "output": "Rain."}, + ], + ), + responses_api_request=_COMPACTION_REQUESTED, + )["messages"] + + assert [message["role"] for message in messages] == ["user", "assistant", "assistant", "tool"] + assert messages[0]["content"] == "Plan a Tokyo trip." + assert messages[1]["provider_specific_fields"] == {"compaction_blocks": [_COMPACTION_BLOCK]} + assert messages[2]["thinking_blocks"] == [_ANSWERED_THINKING] + assert [tool_call["id"] for tool_call in messages[2]["tool_calls"]] == ["call_osaka"] + assert messages[3]["tool_call_id"] == "call_osaka" + + +@pytest.mark.parametrize( + "responses_api_request", + [ + pytest.param(_COMPACTION_REQUESTED, id="native_compaction_block"), + pytest.param(ResponsesAPIOptionalRequestParams(), id="summary_as_text"), + ], +) +def test_compaction_replay_invents_no_call_for_an_output_whose_call_was_never_stored( + responses_api_request: ResponsesAPIOptionalRequestParams, +) -> None: + async def no_stored_turns(previous_response_id: str) -> ChatCompletionSession: + return ChatCompletionSession(messages=[], litellm_session_id=None) + + chat_request: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( + ResponseInputParam, + [ + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"type": "function_call_output", "call_id": "call_never_stored", "output": "Rain."}, + ], + ), + responses_api_request=ResponsesAPIOptionalRequestParams( + **responses_api_request, + tools=[{"type": "function", "name": "get_weather", "parameters": {"type": "object"}, "strict": False}], + ), + ) + + request: Final = asyncio.run( + LiteLLMCompletionResponsesConfig.async_responses_api_session_handler( + previous_response_id="resp_unknown", + litellm_completion_request=chat_request, + instructions=None, + session_lookup=no_stored_turns, + ) + ) + + assert not any(message.get("tool_calls") for message in request["messages"]) + + +@pytest.mark.parametrize( + ("responses_api_request", "expected_turns"), + [ + pytest.param( + _COMPACTION_REQUESTED, + [("user", ()), ("assistant", ("call_0",)), ("tool", "Sunny."), ("assistant", ()), ("assistant", ("call_0",)), ("tool", "Rain.")], + id="native_compaction_block", + ), + pytest.param( + ResponsesAPIOptionalRequestParams(), + [("user", ()), ("assistant", ("call_0",)), ("tool", "Rain.")], + id="summary_as_text", + ), + ], +) +def test_compaction_replay_moves_only_the_call_a_later_output_answers_when_an_earlier_call_reused_its_id( + responses_api_request: ResponsesAPIOptionalRequestParams, expected_turns: list[tuple[str, object]] +) -> None: + messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_request_to_chat_completion_request( + model="anthropic/claude-sonnet-5-5", + input=cast( + ResponseInputParam, + [ + {"role": "user", "content": "Plan a Tokyo trip."}, + {"type": "function_call", "call_id": "call_0", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_0", "output": "Sunny."}, + _reasoning_item("rs_osaka", _ANSWERED_THINKING), + _CHECKING_OSAKA, + {"type": "function_call", "call_id": "call_0", "name": "get_weather", "arguments": "{}"}, + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"type": "function_call_output", "call_id": "call_0", "output": "Rain."}, + ], + ), + responses_api_request=responses_api_request, + )["messages"] + + assert [ + ( + message["role"], + message["content"] + if message["role"] == "tool" + else tuple(tool_call["id"] for tool_call in message.get("tool_calls") or ()), + ) + for message in messages + ] == expected_turns + assert messages[-2]["thinking_blocks"] == [_ANSWERED_THINKING] + + def _continued_chain_request(request_kwargs: Mapping[str, object]) -> Mapping[str, object]: async def session_lookup(previous_response_id: str) -> ChatCompletionSession: assert previous_response_id == "resp_turn_1" @@ -6603,3 +6961,47 @@ async def test_bridge_partial_response_on_disconnect_bills_the_relayed_chat_chun assert incomplete_event.response.usage is not None assert incomplete_event.response.usage.output_tokens == litellm.token_counter(model="gpt-4o-mini", text=relayed_text) assert incomplete_event.response.usage.input_tokens > 0 + + +async def _input_tokens_billed_on_disconnect(request_input: ResponseInputParam) -> int: + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + + chat_stream: Final = await litellm.acompletion( + model="anthropic/claude-sonnet-5-5", + messages=[{"role": "user", "content": "Which day is driest?"}], + mock_response="Thursday looks driest, so plan the walking tour then.", + stream=True, + api_key="test-key", + ) + iterator: Final = LiteLLMCompletionStreamingIterator( + model="anthropic/claude-sonnet-5-5", + litellm_custom_stream_wrapper=chat_stream, + request_input=request_input, + responses_api_request={}, + custom_llm_provider="anthropic", + ) + _ = [await iterator.__anext__() for _ in range(6)] + incomplete_event: Final = iterator.partial_response_on_disconnect() + assert incomplete_event is not None and incomplete_event.response.usage is not None + return incomplete_event.response.usage.input_tokens + + +@pytest.mark.asyncio +async def test_bridge_partial_response_on_disconnect_bills_only_the_input_since_the_newest_compaction_item(): + since_compaction: Final = cast( + ResponseInputParam, + [ + {"type": "compaction", "id": "cmp_litellm", "encrypted_content": json.dumps(_COMPACTION_BLOCK)}, + {"role": "user", "content": "Which day is driest?"}, + ], + ) + history_it_replaces: Final = cast( + ResponseInputParam, + [{"role": "user", "content": "Plan a Tokyo trip, day by day, with every museum and food stall. " * 200}], + ) + + assert await _input_tokens_billed_on_disconnect( + cast(ResponseInputParam, [*history_it_replaces, *since_compaction]) + ) == await _input_tokens_billed_on_disconnect(since_compaction) 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 7a146c1b665..0eb1aa26a65 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 @@ -18,6 +19,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 +235,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 +292,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 +358,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 +406,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 +500,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 +529,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 +596,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 +698,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 +716,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 +775,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( @@ -878,6 +859,286 @@ 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: 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 {} + ) + 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, **compaction_fields}}], + }, + } + + +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) + + +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."}] + 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."), + ) + + 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."), + ("user", "And Osaka?"), + ("assistant", "Dotonbori."), + ] + + +@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( + first_turn_input: str | list[dict[str, str]], +): + compaction_blocks: Final = [{"type": "compaction", "content": "Turn one planned Tokyo."}] + spend_logs: Final = ( + _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.", + compaction_blocks, + ), + _turn_spend_log([{"role": "user", "content": "And Osaka?"}], "Dotonbori."), + ) + + session: Final = await _history_from_stored_turns(spend_logs, "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} + + +_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.")] + + +@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."] + + +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.", + ] + + +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"]) + + +@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:") @@ -889,7 +1150,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)] 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 ea3eab7e002..436367bd82d 100644 --- a/tests/unit/responses/test_responses_api_request_body.py +++ b/tests/unit/responses/test_responses_api_request_body.py @@ -7,6 +7,7 @@ in expected_responses_api_request/. import copy import json import logging +from collections.abc import Awaitable, Callable from pathlib import Path from importlib import import_module from typing import Final @@ -1141,6 +1142,61 @@ async def test_a_native_responses_provider_places_every_point_itself(): 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( + ("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"), + [ + ( + _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, + 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(upstream_url).mock( + return_value=httpx.Response(200, json=_minimal_responses_api_payload("resp_compacted", "gpt-5.6")) + ) + + 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 + + def test_non_streamed_response_ignores_a_stream_timeout_it_never_uses() -> None: answer: Final = json.dumps(_minimal_responses_api_payload("resp_unstreamed", "gpt-5.6")).encode() client: Final = HTTPHandler( diff --git a/tests/unit/responses/test_streaming_iterator.py b/tests/unit/responses/test_streaming_iterator.py index 62d89a55e85..6c3e6177503 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -2605,6 +2605,27 @@ def test_faked_stream_builds_each_refusal_summary_and_custom_tool_input_once_for } +_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." + ) + + @pytest.mark.asyncio async def test_shutdown_drain_waits_for_a_pending_logging_task(): release: Final = asyncio.Event() diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index e1f748fcb58..9cc1bf6ce79 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -7889,6 +7889,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), diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index 5a2cefd08b1..b0086662f1c 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -254,6 +254,7 @@ INTERNAL_STATE_NAMES: Final = ( "_context_compaction_state", *AGENTIC_LOOP_STATE_NAMES, "_router_weights", + "_encrypted_content_affinity_pinned", "fallback_depth", "max_fallbacks", "attempted_targets", diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2b174a1141a..bd08b743c7d 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -44279,7 +44279,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 a31c5230de303d342f62e0e2ecd62af83e5fcd4e Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:42:05 +0000 Subject: [PATCH 2/2] fix(proxy): forward dispatched_call_type through _failure_usage_to_lift (#47815) --- litellm/proxy/utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 52b3866d04d..28be11509ae 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1294,6 +1294,7 @@ def _failure_usage_to_lift( provider_billed: bool = False, deployment_id: str | None = None, base_model: str | None = None, + dispatched_call_type: str | None = None, ) -> tuple[object, object] | None: current_attempt: Final = _current_failure_usage_to_lift( model_call_details, @@ -1302,6 +1303,7 @@ def _failure_usage_to_lift( provider_billed, deployment_id, base_model, + dispatched_call_type=dispatched_call_type, ) failed_attempts: Final = _FAILED_ATTEMPTS_USAGE.validate_python(model_call_details.get("failed_attempts_usage", ())) if not failed_attempts: