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..fa3cdf70a03 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -266,13 +266,23 @@ 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"), ignore_empty=True) + ) + 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] @@ -493,12 +503,22 @@ 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, ) -> _RequestFields | None: + if isinstance(raw_input, str) and preserve_text_metadata: + return patch_response_request_fields( + [{"role": "user", "content": raw_input}], + instructions, + original_messages, + structured_messages, + preserve_text_metadata=True, + ) if not isinstance(raw_input, list) or len(original_messages) != len(structured_messages): return None offset: Final = 1 if instructions else 0 @@ -528,6 +548,7 @@ def _patched_request_fields( item_for_message[message_index]: _rewritten_input_item( cast("Mapping[str, object]", raw_input[item_for_message[message_index]]), # cast-ok: checked Mappings rewritten, + preserve_text_metadata=preserve_text_metadata, ) for message_index, rewritten in body_changes } @@ -545,11 +566,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 +787,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 7198d74909c..f34ede36471 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -16335,6 +16335,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..2717840d10d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/content_text.py +++ b/litellm/proxy/guardrails/guardrail_hooks/content_text.py @@ -12,13 +12,15 @@ 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 -def content_to_text(content: object) -> str: +def content_to_text(content: object, *, ignore_empty: bool = False) -> str: """Collapse a message ``content`` (str or list-of-parts) to plain text. For the multimodal list shape, joins ``{type: "text", text: ...}`` parts @@ -31,7 +33,7 @@ def content_to_text(content: object) -> str: for part in content: if isinstance(part, dict) and part.get("type") == "text": text = part.get("text") - if isinstance(text, str): + if isinstance(text, str) and (text or not ignore_empty): parts.append(text) return "\n\n".join(parts) return "" @@ -94,3 +96,23 @@ 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 replace_text_content(content: object, text: str) -> object: + if isinstance(content, str): + return text + if not isinstance(content, list) or not content: + return None + try: + parts: Final = _TEXT_PARTS_ADAPTER.validate_python(content) + except ValidationError: + return None + if any( + part.get("type") not in ("text", "input_text", "output_text") or not isinstance(part.get("text"), str) + for part in parts + ): + return None + return [{**part, "text": text if index == 0 else ""} for index, part in enumerate(parts)] 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..c66578747fc 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -37,6 +37,7 @@ from litellm.proxy.guardrails.guardrail_hooks.content_text import ( content_to_text, is_all_text_parts, 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 +142,31 @@ def _restore_content_shapes( return restored +def _compressed_message_content(original: Mapping[str, object], rewritten: Mapping[str, object]) -> object: + content: Final = rewritten.get("content") + if content == original.get("content"): + return 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 +265,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 +286,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( @@ -273,10 +302,10 @@ def _protected_indices( 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 +534,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 +548,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 +561,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 +711,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( @@ -800,20 +837,12 @@ 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") + 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 + ) compressible: Final = [m for i, m in enumerate(messages) if i not in protected_indices] if not compressible: return inputs @@ -847,12 +876,34 @@ 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: + 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 self.compress_latest_turn and "input" in request_data: + from litellm.llms.openai.responses.guardrail_translation.handler import patch_response_request_fields + + if ( + patch_response_request_fields( + typed_request.get("input"), + typed_request.get("instructions"), + messages, + compressed, + preserve_text_metadata=True, + ) + is None + ): + 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/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..444b09600d9 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,74 @@ 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 " * 100, "source": "first"}, + {"type": "input_text", "text": "long " * 100, "source": "second"}, +]}]]) +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"}]} + ) + result: Final = await OpenAIResponsesHandler().process_input_messages({"input": raw_input}, callback) + assert route.called + assert result["input"] == ([{"role": "user", "content": "short"}] if isinstance(raw_input, str) else [ + {"role": "user", "content": [ + {"type": "input_text", "text": "short", "source": "first"}, + {"type": "input_text", "text": "", "source": "second"}, + ]}, + ]) + + +@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 + + 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} + route: Final = respx_mock.post("https://headroom.example.com/v1/compress").respond(200, json={ + "messages": [{"role": "tool", "tool_call_id": "call_1", "content": "short"}], + "ccr_hashes": ["abcdefabcdef"], + }) + result: Final = await OpenAIResponsesHandler().process_input_messages(original, callback) + assert route.called + assert result == {"input": items} + assert "tools" not in result diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py index 5bdcdddf4b7..0f57371e841 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -3075,3 +3075,139 @@ 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": "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}) + result: Final = await callback.apply_guardrail(original, {}, "request") + assert route.called + assert result is original + + +@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_each_text_parts_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 " * 100, "source": "first"}, + {"type": "text", "text": "more " * 100, "source": "second"}, + ]}] + route: Final = respx_mock.post(FAKE_API_BASE + "/v1/compress").respond( + 200, json={"messages": [{"role": "user", "content": "short"}]} + ) + 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", "source": "first"}, + {"type": "text", "text": "", "source": "second"}, + ]}] + + +@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}]) + route: Final = respx_mock.post(FAKE_API_BASE + "/v1/compress").respond( + 200, json={"messages": [{"role": "user", "content": "short"}]} + ) + result: Final = await callback.apply_guardrail(original, {}, "request") + assert route.called + assert result is original + + +@pytest.mark.asyncio +async def test_latest_turn_unchanged_text_retains_original_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) + original: Final = [{"role": "user", "content": "Already short", "name": "reader"}] + route: Final = respx_mock.post(FAKE_API_BASE + "/v1/compress").respond( + 200, json={"messages": [{"role": "user", "content": "Already short"}]} + ) + result: Final = await callback.apply_guardrail(GenericGuardrailAPIInputs(structured_messages=original), {}, "request") + assert route.called + assert result["structured_messages"] == original diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 733a093bcac..530b5bf4d11 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -36419,6 +36419,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.