diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 8bfd5cca58a..18797b0a880 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -5,6 +5,7 @@ import re import time import uuid from collections.abc import Mapping, Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, TypeGuard import httpx @@ -46,7 +47,10 @@ if TYPE_CHECKING: BYPASS_HEADER: Final = "x-headroom-bypass" HEADROOM_RETRIEVE_TOOL_NAME: Final = "headroom_retrieve" -_HASH_PATTERN: Final = re.compile(r"hash=([a-f0-9]{24})") +# Headroom writes a retrieval hash as either `<>` or +# `Retrieve more: hash=HASH`, and its own reader accepts 12 to 24 hex chars +# (the row-drop path emits SHA-256[:12]). +_HASH_PATTERN: Final = re.compile(r"(?:< TypeGuard[list[object]]: # guard-ok: isin return isinstance(value, list) +@dataclass(frozen=True, slots=True) +class _CompressResult: + messages: list[dict[str, object]] + succeeded: bool + stats: dict[str, object] + ccr_hashes: tuple[str, ...] = () + + def _flatten_messages_for_compression(messages: list[dict[str, object]]) -> list[dict[str, object]]: """Collapse all-text list-of-parts content to plain strings for /v1/compress. @@ -169,6 +181,17 @@ def _build_compress_failure_detail(status_code: int, body: str) -> dict[str, obj return {"status_code": status_code, "body": body} +def _read_ccr_hashes(body: Mapping[str, object]) -> tuple[str, ...]: + """Hashes the compression service says it stored, from the documented + ``ccr_hashes`` response field. Scanning the returned text is a fallback: + the service knows exactly which markers it wrote, and marker wording is + its own to change.""" + raw: Final = body.get("ccr_hashes") + if not _is_object_list(raw): + return () + return tuple(item for item in raw if isinstance(item, str) and item) + + def extract_hashes_from_messages(messages: list[dict[str, object]]) -> list[str]: hashes: Final[list[str]] = [] for msg in messages: @@ -191,14 +214,15 @@ def _build_headroom_retrieve_tool() -> dict[str, object]: "name": HEADROOM_RETRIEVE_TOOL_NAME, "description": ( "Retrieve original content that was compressed by Headroom. " - "Call this when you encounter a compression marker containing a hash." + "Call this when you encounter a compression marker such as " + "`<>` or `Retrieve more: hash=HASH`." ), "parameters": { "type": "object", "properties": { "hash": { "type": "string", - "description": "The 24-character hex hash from the compression marker.", + "description": "The hex hash from the compression marker.", }, "query": { "type": "string", @@ -422,7 +446,7 @@ class HeadroomGuardrail(CustomGuardrail): self, messages: list[dict[str, object]], model: str | None, - ) -> tuple[list[dict[str, object]], bool, dict[str, object]]: + ) -> _CompressResult: payload: Final[dict[str, object]] = {"messages": messages} if model: payload["model"] = model @@ -434,7 +458,7 @@ class HeadroomGuardrail(CustomGuardrail): headers=self._request_headers(), ) except httpx.HTTPStatusError as e: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned an error", @@ -444,7 +468,7 @@ class HeadroomGuardrail(CustomGuardrail): {}, ) except (httpx.ConnectError, httpx.TimeoutException, httpx.TransportError, litellm.Timeout) as e: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service unreachable", @@ -454,7 +478,7 @@ class HeadroomGuardrail(CustomGuardrail): {}, ) if raw_response is None: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned no response", @@ -466,7 +490,7 @@ class HeadroomGuardrail(CustomGuardrail): response: Final[HttpxResponse] = raw_response if response.status_code != 200: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned an error", @@ -479,7 +503,7 @@ class HeadroomGuardrail(CustomGuardrail): try: body: Final[object] = response.json() except ValueError: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned non-JSON response", @@ -489,7 +513,7 @@ class HeadroomGuardrail(CustomGuardrail): {}, ) if not _is_str_object_dict(body): - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned unexpected response shape", @@ -501,7 +525,7 @@ class HeadroomGuardrail(CustomGuardrail): compressed_messages: Final = body.get("messages") if not _is_object_list(compressed_messages): - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service response missing 'messages'", @@ -513,7 +537,7 @@ class HeadroomGuardrail(CustomGuardrail): filtered: Final = [item for item in compressed_messages if _is_str_object_dict(item)] if not filtered: - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service returned empty message list", @@ -526,7 +550,7 @@ class HeadroomGuardrail(CustomGuardrail): if len(filtered) != len(messages): # Rows are matched positionally when the never-compressed messages # are put back, so a reshaped conversation cannot be applied at all. - return ( + return _CompressResult( self._handle_compress_failure( messages, "Headroom compression service changed the message count", @@ -567,7 +591,7 @@ class HeadroomGuardrail(CustomGuardrail): # tokens_saved, which the live compression service omits; derive it # so savings are counted, but let a service-sent value win. stats["tokens_saved"] = tokens_before - tokens_after - return filtered, True, stats + return _CompressResult(filtered, True, stats, _read_ccr_hashes(body)) async def _call_retrieve(self, hash_value: str, query: str | None = None) -> str: params: Final[dict[str, str]] = {} @@ -646,7 +670,7 @@ class HeadroomGuardrail(CustomGuardrail): model: Final = self.headroom_model or request_data.get("model") start_time: Final = time.time() - returned, compression_succeeded, stats = await self._call_compress( + result: Final = await self._call_compress( messages=_flatten_messages_for_compression(compressible), model=model if isinstance(model, str) else None, ) @@ -656,7 +680,7 @@ class HeadroomGuardrail(CustomGuardrail): add_guardrail_to_applied_guardrails_header, ) - if not compression_succeeded: + if not result.succeeded: self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response={"error": "headroom compression unavailable; request forwarded uncompressed"}, request_data=request_data, @@ -675,12 +699,12 @@ class HeadroomGuardrail(CustomGuardrail): compressed: Final = _restore_protected_messages( messages=messages, - compressed=_restore_content_shapes(originals=compressible, returned=returned), + compressed=_restore_content_shapes(originals=compressible, returned=result.messages), protected_indices=protected_indices, ) self.add_standard_logging_guardrail_information_to_request_data( - guardrail_json_response=stats, + guardrail_json_response=result.stats, request_data=request_data, guardrail_status="success", guardrail_provider=HEADROOM_GUARDRAIL_PROVIDER, @@ -690,7 +714,7 @@ class HeadroomGuardrail(CustomGuardrail): ) add_guardrail_to_applied_guardrails_header(request_data=request_data, guardrail_name=self.guardrail_name) - hashes: Final = extract_hashes_from_messages(compressed) + hashes: Final = frozenset(result.ccr_hashes) | frozenset(extract_hashes_from_messages(compressed)) if not hashes: return {**inputs, "structured_messages": compressed} # pyright: ignore[reportReturnType] @@ -699,7 +723,7 @@ class HeadroomGuardrail(CustomGuardrail): if not call_id: call_id = str(uuid.uuid4()) request_data["litellm_call_id"] = call_id - self._issued_hashes_by_call_id[call_id] = (frozenset(hashes), time.monotonic() + _HASH_CACHE_TTL_SECONDS) + self._issued_hashes_by_call_id[call_id] = (hashes, time.monotonic() + _HASH_CACHE_TTL_SECONDS) existing_tools: Final = inputs.get("tools") retrieve_tool: Final = _build_headroom_retrieve_tool() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 7a2772ce78c..3f2df420540 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -85,7 +85,7 @@ def _make_guardrail(**kwargs) -> HeadroomGuardrail: return HeadroomGuardrail(**defaults) -def _make_compress_response(messages: list, status: int = 200) -> MagicMock: +def _make_compress_response(messages: list, status: int = 200, ccr_hashes: list | None = None) -> MagicMock: mock = MagicMock() mock.status_code = status mock.json.return_value = { @@ -94,6 +94,7 @@ def _make_compress_response(messages: list, status: int = 200) -> MagicMock: "tokens_after": 100, "compression_ratio": 0.1, "transforms_applied": ["router:smart_crusher:0.35"], + "ccr_hashes": ccr_hashes if ccr_hashes is not None else [], } mock.text = "" return mock @@ -319,6 +320,105 @@ async def test_apply_guardrail_injects_retrieve_tool_when_hashes_present( assert has_headroom_retrieve_tool(tools) +@pytest.mark.asyncio +async def test_apply_guardrail_injects_retrieve_tool_for_ccr_marker( + guardrail: HeadroomGuardrail, +): + """`<>` is what Headroom writes in CCR mode. Missing it + left the marker in the prompt with no tool to redeem it.""" + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + mock_response = _make_compress_response([{"role": "user", "content": "<>"}]) + request_data = {"model": "gpt-4o", "litellm_call_id": "call-ccr"} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert has_headroom_retrieve_tool(result.get("tools")) + assert "f3b3d2ef3049" in guardrail._issued_hashes_by_call_id["call-ccr"][0] + + +@pytest.mark.asyncio +async def test_apply_guardrail_trusts_ccr_hashes_from_compress_response( + guardrail: HeadroomGuardrail, +): + """The service reports the hashes it stored in `ccr_hashes`, so retrieval + keeps working when it words a marker in a way this guardrail cannot parse.""" + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + mock_response = _make_compress_response( + [{"role": "user", "content": "[22.8KB elided, ask for {f3b3d2ef3049}]"}], + ccr_hashes=["f3b3d2ef3049"], + ) + request_data = {"model": "gpt-4o", "litellm_call_id": "call-ccr"} + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data=request_data, + input_type="request", + ) + + assert has_headroom_retrieve_tool(result.get("tools")) + + original_content = "the original tool output" + with patch.object( + guardrail.async_handler, + "get", + new_callable=AsyncMock, + return_value=_make_retrieve_response(original_content), + ) as mock_get: + plan = await guardrail.async_build_agentic_loop_plan( + tools={ + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "name": HEADROOM_RETRIEVE_TOOL_NAME, + "arguments": {"hash": "f3b3d2ef3049"}, + } + ] + }, + model="gpt-4o", + messages=[{"role": "user", "content": "what did it say?"}], + response=_make_openai_response_with_tool_call( + tool_name=HEADROOM_RETRIEVE_TOOL_NAME, + arguments={"hash": "f3b3d2ef3049"}, + tool_id="call_1", + ), + anthropic_messages_provider_config=None, + anthropic_messages_optional_request_params={}, + logging_obj=None, + stream=False, + kwargs={"litellm_call_id": "call-ccr"}, + ) + + mock_get.assert_called_once() + assert plan.request_patch is not None + follow_up = plan.request_patch.messages + assert follow_up is not None + tool_result = next(m for m in follow_up if m.get("role") == "tool") + assert tool_result["content"] == original_content + + @pytest.mark.asyncio async def test_apply_guardrail_no_tool_injected_when_no_hashes( guardrail: HeadroomGuardrail, @@ -818,6 +918,25 @@ def test_extract_hashes_from_messages_finds_hashes(): assert "aabbccdd001122334455aabb" in hashes +def test_extract_hashes_from_messages_finds_ccr_markers(): + """Headroom's own marker form. Its row-drop path emits SHA-256[:12], so a + 12 hex hash counts just as much as the 24 hex one.""" + messages = [ + {"role": "tool", "content": "<>"}, + {"role": "user", "content": "<>"}, + ] + hashes = extract_hashes_from_messages(messages) + assert "f3b3d2ef3049" in hashes + assert "b573993006976af767214fac" in hashes + + +def test_extract_hashes_from_messages_ignores_overlong_hashes(): + """A longer hex run is some other identifier, not a Headroom hash. Matching + its first 24 chars would issue a hash that retrieval can never resolve.""" + messages = [{"role": "user", "content": "hash=" + "a" * 32}] + assert not extract_hashes_from_messages(messages) + + def test_extract_hashes_from_messages_ignores_short_hashes(): messages = [{"role": "user", "content": "hash=tooshort"}] hashes = extract_hashes_from_messages(messages)