diff --git a/litellm/constants.py b/litellm/constants.py index da95b97036a..dd6f3d8f411 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -583,6 +583,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 c2cca5739cf..0cff0a0b0ef 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) from litellm.litellm_core_utils.served_output_texts import SERVED_OUTPUT_TEXTS_KEY +from litellm.llms.base_llm.guardrail_translation.utils import stream_item_field, stream_item_items from litellm.llms.vertex_ai.common_utils import ( redact_vertex_ai_metadata_from_litellm_params, redact_vertex_ai_metadata_from_logged_object, @@ -39,6 +40,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObject = Any +MESSAGE_REDACTION_ENABLE_HEADERS: Final = ("litellm-enable-message-redaction", "x-litellm-enable-message-redaction") + def redact_message_input_output_from_custom_logger( litellm_logging_obj: LiteLLMLoggingObject, result, custom_logger: CustomLogger @@ -147,33 +150,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 @@ -182,6 +197,7 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str): for output_item in output_items: if not isinstance(output_item, dict): continue + output_type = stream_item_field(output_item, "type") if output_item.get("text") is not None: output_item["text"] = redacted_str @@ -193,14 +209,16 @@ def _redact_responses_api_output_dict(output_items, redacted_str: str): if isinstance(content_item, dict) and content_item.get("refusal") is not None: content_item["refusal"] = redacted_str - if output_item.get("type") == "reasoning" and isinstance(output_item.get("summary"), list): + if output_type == "reasoning" and isinstance(output_item.get("summary"), list): for summary_item in output_item["summary"]: if isinstance(summary_item, dict) and summary_item.get("text") is not None: summary_item["text"] = redacted_str - if output_item.get("type") == "function_call" and "arguments" in output_item: + if output_type == "compaction" and "encrypted_content" in output_item: + output_item["encrypted_content"] = redacted_str + if output_type == "function_call" and "arguments" in output_item: output_item["arguments"] = redacted_str - if output_item.get("type") == "custom_tool_call" and "input" in output_item: + if output_type == "custom_tool_call" and "input" in output_item: output_item["input"] = redacted_str @@ -442,13 +460,8 @@ def should_redact_message_logging(model_call_details: dict) -> bool: # User explicitly disabled redaction via header return False - possible_enable_headers: Final = [ - "litellm-enable-message-redaction", # old header. maintain backwards compatibility - "x-litellm-enable-message-redaction", # new header - ] - is_redaction_enabled_via_header = False - for header in possible_enable_headers: + for header in MESSAGE_REDACTION_ENABLE_HEADERS: if bool(request_headers.get(header, False)): is_redaction_enabled_via_header = True break diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 718536120b6..9a335ab4b4a 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -855,11 +855,21 @@ class ModelResponseIterator: if reasoning_content is None: reasoning_content = "" elif "content" in content_block["delta"] and content_block["delta"].get("type") == "compaction_delta": - # Handle compaction delta + compaction_delta: Final = content_block["delta"]["content"] provider_specific_fields["compaction_delta"] = { "type": "compaction_delta", - "content": content_block["delta"]["content"], + "content": compaction_delta, } + if self.compaction_blocks: + streamed_summary: Final = self.compaction_blocks[-1].get("content") + self.compaction_blocks = [ + *self.compaction_blocks[:-1], + { + **self.compaction_blocks[-1], + "content": (streamed_summary if isinstance(streamed_summary, str) else "") + compaction_delta, + }, + ] + provider_specific_fields["compaction_blocks"] = self.compaction_blocks return ( text, diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 84ff790ae48..0b0ee4e3357 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 @@ -804,6 +819,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 +869,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 +1142,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 +1244,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 +1294,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 +1521,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 +1650,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 +1723,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 7198d74909c..7a239c7d343 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -27203,6 +27203,9 @@ }, { "$ref": "#/components/schemas/OutputToolSearchCall" + }, + { + "$ref": "#/components/schemas/ResponseCompactionItem" } ] }, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 81e77283ed8..e92c78cff25 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -71,6 +71,7 @@ from litellm.proxy.hooks.batch_enqueued_tokens import ( canonical_provider_batch_id, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.responses.litellm_completion_transformation.compaction_items import billable_input_since_newest_compaction from litellm.router_utils.add_retry_fallback_headers import ( ensure_response_additional_headers, response_has_hidden_params, @@ -1225,7 +1226,10 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): estimable_data.get("input"), ) ) - messages, prompt, input_text = selected_fields + messages, prompt, raw_input = selected_fields + input_text: Final = ( + 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)) @@ -3799,6 +3803,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=data.get("input") or "", responses_api_request=data, + only_since_compaction=True, ) @staticmethod diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 6fbd0b8f1b3..12d71e0f3e6 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -4,6 +4,8 @@ from collections.abc import Callable, Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Final, Protocol, cast +from pydantic import TypeAdapter + import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import batch_cost_is_final @@ -17,6 +19,7 @@ from litellm.litellm_core_utils.core_helpers import ( ) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import guardrail_information_cost +from litellm.litellm_core_utils.redact_messages import MESSAGE_REDACTION_ENABLE_HEADERS from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_checks import ( get_key_object, @@ -87,6 +90,9 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( ) +_OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping) + + def _proxy_stamped_used_client_oauth_token( request_data: Mapping[str, object], request_route: str | None ) -> bool | None: @@ -112,6 +118,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 @@ -226,13 +252,17 @@ class _ProxyDBLogger(CustomLogger): existing_metadata.update(_metadata) litellm_metadata_bucket: Final = request_data.get("litellm_metadata") - if ( - isinstance(litellm_metadata_bucket, dict) - and "standard_logging_guardrail_information" not in existing_metadata - ): - guardrail_info: Final = litellm_metadata_bucket.get("standard_logging_guardrail_information") - if guardrail_info is not None: - existing_metadata["standard_logging_guardrail_information"] = guardrail_info + 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"] = {} @@ -247,7 +277,12 @@ class _ProxyDBLogger(CustomLogger): request_data["litellm_params"]["proxy_server_request"] = ( request_data.get("proxy_server_request") or existing_litellm_params.get("proxy_server_request") or {} ) - request_data["litellm_params"]["metadata"] = existing_metadata + request_data["litellm_params"]["metadata"] = { + **existing_metadata, + **_redaction_opt_in_lifted_from_litellm_metadata( + litellm_metadata, caller_set_headers="headers" in existing_metadata + ), + } # Preserve model name and custom_llm_provider if "model" not in request_data: diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index f032cd79783..39d8fd76b76 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1519,6 +1519,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 d1303b012c8..536950aff45 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -25,6 +25,7 @@ from litellm.constants import ( LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, + MAX_STRING_LENGTH_COMPACTION_SUMMARY_IN_DB, MCP_SPEND_LOG_MODEL_PREFIX, NON_INFERENCE_CALL_TYPES, REDACTED_BY_LITELM_STRING, @@ -67,6 +68,10 @@ from litellm.proxy._types import ( from litellm.proxy.route_llm_request import ProxyModelNotFoundError from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.proxy.utils import PrismaClient, hash_token +from litellm.responses.litellm_completion_transformation.compaction_items import ( + compaction_summary, + summary_only_compaction_content, +) from litellm.types.router import DeploymentTypedDict, LiteLLM_Params from litellm.types.utils import ( BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN, @@ -1187,11 +1192,35 @@ def _is_spend_log_content(parent: Mapping[str, object], key: str) -> bool: ) +def _is_compaction_block_summary(parent: Mapping[str, object], key: str, keep_compaction_blocks: bool) -> bool: + return keep_compaction_blocks and key == "content" and parent.get("type") == "compaction" + + +def _truncated_for_db(value: str, max_length: int) -> str: + if len(value) <= max_length: + return value + start_chars: Final = int(max_length * 0.35) + end_chars: Final = int(max_length * 0.65) + return ( + 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 _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. @@ -1202,11 +1231,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: @@ -1221,47 +1245,19 @@ 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 _compaction_summary_for_db(v) + if isinstance(v, str) and _is_compaction_block_summary(request_body, k, keep_compaction_blocks) else _sanitize_value(v, mask_credentials and not _is_spend_log_content(request_body, k)) for k, v in request_body.items() if k not in _SENSITIVE_REQUEST_BODY_KEYS @@ -1780,6 +1776,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], @@ -1824,8 +1847,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 @@ -1899,7 +1921,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/responses/litellm_completion_transformation/compaction_items.py b/litellm/responses/litellm_completion_transformation/compaction_items.py new file mode 100644 index 00000000000..ecab21ccce4 --- /dev/null +++ b/litellm/responses/litellm_completion_transformation/compaction_items.py @@ -0,0 +1,217 @@ +import json +import uuid +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Final, TypeVar + +from openai.types.responses.easy_input_message_param import EasyInputMessageParam +from openai.types.responses.response_compaction_item import ResponseCompactionItem +from openai.types.responses.response_create_params import ResponseInputParam +from openai.types.responses.response_input_param import ResponseInputItemParam +from pydantic import TypeAdapter, ValidationError + +from litellm.responses.utils import ResponsesAPIRequestUtils + +from .tool_search import chat_config_builds_anthropic_request, provider_chat_config + +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) +_OBJECT_LIST: Final = TypeAdapter(list[object]) +_Item: Final = TypeVar("_Item") +_INSTRUCTION_ROLES: Final = frozenset(("system", "developer")) + + +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 + + +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 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 _items_from_newest_compaction(input: Sequence[_Item]) -> tuple[_Item, ...]: + newest: Final = next( + (index for index in range(len(input) - 1, 0, -1) if compaction_summary(input[index]) is not None), 0 + ) + return (*(item for item in input[:newest] if _is_instruction_item(item)), *input[newest:]) + + +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: ResponseInputItemParam) -> ResponseInputItemParam: + summary: Final = compaction_summary(item) + return item if summary is None else EasyInputMessageParam(role="user", content=summary, type="message") + + +def with_compaction_items_as_summaries(input: str | ResponseInputParam) -> str | ResponseInputParam: + if not isinstance(input, list) or not holds_compaction_summary(input): + return input + return [_as_summary_message(item) for item in _items_from_newest_compaction(input)] diff --git a/litellm/responses/litellm_completion_transformation/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 615bdb5bbfb..924101b39ff 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 +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 @@ -83,6 +181,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") @@ -99,26 +198,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", @@ -137,6 +244,24 @@ class ResponsesSessionHandler: ), ) + @staticmethod + def _history_of_turns( + spend_logs: Sequence["SpendLogsPayload"], + proxy_server_requests: Sequence[Mapping[str, object] | None], + item_cache_control: bool, + ) -> list["SessionMessage"]: + return list( + chain.from_iterable( + ResponsesSessionHandler.extend_chat_completion_message_with_spend_log_payload( + spend_log=spend_log, + proxy_server_request_dict=proxy_server_request, + chat_completion_message_history=[], + item_cache_control=item_cache_control, + ) + for spend_log, proxy_server_request in zip(spend_logs, proxy_server_requests, strict=True) + ) + ) + @staticmethod def extend_chat_completion_message_with_spend_log_payload( spend_log: "SpendLogsPayload", @@ -181,6 +306,7 @@ class ResponsesSessionHandler: responses_api_request={}, replay_reasoning=True, item_cache_control=item_cache_control, + only_since_compaction=True, ) ) @@ -195,6 +321,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 != {}: @@ -341,12 +468,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 5c9e5cac5ab..a3dfe16c8f7 100644 --- a/litellm/responses/litellm_completion_transformation/streaming_iterator.py +++ b/litellm/responses/litellm_completion_transformation/streaming_iterator.py @@ -7,6 +7,7 @@ from typing import Any, Final, cast import litellm from litellm.litellm_core_utils.hidden_params import get_or_create_hidden_params from litellm.main import stream_chunk_builder +from litellm.responses.litellm_completion_transformation.compaction_items import compaction_output_item from litellm.responses.litellm_completion_transformation.custom_tools import ( build_tool_call_item_kwargs, extract_custom_tool_names, @@ -82,6 +83,11 @@ def _output_items_with_id(items: tuple[Any, ...], item_type: str, item_id: str | ) +def _streams_compaction(chunk: ModelResponseStream) -> bool: + delta_fields: Final = chunk.choices[0].delta.provider_specific_fields if chunk.choices else None + return delta_fields is not None and "compaction_blocks" in delta_fields + + def _delta_has_signed_thinking_block(delta: object) -> bool: blocks: Final = getattr(delta, "thinking_blocks", None) or () return any(isinstance(b, dict) and (b.get("signature") or b.get("data")) for b in blocks) @@ -145,6 +151,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 @@ -185,7 +194,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, @@ -193,7 +202,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: @@ -247,6 +259,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)) @@ -255,6 +269,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: @@ -591,6 +644,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, @@ -1003,6 +1062,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) @@ -1090,6 +1153,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, @@ -1138,7 +1202,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( @@ -1194,11 +1258,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( @@ -1355,7 +1419,8 @@ class LiteLLMCompletionStreamingIterator(ResponsesAPIStreamingIterator): "reasoning", self._cached_reasoning_item_id, ) - streamed_items: Final = filterfalse(self._is_unstreamed_empty_message, reasoning_aligned) + compaction_aligned: Final = _output_items_with_id(reasoning_aligned, "compaction", self._compaction_item_id) + streamed_items: Final = filterfalse(self._is_unstreamed_empty_message, compaction_aligned) return tuple(sorted(streamed_items, key=self._streamed_output_position)) def _is_unstreamed_empty_message(self, item: object) -> bool: diff --git a/litellm/responses/litellm_completion_transformation/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 842230699b7..16537f35ec6 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -29,6 +29,7 @@ from openai.types.chat.chat_completion_named_tool_choice_param import ( Function as NamedToolChoiceFunction, ) from openai.types.responses import ResponseFunctionToolCall, ResponseFunctionWebSearch +from openai.types.responses.response_compaction_item import ResponseCompactionItem from openai.types.responses.response_create_params import ResponseInputParam from openai.types.responses.tool_choice_custom_param import ToolChoiceCustomParam from openai.types.responses.tool_choice_function_param import ToolChoiceFunctionParam @@ -120,6 +121,15 @@ from litellm.types.utils import ( Usage, ) +from .compaction_items import ( + compaction_output_item, + decode_compaction_item, + holds_compaction_summary, + provider_compaction_summary, + provider_takes_compaction_blocks, + requests_compaction, + without_items_before_compaction, +) from .custom_tools import ( convert_custom_tool_to_function_tool, extract_custom_tool_names, @@ -224,6 +234,7 @@ TOOL_CALLS_CACHE: Final = InMemoryCache() _ANY_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[object, object]) _STR_KEY_DICT_ADAPTER: Final = TypeAdapter(dict[str, object]) _OBJECT_LIST_ADAPTER: Final = TypeAdapter(list[object]) +_MessageT = TypeVar("_MessageT") _CACHED_CONTENT_ADAPTER: Final = TypeAdapter(ChatCompletionCachedContent) _DICT_ITEMS_LIST_ADAPTER: Final = TypeAdapter(list[dict[object, object]]) _TEXT_ADAPTER: Final = TypeAdapter(str) @@ -306,6 +317,100 @@ 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 with_stored_calls_the_input_answers( + new_messages: Sequence[_MessageT], session_history: Sequence[SessionMessage] +) -> tuple[_MessageT | SessionMessage, ...]: + answered: Final = _answered_tool_call_ids(new_messages) + stored_calls: Final = tuple( + turn for turn in (_stored_turn_with_calls(message, answered) for message in session_history) if turn is not None + ) + if not stored_calls: + return tuple(new_messages) + first_output: Final = next( + index for index, message in enumerate(new_messages) if _tool_output_call_id(message) in answered + ) + return (*new_messages[:first_output], *stored_calls, *new_messages[first_output:]) + + +def _stored_compaction_summary(message: SessionMessage) -> str | None: + return provider_compaction_summary(message.provider_specific_fields) if isinstance(message, Message) else None + + +def _provider_fields_without_compaction(message: Message) -> Mapping[str, object] | None: + if not message.provider_specific_fields: + return None + other_fields: Final = { + key: value + for key, value in _STR_KEY_DICT_ADAPTER.validate_python(message.provider_specific_fields).items() + if key != "compaction_blocks" + } + return other_fields or None + + +def _with_stored_compaction_as_summary(message: SessionMessage) -> tuple[SessionMessage, ...]: + if not isinstance(message, Message): + return (message,) + summary: Final = provider_compaction_summary(message.provider_specific_fields) + if summary is None: + return (message,) + 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: @@ -344,6 +449,22 @@ 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 _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")) @@ -430,9 +551,6 @@ def _block_with_cache_control(block: object, rewrite: _CacheControlRewrite) -> o return _fields_with_cache_control(_STR_KEY_DICT_ADAPTER.validate_python(block), rewrite) -_MessageT = TypeVar("_MessageT") - - def _with_cache_control(message: _MessageT, rewrite: _CacheControlRewrite) -> _MessageT: if not isinstance(message, dict): return message @@ -874,14 +992,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 = { @@ -913,7 +1038,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, @@ -943,6 +1068,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 @@ -964,6 +1091,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 @@ -979,10 +1119,13 @@ class LiteLLMCompletionResponsesConfig: ) ) + dumped_input: Final = with_compaction_items_as_params(input) + read_input: Final = without_items_before_compaction(dumped_input) if only_since_compaction else dumped_input messages.extend( LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message( - input=input if item_cache_control else without_item_cache_control(input), + input=read_input if item_cache_control else without_item_cache_control(read_input), replay_reasoning=replay_reasoning, + native_compaction=native_compaction, ) ) @@ -1001,6 +1144,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, @@ -1009,6 +1164,7 @@ class LiteLLMCompletionResponsesConfig: session_lookup: Callable[[str], Awaitable[ChatCompletionSession]] | None = None, request_kwargs: Mapping[str, object] = MappingProxyType({}), previous_response_client_ids: tuple[str, ...] = (), + input_replaces_history: bool = False, ) -> dict: """ Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages @@ -1031,11 +1187,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 ) @@ -1066,11 +1226,16 @@ class LiteLLMCompletionResponsesConfig: # Store original _messages before combining for safety check original_new_messages: Final = _messages.copy() if _messages else [] - combined_messages = [*leading_system_messages, *session_messages, *_messages[instructions_end:]] + new_messages: Final = ( + with_stored_calls_the_input_answers(_messages[instructions_end:], session_history) + if input_replaces_history + else _messages[instructions_end:] + ) + combined_messages = [*leading_system_messages, *session_messages, *new_messages] # Fix: Ensure tool_results have corresponding tool_calls in previous assistant message # Pass tools parameter to help reconstruct tool_calls if not in cache - tools: Final = litellm_completion_request.get("tools") or [] + tools: Final = _OBJECT_LIST_ADAPTER.validate_python(litellm_completion_request.get("tools") or []) combined_messages = LiteLLMCompletionResponsesConfig._ensure_tool_results_have_corresponding_tool_calls( messages=combined_messages, tools=tools ) @@ -1119,6 +1284,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 ]: @@ -1144,6 +1310,7 @@ class LiteLLMCompletionResponsesConfig: LiteLLMCompletionResponsesConfig._transform_responses_api_input_item_to_chat_completion_message( input_item=_input, replay_reasoning=replay_reasoning, + native_compaction=native_compaction, ) ) @@ -1843,6 +2010,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 @@ -1881,6 +2049,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 @@ -1932,6 +2106,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 @@ -3042,6 +3233,7 @@ class LiteLLMCompletionResponsesConfig: | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | ResponseCompactionItem ]: responses_output: list[ GenericResponseOutputItem @@ -3052,8 +3244,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) ) @@ -3119,6 +3313,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 e775e155dea..028d09a4e9f 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, @@ -1649,7 +1655,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: @@ -2496,7 +2504,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 3d07f22c81a..a9ec5ab4e63 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -1624,6 +1624,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 65ba67d9ab8..f2e4897e9bc 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -13506,6 +13506,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/llms/openai.py b/litellm/types/llms/openai.py index 8cc9e912e80..520edb2b7dd 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, @@ -1481,6 +1482,7 @@ class ResponsesAPIResponse(BaseLiteLLMOpenAIResponseObject): | ResponseFunctionWebSearch | CustomToolCallOutputItem | OutputToolSearchCall + | ResponseCompactionItem ] ) parallel_tool_calls: bool | None = None 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..f309c374645 --- /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) + +_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 d14f236f765..b98bf974f52 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -55,6 +55,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}, } @@ -163,6 +165,26 @@ def test_streaming_redaction_redacts_delta_provider_fields_without_touching_the_ assert fields == _provider_fields() +@pytest.mark.parametrize("shape", ["object", "dict"]) +def test_redaction_hides_a_compaction_summary_in_responses_api_output(shape: str) -> None: + response: Final = ResponsesAPIResponse( + id="resp_1", + created_at=1, + model="claude-sonnet-5-5", + object="response", + status="completed", + output=[ + {"type": "compaction", "id": "cmp_1", "encrypted_content": json.dumps({"type": "compaction", "content": SECRET})} + ], + ) + result: Final = response if shape == "object" else response.model_dump() + + redacted: Final = perform_redaction({}, result) + + assert SECRET not in str(redacted), redacted + assert SECRET in str(result) + + @pytest.fixture(autouse=True) def _reset_global_redaction(): """Ensure the global setting is off for every test.""" diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index bf871210074..7b0dc1f9922 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -298,6 +298,32 @@ class TestOpenAIResponsesHandlerInputProcessing: assert result["input"][1]["content"] == "World [GUARDRAILED]" assert result["model"] == "gpt-4" + @pytest.mark.asyncio + async def test_compaction_summary_is_scanned_and_rewritten_inside_its_item(self) -> None: + affinity_prefix: Final = "litellm_enc:bW9kZWxfaWQ6ZGVwLWE=;" + data: Final = { + "input": [ + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": affinity_prefix + + json.dumps({"type": "compaction", "content": "Cat is Biscuit."}), + }, + {"role": "user", "content": "Cat name?", "type": "message"}, + ], + "model": "claude-sonnet-5-5", + } + + result: Final = await OpenAIResponsesHandler().process_input_messages(data, MockGuardrail(guardrail_name="t")) + + rewritten: Final = result["input"][0]["encrypted_content"] + assert rewritten.startswith(affinity_prefix) + assert json.loads(rewritten.removeprefix(affinity_prefix)) == { + "type": "compaction", + "content": "Cat is Biscuit. [GUARDRAILED]", + } + assert result["input"][1]["content"] == "Cat name? [GUARDRAILED]" + @pytest.mark.asyncio async def test_process_input_list_with_multimodal_content(self): """Test processing list input with multimodal content""" @@ -4174,3 +4200,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/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 2f0e8a268c2..0a76f7e4a4d 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -162,6 +162,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 fc329eaec5f..0f3ffb834bc 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -4336,3 +4336,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 262357fe1bd..2201150c74c 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -14,8 +14,22 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm import litellm.constants as litellm_constants import litellm.proxy.spend_tracking.spend_tracking_utils as spend_tracking_utils -from litellm.constants import LITELLM_TRUNCATED_PAYLOAD_FIELD, LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, LITTELM_CLI_SERVICE_ACCOUNT_NAME, LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, MAX_SPEND_LOG_MODEL_NAME_LENGTH, REDACTED_BY_LITELM_STRING, SESSION_ID_GENERATED_METADATA_KEY, SESSION_ID_OMITTED_METADATA_KEY, UNKNOWN_MODEL_SPEND_LOG_MODEL -from litellm.constants import CALLER_SESSION_ID_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY +from litellm.constants import ( + LITELLM_TRUNCATED_PAYLOAD_FIELD, + LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE, + LITTELM_CLI_SERVICE_ACCOUNT_NAME, + LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, + MAX_SPEND_LOG_MODEL_NAME_LENGTH, + REDACTED_BY_LITELM_STRING, + SESSION_ID_GENERATED_METADATA_KEY, + SESSION_ID_OMITTED_METADATA_KEY, + UNKNOWN_MODEL_SPEND_LOG_MODEL, +) +from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, + INTERNAL_CALL_ORIGIN_METADATA_KEY, + MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, +) from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.llms.base_llm.ocr.transformation import OCRResponse, OCRUsageInfo from litellm.litellm_core_utils.safe_json_dumps import safe_dumps @@ -279,7 +293,10 @@ def test_caller_session_keeps_the_legacy_trace_id_with_correlation_logging_off(m def test_session_displaced_by_a_trace_header_wins_only_with_correlation_logging( monkeypatch: pytest.MonkeyPatch, correlation_in_logs: bool, expected_session_id: str ): - header_chain: Final = {"litellm_session_id": "trace-abc", "metadata": {"trace_id": "trace-abc", "session_id": "trace-abc"}} + header_chain: Final = { + "litellm_session_id": "trace-abc", + "metadata": {"trace_id": "trace-abc", "session_id": "trace-abc"}, + } session_id: Final = _chat_session_id( monkeypatch, header_chain, @@ -293,7 +310,11 @@ def test_session_displaced_by_a_trace_header_wins_only_with_correlation_logging( ("litellm_params", "metadata"), [ pytest.param({"litellm_session_id": "convo\x00x"}, {}, id="body-session"), - pytest.param({"litellm_session_id": "trace-abc"}, {CALLER_SESSION_ID_METADATA_KEY: "convo\x00x"}, id="beside-trace-header"), + pytest.param( + {"litellm_session_id": "trace-abc"}, + {CALLER_SESSION_ID_METADATA_KEY: "convo\x00x"}, + id="beside-trace-header", + ), ], ) def test_caller_session_postgres_cannot_store_keeps_the_trace_id( @@ -307,7 +328,10 @@ def test_caller_session_postgres_cannot_store_keeps_the_trace_id( [ pytest.param({"metadata": {}}, id="no-session"), pytest.param( - {"litellm_session_id": "minted-1", "metadata": {"session_id": "minted-1", SESSION_ID_GENERATED_METADATA_KEY: True}}, + { + "litellm_session_id": "minted-1", + "metadata": {"session_id": "minted-1", SESSION_ID_GENERATED_METADATA_KEY: True}, + }, id="minted-by-generate-policy", ), pytest.param({"litellm_session_id": "job-7", "metadata": {"trace_id": "job-7"}}, id="copied-from-trace-id"), @@ -644,6 +668,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 @@ -662,9 +699,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): @@ -1056,6 +1091,142 @@ 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) + + @patch("litellm.proxy.spend_tracking.spend_tracking_utils.should_store_prompts_and_responses_in_spend_logs") def test_get_response_for_spend_logs_payload_strips_null_bytes(mock_should_store): """Regression for PostgreSQL 22P05: NUL bytes must be stripped from response.""" diff --git a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py index 025f1315aab..afc0bf644de 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py +++ b/tests/unit/responses/litellm_completion_transformation/test_anthropic_responses_bridge.py @@ -1,15 +1,17 @@ import json from collections.abc import Mapping, Sequence from itertools import chain -from typing import Final +from typing import Final, cast from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx +from openai.types.responses.response_compaction_item import ResponseCompactionItem from pydantic import BaseModel import litellm +from litellm import CustomLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, @@ -51,6 +53,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 @@ -541,3 +544,531 @@ 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."}] + + +@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."}] + + +_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 + + +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_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index cc5d5c62739..d646d597d3d 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -6502,6 +6502,169 @@ 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?"}, + ] + + 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" 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 4a2fc0627ab..b7af1ee4845 100644 --- a/tests/unit/responses/test_responses_api_request_body.py +++ b/tests/unit/responses/test_responses_api_request_body.py @@ -6,6 +6,7 @@ in expected_responses_api_request/. import copy import json +from collections.abc import Awaitable, Callable from pathlib import Path from importlib import import_module from typing import Final @@ -1042,6 +1043,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 7eb3115117d..0a462f819ca 100644 --- a/tests/unit/responses/test_streaming_iterator.py +++ b/tests/unit/responses/test_streaming_iterator.py @@ -2594,3 +2594,24 @@ def test_faked_stream_builds_each_refusal_summary_and_custom_tool_input_once_for (1, "content", 1): "No forecasts for Mars.", (2, "input", 0): "curl wttr.in/Berlin", } + + +_COMPACTION_SUMMARY: Final = "The user planned a Tokyo trip." +_NEXT_TURN: Final = {"type": "message", "role": "user", "content": "And Osaka?"} +_COMPACTED_INPUT: Final = [ + {"type": "message", "role": "user", "content": "Plan a Tokyo trip. " * 200}, + {"type": "message", "role": "assistant", "content": "Day one, Asakusa. " * 200}, + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": _COMPACTION_SUMMARY}), + }, + _NEXT_TURN, +] +_SUMMARY_INPUT: Final = [{"type": "message", "role": "user", "content": _COMPACTION_SUMMARY}, _NEXT_TURN] + + +def test_usage_estimate_counts_a_compaction_item_as_its_summary(): + assert _estimate_usage_from_text("gpt-4o-mini", _COMPACTED_INPUT, {}, "Dotonbori.") == _estimate_usage_from_text( + "gpt-4o-mini", _SUMMARY_INPUT, {}, "Dotonbori." + ) diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 1dec4de45f4..635aa670c63 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -7755,6 +7755,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/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index a111ed1d773..15cdc4b274c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -44043,7 +44043,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 */