scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

fix(guardrails/headroom): recognise <<ccr:...>> markers and ccr_hashes so compressed content is retrievable

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-08-11 23:48:45 +00:00
parent bea31871fc
commit 0cb5a34b98
2 changed files with 164 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

64
litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py
View file

@ -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 `<<ccr:HASH,type,size>>` 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"(?:<<ccr:|hash=)([a-f0-9]{12,24})(?![a-f0-9])")
_HASH_CACHE_TTL_SECONDS: Final = 15 * 60
@ -58,6 +62,14 @@ def _is_object_list(value: object) -> 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 "
"`<<ccr:HASH,type,size>>` 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()

121
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py
View file

@ -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,
):
"""`<<ccr:HASH,type,size>>` 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": "<<ccr:f3b3d2ef3049,string,22.8KB>>"}])
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": "<<ccr:f3b3d2ef3049,string,22.8KB>>"},
{"role": "user", "content": "<<ccr:b573993006976af767214fac>>"},
]
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)