diff --git a/litellm/compression/compress.py b/litellm/compression/compress.py index 6a3e44bda65..fd33edb277f 100644 --- a/litellm/compression/compress.py +++ b/litellm/compression/compress.py @@ -222,12 +222,14 @@ def _cached_prefix_indices(messages: Sequence[Mapping[str, object]]) -> tuple[in return tuple(range(last_breakpoint + 1)) -def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int, ...]: +def get_protected_indices( + messages: Sequence[Mapping[str, object]], *, protect_latest_turn: bool = True +) -> tuple[int, ...]: """ Return indices of messages that must never be compressed: - All system messages - - The last user message - - The last assistant message + - The last user message when protect_latest_turn is enabled + - The last assistant message when protect_latest_turn is enabled - Every message up to and including the last one carrying an Anthropic cache_control breakpoint The last user message is what the model is being asked to act on right now, @@ -240,7 +242,8 @@ def get_protected_indices(messages: Sequence[Mapping[str, object]]) -> tuple[int system_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "system") last_user: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "user")[-1:] assistant_indices: Final = tuple(index for index, msg in enumerate(messages) if msg.get("role", "") == "assistant") - return tuple(dict.fromkeys(system_indices + last_user + assistant_indices[-1:] + _cached_prefix_indices(messages))) + latest_indices: Final = last_user + assistant_indices[-1:] if protect_latest_turn else () + return tuple(dict.fromkeys(system_indices + latest_indices + _cached_prefix_indices(messages))) def _combine_scores( diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 1c12ac395c7..2bdefd35a94 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1126,6 +1126,9 @@ class CustomGuardrail(CustomLogger): """ return True + def preserve_request_text_metadata(self) -> bool: + return False + def structured_messages_cover_full_request(self) -> bool: """Whether returned ``structured_messages`` span the whole request. diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 84ff790ae48..c1d81ac286f 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -32,6 +32,7 @@ import copy import json import time import uuid +from collections import Counter from collections.abc import Mapping, MutableSequence, Sequence from dataclasses import dataclass from itertools import accumulate, chain, repeat @@ -266,13 +267,21 @@ def _item_rewrite_field(item: Mapping[str, object]) -> str | None: return _PATCHABLE_ITEM_FIELDS.get(item_type) -def _rewritten_input_item(item: Mapping[str, object], rewritten: object) -> Mapping[str, object] | None: +def _rewritten_input_item( + item: Mapping[str, object], rewritten: object, *, preserve_text_metadata: bool = False +) -> Mapping[str, object] | None: field: Final = _item_rewrite_field(item) if field is None or not isinstance(rewritten, Mapping): return None rewritten_content: Final = rewritten.get("content") if isinstance(item.get(field), str) and isinstance(rewritten_content, str): return {**item, field: rewritten_content} + if preserve_text_metadata: + from litellm.proxy.guardrails.guardrail_hooks.content_text import content_to_text, replace_text_content + + typed_rewritten: Final = _item_fields_adapter.validate_python(rewritten) + replaced: Final = replace_text_content(item.get(field), content_to_text(typed_rewritten.get("content"))) + return {**item, field: replaced} if replaced is not None else None rewritten_row: Final = cast("AllMessageValues", rewritten) # cast-ok: guardrails hand back chat-shaped rows converted_items, _ = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api( [rewritten_row] @@ -472,6 +481,42 @@ def _input_item_provenance( return item_for_message, tainted +_RESPONSE_INPUT_ITEMS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +def rewriteable_response_message_indices( + raw_input: object, instructions: object, original_messages: Sequence[object] +) -> frozenset[int]: + from litellm.proxy.guardrails.guardrail_hooks.content_text import is_rewriteable_text_content + + input_items: Final = [{"role": "user", "content": raw_input}] if isinstance(raw_input, str) else raw_input + if not isinstance(input_items, list): + return frozenset() + try: + items: Final = _RESPONSE_INPUT_ITEMS_ADAPTER.validate_python(input_items) + except ValidationError: + return frozenset() + offset: Final = 1 if instructions else 0 + provenance: Final = _input_item_provenance(items, tuple(original_messages)[offset:]) + if provenance is None: + return frozenset() + item_for_message, tainted = provenance + text_rows: Final = tuple( + (message_index, item_index) + for message_index, item_index in item_for_message.items() + if stream_item_field(original_messages[message_index + offset], "role") != "assistant" + ) + item_counts: Final = Counter(item_index for _, item_index in text_rows) + return frozenset( + message_index + offset + for message_index, item_index in text_rows + if message_index not in tainted + and item_counts[item_index] == 1 + and (field := _item_rewrite_field(items[item_index])) is not None + and is_rewriteable_text_content(items[item_index].get(field)) + ) + + class _RequestFields(NamedTuple): input: tuple[object, ...] instructions: str | None @@ -493,23 +538,30 @@ def _input_item_role(item: object) -> str: return role.lower() if isinstance(role, str) else "" -def _patched_request_fields( +def patch_response_request_fields( raw_input: object, instructions: object, original_messages: Sequence[object], structured_messages: Sequence[object], + *, + preserve_text_metadata: bool = False, + scoped_indices: Sequence[int] | None = None, ) -> _RequestFields | None: - if not isinstance(raw_input, list) or len(original_messages) != len(structured_messages): + input_items: Final = ( + [{"role": "user", "content": raw_input}] if isinstance(raw_input, str) and preserve_text_metadata else raw_input + ) + message_indices: Final = range(len(original_messages)) if scoped_indices is None else scoped_indices + if not isinstance(input_items, list) or len(message_indices) != len(structured_messages): return None offset: Final = 1 if instructions else 0 - provenance: Final = _input_item_provenance(raw_input, tuple(original_messages)[offset:]) + provenance: Final = _input_item_provenance(input_items, tuple(original_messages)[offset:]) if provenance is None: return None item_for_message, tainted = provenance changed: Final = tuple( (index, rewritten) - for index, (original, rewritten) in enumerate(zip(original_messages, structured_messages)) - if original != rewritten + for index, rewritten in zip(message_indices, structured_messages) + if original_messages[index] != rewritten ) instruction_rewrites: Final = tuple(rewritten for index, rewritten in changed if index < offset) rewritten_instructions: Final = ( @@ -526,8 +578,9 @@ def _patched_request_fields( replacements: Final = MappingProxyType( { item_for_message[message_index]: _rewritten_input_item( - cast("Mapping[str, object]", raw_input[item_for_message[message_index]]), # cast-ok: checked Mappings + cast("Mapping[str, object]", input_items[item_for_message[message_index]]), # cast-ok: checked Mappings rewritten, + preserve_text_metadata=preserve_text_metadata, ) for message_index, rewritten in body_changes } @@ -535,7 +588,7 @@ def _patched_request_fields( if len(replacements) != len(body_changes) or any(item is None for item in replacements.values()): return None return _RequestFields( - input=tuple(replacements.get(index, item) for index, item in enumerate(raw_input)), + input=tuple(replacements.get(index, item) for index, item in enumerate(input_items)), instructions=instructions_value, ) @@ -545,11 +598,15 @@ def _patch_or_convert_request_fields( instructions: object, original_messages: Sequence[object], structured_messages: Sequence[AllMessageValues], + *, + preserve_text_metadata: bool = False, ) -> _RequestFields | None: if not isinstance(structured_messages, list): return None - patched: Final = _patched_request_fields(raw_input, instructions, original_messages, structured_messages) - if patched is not None: + patched: Final = patch_response_request_fields( + raw_input, instructions, original_messages, structured_messages, preserve_text_metadata=preserve_text_metadata + ) + if patched is not None or preserve_text_metadata: return patched input_items, converted_instructions = ( LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api(structured_messages) @@ -762,7 +819,11 @@ class OpenAIResponsesHandler(BaseTranslation): ) ) return _patch_or_convert_request_fields( - data.get("input"), data.get("instructions"), structured_messages, merged + data.get("input"), + data.get("instructions"), + structured_messages, + merged, + preserve_text_metadata=guardrail_to_apply.preserve_request_text_metadata(), ) def extract_request_tool_names(self, data: Mapping[str, object]) -> list[str]: diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 545b317055a..f40478993a1 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -17372,6 +17372,12 @@ "description": "Client secret of the gateway's Entra app registration, used to perform the On-Behalf-Of exchange. Falls back to the AGENT365_CLIENT_SECRET environment variable.", "title": "Client Secret" }, + "compress_latest_turn": { + "default": false, + "description": "Allow compression of latest user and tool-result text, which can lose instructions or details. System prompts, assistant messages, cached prefixes and retrieved content remain protected. Unsafe rewrites retain the original input.", + "title": "Compress Latest Turn", + "type": "boolean" + }, "confidence_threshold": { "default": 0.5, "default_value": 0.5, diff --git a/litellm/proxy/guardrails/guardrail_hooks/content_text.py b/litellm/proxy/guardrails/guardrail_hooks/content_text.py index 2af3b0233f2..4dd8b1830d0 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/content_text.py +++ b/litellm/proxy/guardrails/guardrail_hooks/content_text.py @@ -12,9 +12,11 @@ therefore only be written back over a run of text parts, never across a non-text part, which is what ``is_all_text_parts`` gates. """ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Final +from pydantic import TypeAdapter, ValidationError + from litellm.litellm_core_utils.prompt_templates.factory import get_attribute_or_key @@ -94,3 +96,27 @@ def assistant_text_from_response(response: object) -> str | None: if output_parts: return "".join(output_parts) return None + + +_TEXT_PARTS_ADAPTER: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +def is_rewriteable_text_content(content: object) -> bool: + if isinstance(content, str): + return True + if not isinstance(content, list) or len(content) != 1: + return False + try: + parts: Final = _TEXT_PARTS_ADAPTER.validate_python(content) + except ValidationError: + return False + return parts[0].get("type") in ("text", "input_text", "output_text") and isinstance(parts[0].get("text"), str) + + +def replace_text_content(content: object, text: str) -> object: + if not text.strip() or not is_rewriteable_text_content(content): + return None + if isinstance(content, str): + return text + parts: Final = _TEXT_PARTS_ADAPTER.validate_python(content) + return [{**parts[0], "text": text}] diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py index b7e4d9275fa..92c8fef9673 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py @@ -37,6 +37,7 @@ def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail) -> unreachable_fallback=litellm_params.unreachable_fallback, timeout=litellm_params.timeout, ccr_retrieval=litellm_params.ccr_retrieval, + compress_latest_turn=litellm_params.compress_latest_turn, ) litellm.logging_callback_manager.add_litellm_callback( # pyright: ignore[reportUnknownMemberType] _callback diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index dac802a1ddc..cf72b951d76 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -36,7 +36,9 @@ from litellm.proxy.guardrails.guardrail_hooks.content_text import ( assistant_text_from_response, content_to_text, is_all_text_parts, + is_rewriteable_text_content, merge_rewritten_text_parts, + replace_text_content, ) from litellm.proxy.spend_tracking.compression_savings import HEADROOM_GUARDRAIL_PROVIDER from litellm.secret_managers.main import get_secret_str @@ -141,6 +143,34 @@ def _restore_content_shapes( return restored +def _compressed_message_content(original: Mapping[str, object], rewritten: Mapping[str, object]) -> object: + content: Final = rewritten.get("content") + original_content: Final = original.get("content") + if content == original_content or ( + is_all_text_parts(original_content) and content == content_to_text(original_content) + ): + return original_content + if not isinstance(content, str): + return None + return replace_text_content(original.get("content"), content) + + +def _restore_compressed_text( + originals: Sequence[Mapping[str, object]], returned: Sequence[Mapping[str, object]] +) -> tuple[Mapping[str, object], ...] | None: + if any( + original.get("role") != rewritten.get("role") or original.get("tool_call_id") != rewritten.get("tool_call_id") + for original, rewritten in zip(originals, returned) + ): + return None + contents: Final = tuple( + _compressed_message_content(original, rewritten) for original, rewritten in zip(originals, returned) + ) + if any(content is None for content in contents): + return None + return tuple({**original, "content": content} for original, content in zip(originals, contents)) + + def _tool_call_name(tool_call: Mapping[str, object]) -> str | None: function: Final = tool_call.get("function") if not _is_str_object_dict(function): @@ -239,7 +269,10 @@ def _retrieval_result_indices( def _protected_indices( - messages: Sequence[Mapping[str, object]], extra_retrieve_call_ids: frozenset[str] = frozenset() + messages: Sequence[Mapping[str, object]], + extra_retrieve_call_ids: frozenset[str] = frozenset(), + *, + compress_latest_turn: bool = False, ) -> frozenset[int]: """Indices headroom must not send to the compression service. @@ -257,9 +290,9 @@ def _protected_indices( so the model's own earlier tables came back rewritten and it imitated the shape. The tool results those turns asked for stay compressible. """ - protected: Final = frozenset(get_protected_indices(messages)) | _retrieval_result_indices( - messages, extra_retrieve_call_ids - ) + protected: Final = frozenset( + get_protected_indices(messages, protect_latest_turn=not compress_latest_turn) + ) | _retrieval_result_indices(messages, extra_retrieve_call_ids) return ( protected | frozenset( @@ -269,14 +302,19 @@ def _protected_indices( for index in group ) | frozenset(index for index, message in enumerate(messages) if message.get("role") == "assistant") + | frozenset( + index + for index, message in enumerate(messages) + if compress_latest_turn and not is_rewriteable_text_content(message.get("content")) + ) ) def _restore_protected_messages( - messages: Sequence[dict[str, object]], - compressed: Sequence[dict[str, object]], + messages: Sequence[Mapping[str, object]], + compressed: Sequence[Mapping[str, object]], protected_indices: frozenset[int], -) -> Sequence[dict[str, object]]: +) -> Sequence[Mapping[str, object]]: """Put the rows that were held back from compression at their original positions. Requires one returned row per row actually sent, which ``_call_compress`` @@ -505,6 +543,7 @@ class HeadroomGuardrail(CustomGuardrail): unreachable_fallback: str | None = None, timeout: float | None = None, ccr_retrieval: bool = True, + compress_latest_turn: bool = False, ): self.headroom_api_base = (api_base or get_secret_str("HEADROOM_API_BASE") or "").rstrip("/") if not self.headroom_api_base: @@ -518,6 +557,7 @@ class HeadroomGuardrail(CustomGuardrail): "fail_open" if unreachable_fallback == "fail_open" else "fail_closed" ) self.ccr_retrieval = ccr_retrieval + self.compress_latest_turn = compress_latest_turn self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback, ) @@ -530,6 +570,9 @@ class HeadroomGuardrail(CustomGuardrail): ) self.timeout = self._resolve_timeout(timeout) + def preserve_request_text_metadata(self) -> bool: + return self.compress_latest_turn + def _should_bypass(self, request_data: dict) -> bool: psr: Final = request_data.get("proxy_server_request") if not _is_str_object_dict(psr): @@ -677,6 +720,9 @@ class HeadroomGuardrail(CustomGuardrail): ) filtered: Final = [item for item in compressed_messages if _is_str_object_dict(item)] + if self.compress_latest_turn and (len(filtered) != len(messages) or len(filtered) != len(compressed_messages)): + return _CompressResult(messages, False, {}) + if not filtered: return _CompressResult( self._handle_compress_failure( @@ -773,11 +819,28 @@ class HeadroomGuardrail(CustomGuardrail): return str(body) + def _record_skipped_rewrite(self, request_data: dict[str, object], start_time: float, end_time: float) -> None: + from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header + + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={ + "error": "Headroom rewrite could not be mapped safely", + "compression_applied": False, + }, + request_data=request_data, + guardrail_status="guardrail_failed_to_respond", + guardrail_provider=HEADROOM_GUARDRAIL_PROVIDER, + start_time=start_time, + end_time=end_time, + duration=end_time - start_time, + ) + add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) + @log_guardrail_information async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: @@ -800,20 +863,44 @@ class HeadroomGuardrail(CustomGuardrail): if not messages: return inputs - # The last user message is the instruction the model is being asked to - # act on, so replacing it with a marker means the model answers a - # retrieval result instead of the request. Protected rows are held back - # from the payload rather than pinned after the fact, so their tokens - # are not counted as savings we never apply; the Anthropic write-back - # discards a compressed system prompt outright. Keep it that way unless - # /v1/compress grows a field for sending the live turn as the retrieval - # query without compressing it: query-aware compression reads the newest - # user message, so it is withheld here at some cost to history ranking. - # request_data is a bare dict on the base signature; narrow it before - # reading the untranslated messages so long tool names can be recovered. - raw_messages: Final = _REQUEST_DATA_ADAPTER.validate_python(request_data).get("messages") + from litellm.llms.base_llm.guardrail_translation.utils import ( + effective_skip_system_message_for_guardrail, + scoped_structured_message_indices, + ) + from litellm.llms.openai.responses.guardrail_translation.handler import ( + OpenAIResponsesHandler, + patch_response_request_fields, + rewriteable_response_message_indices, + ) + + native_request: Final = self.compress_latest_turn and "input" in request_data + full_messages: Final = ( + OpenAIResponsesHandler().get_structured_messages(request_data) or [] if native_request else [] + ) + scoped_indices: Final = scoped_structured_message_indices( + full_messages, + scan_only_tool_results=False, + skip_system=effective_skip_system_message_for_guardrail(self), + skip_tool=False, + ) + rewriteable: Final = ( + rewriteable_response_message_indices( + request_data.get("input"), request_data.get("instructions"), full_messages + ) + if native_request + else frozenset() + ) + native_protected: Final = frozenset( + index for index, original_index in enumerate(scoped_indices) if original_index not in rewriteable + ) + + typed_request: Final = _REQUEST_DATA_ADAPTER.validate_python(request_data) + raw_messages: Final = typed_request.get("messages") raw_retrieve_call_ids: Final = _raw_retrieve_call_ids(raw_messages) - protected_indices: Final = _protected_indices(messages, raw_retrieve_call_ids) + protected_indices: Final = ( + _protected_indices(messages, raw_retrieve_call_ids, compress_latest_turn=self.compress_latest_turn) + | native_protected + ) compressible: Final = [m for i, m in enumerate(messages) if i not in protected_indices] if not compressible: return inputs @@ -847,12 +934,35 @@ class HeadroomGuardrail(CustomGuardrail): # write-back and restructures it for nothing. return inputs + restored: Final = ( + _restore_compressed_text(compressible, result.messages) + if self.compress_latest_turn + else _restore_content_shapes(originals=compressible, returned=result.messages) + ) + if restored is None: + self._record_skipped_rewrite(request_data, start_time, end_time) + return inputs compressed: Final = _restore_protected_messages( messages=messages, - compressed=_restore_content_shapes(originals=compressible, returned=result.messages), + compressed=restored, protected_indices=protected_indices, ) + if ( + native_request + and patch_response_request_fields( + typed_request.get("input"), + typed_request.get("instructions"), + full_messages, + compressed, + preserve_text_metadata=True, + scoped_indices=scoped_indices, + ) + is None + ): + self._record_skipped_rewrite(request_data, start_time, end_time) + return inputs + self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=result.stats, request_data=request_data, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py index 0df0f0c80ef..988177d903a 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py @@ -31,6 +31,15 @@ class HeadroomGuardrailConfigModel(GuardrailConfigModel[BaseModel]): description="Inject the Headroom retrieval tool for hashes declared by the compression service.", ) + compress_latest_turn: bool = Field( + default=False, + description=( + "Allow compression of latest user and tool-result text, which can lose instructions or details. " + "System prompts, assistant messages, " + "cached prefixes and retrieved content remain protected. Unsafe rewrites retain the original input." + ), + ) + @staticmethod def ui_friendly_name() -> str: return "Headroom" diff --git a/tests/unit/compression/test_compress.py b/tests/unit/compression/test_compress.py index 1765365af65..422794de0e5 100644 --- a/tests/unit/compression/test_compress.py +++ b/tests/unit/compression/test_compress.py @@ -280,3 +280,14 @@ def test_embedding_scorer(respx_mock: respx.MockRouter) -> None: {"role": "user", "content": "Fix auth"}, ] assert result["cache"] == {"message_1": recipe_context} + + +def test_relaxing_latest_turn_protection_retains_system_and_cached_prefix() -> None: + messages: Final = [ + {"role": "system", "content": "instruction"}, + {"role": "user", "content": "cached", "cache_control": {"type": "ephemeral"}}, + {"role": "assistant", "content": "answer"}, + {"role": "user", "content": "latest"}, + ] + assert get_protected_indices(messages, protect_latest_turn=False) == (0, 1) + assert set(get_protected_indices(messages)) == {0, 1, 2, 3} diff --git a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py index 8218c120153..798bea9327c 100644 --- a/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py +++ b/tests/unit/llms/anthropic/chat/guardrail_translation/test_anthropic_guardrail_handler.py @@ -10,6 +10,8 @@ from typing import Any, Literal, Optional, Final from unittest.mock import MagicMock, patch import pytest +import respx +import litellm from litellm.integrations.custom_guardrail import CustomGuardrail @@ -3043,3 +3045,22 @@ class TestAnthropicMessagesHandlerPostCallHookResponse: native = {"type": "message", "role": "assistant", "content": [{"type": "text", "text": "hi"}]} assert AnthropicMessagesHandler().post_call_hook_response(native) is native + + +@pytest.mark.asyncio +async def test_headroom_latest_multipart_turn_preserves_valid_text_blocks( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import HeadroomGuardrail + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + callback: Final = HeadroomGuardrail(api_base="https://headroom.example.com", compress_latest_turn=True) + messages: Final = [{"role": "user", "content": [ + {"type": "text", "text": "Context\n\nparagraph " * 100}, + {"type": "text", "text": "Summarize it"}, + ]}] + result: Final = await AnthropicMessagesHandler().process_input_messages( + {"model": "test-model", "max_tokens": 200, "messages": messages}, callback + ) + assert not respx_mock.calls + assert result["messages"] == messages diff --git a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py index 5d71cbe9afc..6c14d15370e 100644 --- a/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py +++ b/tests/unit/llms/openai/chat/guardrail_translation/test_openai_guardrail_handler.py @@ -7,9 +7,10 @@ with guardrail transformations, including tool calls. import json from collections.abc import Mapping -from typing import Any, Literal, Optional +from typing import Any, Final, Literal, Optional import pytest +import respx from litellm.integrations.custom_guardrail import CustomGuardrail @@ -2365,3 +2366,22 @@ class TestStreamingScanKey: for released in (text_only, finished_tool_call): ended = handler.released_stream_as_ended(released) assert len(ended) == len(released) and all(a is b for a, b in zip(ended, released, strict=True)) + + +@pytest.mark.asyncio +async def test_headroom_latest_user_rewrite_reaches_chat_request( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + import litellm + from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import HeadroomGuardrail + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + callback: Final = HeadroomGuardrail(api_base="https://headroom.example.com", compress_latest_turn=True) + route: Final = respx_mock.post("https://headroom.example.com/v1/compress").respond( + 200, json={"messages": [{"role": "user", "content": "short"}]} + ) + result: Final = await OpenAIChatCompletionsHandler().process_input_messages( + {"messages": [{"role": "user", "content": "long " * 100, "name": "reader"}]}, callback + ) + assert route.called + assert result["messages"] == [{"role": "user", "content": "short", "name": "reader"}] 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..e1836e3dd5f 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 @@ -14,6 +14,8 @@ from unittest.mock import AsyncMock, MagicMock, patch import logging import pytest +import respx +import litellm from fastapi import HTTPException @@ -4174,3 +4176,193 @@ 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_headroom_latest_turn_patches_native_tool_result_with_metadata( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import HeadroomGuardrail + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + callback: Final = HeadroomGuardrail(api_base="https://headroom.example.com", compress_latest_turn=True) + items: Final = [ + {"role": "user", "content": "Summarize this data"}, + {"type": "function_call", "id": "fc_original", "call_id": "call_original", "name": "read", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_original", "output": "data " * 500, "status": "completed"}, + ] + route: Final = respx_mock.post("https://headroom.example.com/v1/compress").respond(200, json={"messages": [ + {"role": "user", "content": "Summarize data"}, + {"role": "tool", "tool_call_id": "call_original", "content": "short data"}, + ]}) + result: Final = await OpenAIResponsesHandler().process_input_messages({"input": items}, callback) + assert route.called + assert result["input"] == [ + {**items[0], "content": "Summarize data"}, items[1], {**items[2], "output": "short data"} + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raw_input", ["long " * 100, [{"role": "user", "content": [ + {"type": "input_text", "text": "long\n\nparagraph " * 100, "source": "first"}, +]}]]) +async def test_headroom_latest_user_input_preserves_native_text_metadata( + raw_input: object, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import HeadroomGuardrail + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + callback: Final = HeadroomGuardrail(api_base="https://headroom.example.com", compress_latest_turn=True) + route: Final = respx_mock.post("https://headroom.example.com/v1/compress").respond( + 200, json={"messages": [{"role": "user", "content": "short\n\nsummary"}]} + ) + result: Final = await OpenAIResponsesHandler().process_input_messages({"input": raw_input}, callback) + assert route.called + assert result["input"] == ([{"role": "user", "content": "short\n\nsummary"}] if isinstance(raw_input, str) else [ + {"role": "user", "content": [ + {"type": "input_text", "text": "short\n\nsummary", "source": "first"}, + ]}, + ]) + + +@pytest.mark.asyncio +async def test_headroom_unsafe_native_mapping_preserves_request_and_tools( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import HeadroomGuardrail + from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket + from litellm.proxy.spend_tracking.compression_savings import extract_compression_saved_tokens + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + callback: Final = HeadroomGuardrail(api_base="https://headroom.example.com", compress_latest_turn=True) + items: Final = [ + {"id": "fc_1", "type": "function_call", "call_id": "call_1", "name": "read", "arguments": "{}"}, + {"role": "assistant", "content": "Reading now"}, + {"type": "function_call_output", "call_id": "call_1", "output": "long " * 100}, + ] + original: Final = {"input": items} + result: Final = await OpenAIResponsesHandler().process_input_messages(original, callback) + assert not respx_mock.calls + assert result["input"] == items + assert "tools" not in result + _, metadata = get_or_create_metadata_bucket(original) + entries: Final = metadata.get("standard_logging_guardrail_information", []) + assert entries == [] + assert extract_compression_saved_tokens({"guardrail_information": entries}) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("skip_source", ["guardrail", "global"]) +@pytest.mark.parametrize("include_system_item", [False, True]) +async def test_headroom_skipped_system_messages_preserve_response_compression( + skip_source: str, include_system_item: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import HeadroomGuardrail + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "skip_system_message_in_guardrail", skip_source == "global") + callback: Final = HeadroomGuardrail( + api_base="https://headroom.example.com", compress_latest_turn=True, + ) + callback.skip_system_message_in_guardrail = True if skip_source == "guardrail" else None + prefix: Final = [{"role": "system", "content": "Keep all identifiers"}] if include_system_item else [] + items: Final = [*prefix, + {"role": "user", "content": "Read old data"}, + {"type": "function_call", "call_id": "old_call", "name": "read", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "old_call", "output": "old " * 500}, + {"role": "assistant", "content": "Done"}, + {"role": "user", "content": "Read new data"}, + {"type": "function_call", "call_id": "new_call", "name": "read", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "new_call", "output": "new " * 500}, + ] + route: Final = respx_mock.post("https://headroom.example.com/v1/compress").respond(200, json={"messages": [ + {"role": "user", "content": "Read old"}, + {"role": "tool", "tool_call_id": "old_call", "content": "old summary"}, + {"role": "user", "content": "Read new"}, + {"role": "tool", "tool_call_id": "new_call", "content": "new summary"}, + ]}) + result: Final = await OpenAIResponsesHandler().process_input_messages( + {"instructions": "Be terse", "input": items}, callback + ) + assert route.called + assert result["instructions"] == "Be terse" + assert result["input"] == [*prefix, + {"role": "user", "content": "Read old"}, items[len(prefix) + 1], + {"type": "function_call_output", "call_id": "old_call", "output": "old summary"}, + {"role": "assistant", "content": "Done"}, + {"role": "user", "content": "Read new"}, items[len(prefix) + 5], + {"type": "function_call_output", "call_id": "new_call", "output": "new summary"}, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("output", [ + [{"type": "input_text", "text": "First\n\nparagraph", "source": "one"}, + {"type": "input_text", "text": "Second", "source": "two"}], + [{"type": "input_text", "text": "Image context"}, + {"type": "input_image", "image_url": "https://example.com/image.png"}], +]) +async def test_headroom_native_tool_parts_remain_intact_while_latest_output_compresses( + output: list[dict[str, object]], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import HeadroomGuardrail + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + callback: Final = HeadroomGuardrail(api_base="https://headroom.example.com", compress_latest_turn=True) + items: Final = [ + {"type": "function_call", "call_id": "old", "name": "read", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "old", "output": output}, + {"role": "user", "content": "Read new data"}, + {"type": "function_call", "call_id": "new", "name": "read", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "new", "output": "new " * 500}, + ] + route: Final = respx_mock.post("https://headroom.example.com/v1/compress").respond(200, json={"messages": [ + {"role": "user", "content": "Read new data"}, + {"role": "tool", "tool_call_id": "new", "content": "new summary"}, + ]}) + result: Final = await OpenAIResponsesHandler().process_input_messages({"input": items}, callback) + assert json.loads(route.calls.last.request.content)["messages"] == [ + {"role": "user", "content": "Read new data"}, + {"role": "tool", "tool_call_id": "new", "content": "new " * 500}, + ] + assert result["input"] == [*items[:-1], {**items[-1], "output": "new summary"}] + + +@pytest.mark.parametrize("raw_input", [None, 42, ["invalid item"], [{1: "invalid key"}]]) +def test_native_compression_excludes_invalid_input_shapes(raw_input: object) -> None: + from litellm.llms.openai.responses.guardrail_translation.handler import rewriteable_response_message_indices + + assert rewriteable_response_message_indices(raw_input, None, []) == frozenset() + + +@pytest.mark.asyncio +async def test_headroom_cached_tool_definition_preserves_latest_output_compression( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import HeadroomGuardrail + from litellm.responses.litellm_completion_transformation.transformation import TOOL_CALLS_CACHE + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + call_id: Final = "call_headroom_cached_definition" + TOOL_CALLS_CACHE.set_cache(key=call_id, value={ + "id": call_id, "type": "function", "function": {"name": "read", "arguments": "{}"}, + }) + callback: Final = HeadroomGuardrail(api_base="https://headroom.example.com", compress_latest_turn=True) + items: Final = [ + {"role": "user", "content": "Read data"}, + {"type": "function_call_output", "call_id": call_id, "output": "long " * 500}, + ] + route: Final = respx_mock.post("https://headroom.example.com/v1/compress").respond(200, json={"messages": [ + {"role": "user", "content": "Read data"}, + {"role": "tool", "tool_call_id": call_id, "content": "short summary"}, + ]}) + try: + result: Final = await OpenAIResponsesHandler().process_input_messages({"input": items}, callback) + assert route.called + assert json.loads(route.calls.last.request.content)["messages"] == [ + {"role": "user", "content": "Read data"}, + {"role": "tool", "tool_call_id": call_id, "content": "long " * 500}, + ] + assert result["input"] == [items[0], {**items[1], "output": "short summary"}] + finally: + TOOL_CALLS_CACHE.delete_cache(key=call_id) diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py index 5bdcdddf4b7..a6a612dc50e 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -3075,3 +3075,246 @@ def test_unusable_timeout_falls_back_to_the_default(configured: float): assert guardrail.timeout.read == 60.0 assert guardrail.timeout.connect == 5.0 + + +@pytest.mark.asyncio +async def test_configured_latest_turn_compression_preserves_tool_linkage( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.proxy.guardrails.guardrail_hooks.headroom import initialize_guardrail + from litellm.types.guardrails import LitellmParams + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + params: Final = LitellmParams.model_validate({ + "guardrail": "headroom", "mode": "pre_call", "api_base": FAKE_API_BASE, + "compress_latest_turn": True, + }) + callback: Final = initialize_guardrail(params, {"guardrail_name": "compress-latest", "litellm_params": params}) + messages: Final = [ + {"role": "user", "content": "Summarize this listing " * 100}, + {"role": "assistant", "content": None, "tool_calls": [ + {"id": "call_listing", "type": "function", "function": {"name": "list_files", "arguments": "{}"}}, + ]}, + {"role": "tool", "tool_call_id": "call_listing", "content": "file.py\n" * 1000, "name": "list_files"}, + ] + compressed: Final = [ + {"role": "user", "content": "Summarize the listing"}, + {"role": "tool", "tool_call_id": "call_listing", "content": "1000 Python files"}, + ] + route: Final = respx_mock.post(FAKE_API_BASE + "/v1/compress").respond(200, json={"messages": compressed}) + try: + result: Final = await callback.apply_guardrail( + GenericGuardrailAPIInputs(structured_messages=messages), {}, "request" + ) + assert result["structured_messages"] == [ + {**messages[0], "content": compressed[0]["content"]}, messages[1], + {**messages[2], "content": compressed[1]["content"]}, + ] + assert json.loads(route.calls.last.request.content)["messages"] == [messages[0], messages[2]] + finally: + litellm.logging_callback_manager.remove_callback_from_list_by_object( + litellm.callbacks, callback, require_self=False + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("rewritten", [ + [], + [{"role": "assistant", "content": "short"}], + [{"role": "user", "content": {"unexpected": "shape"}}], + [{"role": "user", "content": ""}], + [{"role": "user", "content": " "}], + [{"role": "user", "content": "short", "tool_call_id": "wrong"}], +]) +async def test_latest_turn_unsafe_rewrite_keeps_original_input( + rewritten: list[dict[str, object]], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + callback: Final = HeadroomGuardrail(api_base=FAKE_API_BASE, compress_latest_turn=True) + original: Final = GenericGuardrailAPIInputs(structured_messages=[{"role": "user", "content": "long " * 100}]) + route: Final = respx_mock.post(FAKE_API_BASE + "/v1/compress").respond(200, json={"messages": rewritten, "ccr_hashes": [CCR_HASH]}) + request_data: Final[dict[str, object]] = {} + result: Final = await callback.apply_guardrail(original, request_data, "request") + assert route.called + assert result is original + assert "tools" not in result + entries: Final = _recorded_guardrail_entries(request_data) + assert len(entries) == 1 + assert entries[0]["guardrail_status"] == "guardrail_failed_to_respond" + assert entries[0]["duration"] >= 0 + assert entries[0]["end_time"] >= entries[0]["start_time"] + assert extract_compression_saved_tokens({"guardrail_information": entries}) == 0 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("enabled", [False, True]) +async def test_latest_turn_preserves_cached_prefix_and_retrieval( + enabled: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + callback: Final = HeadroomGuardrail(api_base=FAKE_API_BASE, compress_latest_turn=enabled) + messages: Final = [ + {"role": "system", "content": "System instruction"}, + {"role": "user", "content": "Cached data", "cache_control": {"type": "ephemeral"}}, + {"role": "assistant", "content": None, "tool_calls": [ + {"id": "retrieve", "type": "function", "function": {"name": "headroom_retrieve", "arguments": "{}"}}, + ]}, + {"role": "tool", "tool_call_id": "retrieve", "content": "Recovered original"}, + {"role": "user", "content": "Latest instruction"}, + ] + with respx.mock(assert_all_called=False) as router: + route: Final = router.post(FAKE_API_BASE + "/v1/compress").respond( + 200, json={"messages": [{"role": "user", "content": "Short instruction"}]} + ) + result: Final = await callback.apply_guardrail( + GenericGuardrailAPIInputs(structured_messages=messages), {}, "request" + ) + assert route.called is enabled + assert result["structured_messages"] == [ + *messages[:-1], {**messages[-1], "content": "Short instruction"} if enabled else messages[-1], + ] + + +@pytest.mark.asyncio +async def test_latest_turn_retains_single_text_block_metadata( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + callback: Final = HeadroomGuardrail(api_base=FAKE_API_BASE, compress_latest_turn=True) + messages: Final = [{"role": "user", "content": [ + {"type": "text", "text": "long\n\nparagraph " * 100, "source": "first"}, + ]}] + route: Final = respx_mock.post(FAKE_API_BASE + "/v1/compress").respond( + 200, json={"messages": [{"role": "user", "content": "short\n\nsummary"}]} + ) + result: Final = await callback.apply_guardrail( + GenericGuardrailAPIInputs(structured_messages=messages), {}, "request" + ) + assert route.called + assert result["structured_messages"] == [{"role": "user", "content": [ + {"type": "text", "text": "short\n\nsummary", "source": "first"}, + ]}] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "content", [[], ["invalid"], [{"type": "image_url", "image_url": {"url": "https://example.com/image"}}]] +) +async def test_latest_turn_does_not_replace_non_text_content( + content: list[object], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + callback: Final = HeadroomGuardrail(api_base=FAKE_API_BASE, compress_latest_turn=True) + original: Final = GenericGuardrailAPIInputs(structured_messages=[{"role": "user", "content": content}]) + result: Final = await callback.apply_guardrail(original, {}, "request") + assert not respx_mock.calls + assert result is original + + +@pytest.mark.asyncio +@pytest.mark.parametrize("content", ["Already short", [ + {"type": "text", "text": "Already", "source": "first"}, + {"type": "text", "text": "short", "source": "second"}, +]]) +async def test_latest_turn_unchanged_text_retains_original_metadata( + content: object, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + callback: Final = HeadroomGuardrail(api_base=FAKE_API_BASE, compress_latest_turn=True) + original: Final = [{"role": "user", "content": content, "name": "reader"}] + if isinstance(content, str): + respx_mock.post(FAKE_API_BASE + "/v1/compress").respond( + 200, json={"messages": [ + {"role": "user", "content": "Already short" if isinstance(content, str) else "Already\n\nshort"} + ]} + ) + result: Final = await callback.apply_guardrail( + GenericGuardrailAPIInputs(structured_messages=original), {}, "request" + ) + assert bool(respx_mock.calls) is isinstance(content, str) + assert result["structured_messages"] == original + + +@pytest.mark.asyncio +@pytest.mark.parametrize("surface", ["chat", "messages", "responses"]) +async def test_ambiguous_history_does_not_prevent_latest_tool_compression( + surface: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler + from litellm.llms.openai.chat.guardrail_translation.handler import OpenAIChatCompletionsHandler + from litellm.llms.openai.responses.guardrail_translation.handler import OpenAIResponsesHandler + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + callback: Final = HeadroomGuardrail(api_base=FAKE_API_BASE, compress_latest_turn=True) + part_type: Final = "input_text" if surface == "responses" else "text" + history: Final = {"role": "user", "content": [ + {"type": part_type, "text": "First paragraph\n\nSecond paragraph", **({} if surface == "messages" else {"source": "first"})}, + {"type": part_type, "text": "Other block", **({} if surface == "messages" else {"source": "second"})}, + ]} + tool_text: Final = "file.py\n" * 1000 + prefix: Final = [history, {"role": "assistant", "content": "Read the history"}] + chat_rows: Final = [*prefix, {"role": "user", "content": "List files"}, + {"role": "assistant", "content": None, "tool_calls": [ + {"id": "call_files", "type": "function", "function": {"name": "list_files", "arguments": "{}"}}, + ]}, {"role": "tool", "tool_call_id": "call_files", "content": tool_text}, + ] + native_items: Final = [*prefix, {"role": "user", "content": "List files"}, + {"type": "function_call", "call_id": "call_files", "name": "list_files", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call_files", "output": tool_text}, + ] + anthropic_rows: Final = [*prefix, {"role": "user", "content": "List files"}, + {"role": "assistant", "content": [{"type": "tool_use", "id": "call_files", "name": "list_files", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_files", "content": tool_text}]}, + ] + request: Final = {"input": native_items} if surface == "responses" else { + "messages": anthropic_rows if surface == "messages" else chat_rows + } + route: Final = respx_mock.post(FAKE_API_BASE + "/v1/compress").respond(200, json={ + "messages": [{"role": "user", "content": "List files"}, + {"role": "tool", "tool_call_id": "call_files", "content": "file summary"}], + "tokens_before": 1000, "tokens_after": 10, "ccr_hashes": [CCR_HASH], + }) + handler: Final = (OpenAIResponsesHandler() if surface == "responses" else + AnthropicMessagesHandler() if surface == "messages" else OpenAIChatCompletionsHandler()) + result: Final = await handler.process_input_messages({**request, "model": "test-model", "max_tokens": 100}, callback) + assert json.loads(route.calls.last.request.content)["messages"] == [ + {"role": "user", "content": "List files"}, + {"role": "tool", "tool_call_id": "call_files", "content": tool_text}, + ] + if surface == "responses": + assert result["input"] == [*native_items[:-1], {**native_items[-1], "output": "file summary"}] + elif surface == "chat": + assert result["messages"] == [*chat_rows[:-1], {**chat_rows[-1], "content": "file summary"}] + else: + assert result["messages"][0] == history + assert result["messages"][-1]["content"] == [ + {"type": "tool_result", "tool_use_id": "call_files", "content": "file summary"} + ] + assert extract_compression_saved_tokens({"guardrail_information": _recorded_guardrail_entries(result)}) == 990 + + assert has_headroom_retrieve_tool(result.get("tools")) + + +@pytest.mark.asyncio +async def test_native_scope_mismatch_discards_rewrite_and_retrieval( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + callback: Final = HeadroomGuardrail(api_base=FAKE_API_BASE, compress_latest_turn=True) + original: Final = GenericGuardrailAPIInputs(structured_messages=[{"role": "user", "content": "long " * 100}]) + request: Final = {"input": [ + {"role": "user", "content": "long " * 100}, {"role": "user", "content": "Additional input"}, + ]} + route: Final = respx_mock.post(FAKE_API_BASE + "/v1/compress").respond(200, json={ + "messages": [{"role": "user", "content": "short"}], + "ccr_hashes": [CCR_HASH], "tokens_before": 1000, "tokens_after": 10, + }) + result: Final = await callback.apply_guardrail(original, request, "request") + assert route.called + assert result is original + assert "tools" not in result + entries: Final = _recorded_guardrail_entries(request) + assert len(entries) == 1 + assert entries[0]["guardrail_status"] == "guardrail_failed_to_respond" + assert entries[0]["duration"] >= 0 + assert extract_compression_saved_tokens({"guardrail_information": entries}) == 0 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 2b174a1141a..e2ba1dbc8a5 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -36572,6 +36572,12 @@ export interface components { * @description Client secret of the gateway's Entra app registration, used to perform the On-Behalf-Of exchange. Falls back to the AGENT365_CLIENT_SECRET environment variable. */ client_secret?: string | null; + /** + * Compress Latest Turn + * @description Allow compression of latest user and tool-result text, which can lose instructions or details. System prompts, assistant messages, cached prefixes and retrieved content remain protected. Unsafe rewrites retain the original input. + * @default false + */ + compress_latest_turn: boolean; /** * Confidence Threshold * @description Only block or mask when detection confidence >= this value; below threshold, allow or log_only.