feat(guardrails): add headroom guardrail for message compression - #31407
Conversation
Adds a headroom guardrail that compresses request messages via POST /v1/compress before they reach the LLM. The guardrail implements apply_guardrail so it runs on the unified guardrail path; it receives pre-built structured_messages (OpenAI format) from the translation layer, calls the headroom compression service, and returns the compressed messages as structured_messages. Set x-headroom-bypass: true on the request to skip compression. Also adds structured_messages write-back support to the OpenAI and Anthropic translation handlers: when apply_guardrail returns structured_messages, those are written to data["messages"] directly (OpenAI) or reverse-translated via anthropic_messages_pt (Anthropic) instead of falling through to the existing text-patch path. This is a prerequisite for any guardrail that needs to replace the full message list rather than patch individual text spans.
…guardrail_information in spend logs
|
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
Greptile SummaryThis PR adds a Headroom guardrail for compressing request messages before model calls. The main changes are:
Confidence Score: 5/5The changes are isolated to the new Headroom guardrail integration, guardrail translation write-back behavior, and metadata propagation for spend logging. The implementation is covered by focused mocked tests for compression behavior, bypass handling, service error handling, structured message write-back, and logging metadata synchronization.
What T-Rex did
Reviews (9): Last reviewed commit: "fix(lint): extract _write_back_structure..." | Re-trigger Greptile |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 45340607f6
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| data["messages"] = anthropic_messages_pt( | ||
| messages=guardrailed_structured_messages, | ||
| model=model, | ||
| llm_provider="anthropic", | ||
| ) |
There was a problem hiding this comment.
Handle Anthropic system messages before converting back
When an Anthropic request has the normal top-level system field, get_structured_messages() translates it into a leading OpenAI role: system message before Headroom sees it. Passing the returned list directly into anthropic_messages_pt() here stalls such requests because that helper only advances over user/tool/function/assistant roles, so a leading system message leaves its loop index unchanged. With the default Headroom config where system messages are not skipped, any Anthropic request with a system prompt can hang instead of reaching the model; the system block needs to be stripped or restored to data["system"] separately before this conversion.
Useful? React with 👍 / 👎.
| guardrailed_structured_messages = guardrailed_inputs.get("structured_messages") | ||
| if guardrailed_structured_messages is not None: | ||
| data["messages"] = guardrailed_structured_messages |
There was a problem hiding this comment.
Preserve skipped OpenAI messages on structured writeback
If this guardrail is configured with skip_system_message_in_guardrail or skip_tool_message_in_guardrail, the structured_messages sent to Headroom above has those roles filtered out, but this new writeback treats the returned subset as the full conversation. For an OpenAI chat request with a system prompt and Headroom skip_system_message_in_guardrail: true, the system message is deleted from data["messages"], whereas the old text patchback path preserved skipped messages. Merge the compressed messages back into the original conversation or avoid full replacement when skip filters were applied.
Useful? React with 👍 / 👎.
|
@greptile review |
|
@greptile review |
|
@greptile review |
|
@greptile review |
|
@greptileai review |
|
@greptile review |
…sages_pt reverse-translation
…tadata for spend log
…r not model_call_details copy
| ) | ||
|
|
||
| if litellm_logging_obj is not None: | ||
| slg_info = ( |
There was a problem hiding this comment.
this is a really complicated if block. What exactly are you trying to do here? is there a simpler way of achieving this?
| "standard_logging_guardrail_information" | ||
| ) | ||
| if existing is None: | ||
| log_meta["standard_logging_guardrail_information"] = ( |
There was a problem hiding this comment.
this is so nested. can you simplify what within the logging object you're trying to extract ?
| isinstance(block, dict) | ||
| and block.get("type") == "thinking" | ||
| ): | ||
| block.pop("cache_control", None) |
There was a problem hiding this comment.
what is the problem you're trying to solve here? please make sure this is a tested flow
krrish-berri-2
left a comment
There was a problem hiding this comment.
Summary
Headroom guardrail + structured_messages write-back look good and should ship. The guardrail logging sync block is the main blocker: it fixes a real platform bug at the wrong layer and needs to move before merge.
Required changes (in order)
1. Remove the inline sync block from the anthropic translation handler
Delete the ~35-line block after apply_guardrail that copies standard_logging_guardrail_information into logging_obj.litellm_params / model_call_details. This is a cross-cutting concern, not anthropic-handler logic.
2. Add a shared helper and call it from @log_guardrail_information
Add _sync_guardrail_info_to_logging_obj(request_data, logging_obj) in litellm/integrations/custom_guardrail.py (next to @log_guardrail_information). It should:
- Read
standard_logging_guardrail_informationfromrequest_data["metadata"]orrequest_data["litellm_metadata"](same precedence asadd_standard_logging_guardrail_information_to_request_data) - Append/merge into both
logging_obj.litellm_params["metadata"]andlogging_obj.model_call_details["litellm_params"]["metadata"]whenlogging_objis not None - Not swallow exceptions silently
Call it at the end of the @log_guardrail_information async/sync wrappers, after _process_response / _process_error, when kwargs.get("logging_obj") is present. That fixes /v1/messages, /v1/responses, and every other LITELLM_METADATA_ROUTES guardrail in one place (headroom, presidio, generic_guardrail_api, etc.), not just anthropic.
3. Add a regression test for the sync helper
New test file e.g. tests/test_litellm/integrations/test_guardrail_logging_sync.py. Model it on tests/test_litellm/proxy/test_proxy_utils.py::test_proxy_only_error_log_keeps_litellm_metadata_in_litellm_params. Assert that when guardrail info is written only to request_data["litellm_metadata"] and logging_obj.litellm_params starts empty, the helper copies it into logging_obj.litellm_params["metadata"] so merge_litellm_metadata surfaces it in spend logs.
4. Fix response.json() error handling in headroom.py
Wrap response.json() in try/except and raise 502 with a clear detail. A 200 with HTML/truncated JSON should not become an unhandled 500.
5. Minor cleanups
headroom.py: drop redundantheaders.get(BYPASS_HEADER.lower())test_headroom.py::test_init_raises_without_api_base: monkeypatch env/secrets so the test is deterministic- Revert the whitespace-only diff in
litellm_logging.py - Squash the debug-print commits before merge
Keep as-is
- Headroom guardrail implementation and tests
- structured_messages identity-check write-back (OpenAI + Anthropic handlers)
- Anthropic system-message strip before
anthropic_messages_ptandcache_controlstrip on thinking blocks
| logging_obj=litellm_logging_obj, | ||
| ) | ||
|
|
||
| if litellm_logging_obj is not None: |
There was a problem hiding this comment.
Remove this block; fix at the decorator layer instead.
This copies standard_logging_guardrail_information from data into logging_obj after every guardrail run. Right symptom, wrong location.
Change: Delete lines 166–200 entirely. Add _sync_guardrail_info_to_logging_obj(request_data, logging_obj) in litellm/integrations/custom_guardrail.py and invoke it from @log_guardrail_information when kwargs.get("logging_obj") is present (after _process_response / _process_error). The decorator already has both request_data and logging_obj in scope for apply_guardrail calls.
Why: On LITELLM_METADATA_ROUTES (/v1/messages, etc.) @log_guardrail_information writes to litellm_metadata on data, but spend logs read from logging_obj.litellm_params via merge_litellm_metadata. That gap affects all unified guardrails on those routes, not just headroom.
Also: Replace except (KeyError, AttributeError): pass with no catch or explicit debug logging.
| }, | ||
| ) | ||
|
|
||
| body: object = response.json() |
There was a problem hiding this comment.
Wrap response.json() and return 502 on decode failure.
If the headroom service returns HTTP 200 with a non-JSON body, response.json() raises and propagates as an unhandled 500.
Change:
try:
body: object = response.json()
except Exception:
raise HTTPException(
status_code=502,
detail={
"error": "Headroom compression service returned non-JSON response",
"body": response.text,
},
)| headers = psr.get("headers") | ||
| if not _is_str_object_dict(headers): | ||
| return False | ||
| value = headers.get(BYPASS_HEADER) or headers.get(BYPASS_HEADER.lower()) |
There was a problem hiding this comment.
Simplify header lookup.
BYPASS_HEADER is already lowercase. The or headers.get(BYPASS_HEADER.lower()) branch is identical.
Change: value = headers.get(BYPASS_HEADER)
| assert "empty message list" in str(exc_info.value.detail) | ||
|
|
||
|
|
||
| def test_init_raises_without_api_base(): |
There was a problem hiding this comment.
Make this test deterministic in CI.
If HEADROOM_API_BASE is set in the environment, HeadroomGuardrail(api_base=None) resolves from env and the test never raises.
Change:
def test_init_raises_without_api_base(monkeypatch):
monkeypatch.delenv("HEADROOM_API_BASE", raising=False)
monkeypatch.setattr(
"litellm.proxy.guardrails.guardrail_hooks.headroom.headroom.get_secret_str",
lambda key: None,
)
with pytest.raises(ValueError, match="API base URL"):
HeadroomGuardrail(api_base=None)…guardrail.py - Add _sync_guardrail_info_to_logging_obj in custom_guardrail.py; call it from both async and sync wrappers in @log_guardrail_information, fixing guardrail_information=null in spend logs for all passthrough routes (/v1/messages, /v1/responses, etc.) in one place - Remove the 35-line inline sync block from the anthropic translation handler - Wrap response.json() in try/except in headroom.py to 502 on HTML/truncated responses - Drop redundant headers.get(BYPASS_HEADER.lower()) — header key already lowercase - Add regression tests for _sync_guardrail_info_to_logging_obj
|
All required changes from the review addressed in caf42dc: 1. Inline sync block removed from anthropic translation handler. The 35-line block that copied 2. Shared helper added and called from 3. Regression test added. 4. 5. Minor cleanups done. Dropped |
krrish-berri-2
left a comment
There was a problem hiding this comment.
Two small follow-ups from the earlier review that are still open in caf42dc6 (the refactor itself looks good)
| @@ -5959,7 +5959,6 @@ def get_standard_logging_object_payload( | |||
| ), | |||
| standard_built_in_tools_params=standard_built_in_tools_params, | |||
| ) | |||
There was a problem hiding this comment.
Revert this whitespace-only change.
This PR removes a blank line in get_standard_logging_object_payload; it is unrelated to headroom or guardrail logging sync.
Change: restore the blank line so this file has zero diff vs base:
)
+
# emit_standard_logging_payload(payload) - Moved to success_handler to prevent double emittingKeeps the PR scope clean and avoids noisy merge conflicts on an unrelated 6k-line file.
| assert "empty message list" in str(exc_info.value.detail) | ||
|
|
||
|
|
||
| def test_init_raises_without_api_base(): |
There was a problem hiding this comment.
Make this test deterministic when HEADROOM_API_BASE is set in CI.
HeadroomGuardrail(api_base=None) falls back to get_secret_str("HEADROOM_API_BASE"), so the test can silently pass without asserting anything if that env var is present.
Change:
def test_init_raises_without_api_base(monkeypatch):
monkeypatch.delenv("HEADROOM_API_BASE", raising=False)
monkeypatch.setattr(
"litellm.proxy.guardrails.guardrail_hooks.headroom.headroom.get_secret_str",
lambda key: None,
)
with pytest.raises(ValueError, match="API base URL"):
HeadroomGuardrail(api_base=None)…input_messages complexity
|
@greptile review |
1 similar comment
|
@greptile review |
Relevant issues
Linear ticket
Pre-Submission checklist
make test-unit@greptileaiand received a Confidence Score of at least 4/5 before requesting a maintainer reviewDelays in PR merge?
If you're seeing a delay in your PR being merged, ping the LiteLLM Team on Slack (#pr-review).
CI (LiteLLM team)
Branch creation CI run
Link:
CI run for the last commit
Link:
Merge / cherry-pick CI run
Links:
Screenshots / Proof of Fix
End-to-end trace on a live proxy with real DB confirmed the full chain works:
To test: run the proxy with a config that includes the headroom guardrail (with
default_on: trueor via key metadata), send a/v1/messagesrequest, and verifyguardrail_informationis non-null in the spend log.Type
🆕 New Feature + 🐛 Bug Fixes
Changes
Core feature: headroom guardrail for message compression
Adds a new
headroomguardrail type that compresses request messages viaPOST /v1/compresson a configurable headroom service before they reach the LLM. Useful for reducing token costs on long-context conversations.Config:
Key behaviors:
apply_guardrail, works for both/v1/chat/completionsand/v1/messagesx-headroom-bypass: truerequest header skips compressionapi_key/api_basein guardrail config orHEADROOM_API_KEY/HEADROOM_API_BASEenv vars@log_guardrail_informationsoguardrail_informationis populated in spend logsBug fixes made during development:
structured_messageswrite-back regression: the unified guardrail handler for both OpenAI and Anthropic paths was unconditionally writing backstructured_messageseven when the guardrail returned inputs unchanged. This clobbereddata["messages"]with a filtered copy (e.g. skip-system filtered messages). Fixed with an identity check — only write back when the guardrail actually returned a new object.anthropic_messages_ptcrash on system messages with cache_control:translate_anthropic_to_openaipreserves system messages with list content (cache_control blocks) as-is.anthropic_messages_ptrejects these. Fixed by stripping system messages before the reverse-translation call — they live indata["system"]in the Anthropic passthrough path, notdata["messages"].thinking.cache_controlrejected by Anthropic: theprompt-caching-scopebeta addscache_controlto all content blocks includingthinkingblocks, which Anthropic rejects. Fixed by strippingcache_controlfromthinkingtype content blocks after the reverse-translation.guardrail_informationnull in spend logs for/v1/messages: multiple root causes tracked down and fixed. The@log_guardrail_informationdecorator writes torequest_data["metadata"]orrequest_data["litellm_metadata"]depending on which key is present. The spend log is built fromlitellm_params["metadata"]viamerge_litellm_metadata. For/v1/messages,dataonly haslitellm_metadata(nometadatakey), andLogging.litellm_paramsgets reassigned during the request (creating a new dict that diverges frommodel_call_details["litellm_params"]). Fix: afterapply_guardrailreturns, explicitly copy the guardrail info from whichever metadata key it was written to, into bothlogging_obj.litellm_params["metadata"]andlogging_obj.model_call_details["litellm_params"]["metadata"].New files:
litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.pylitellm/proxy/guardrails/guardrail_hooks/headroom/__init__.pylitellm/types/proxy/guardrails/guardrail_hooks/headroom.pytests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.pytests/test_litellm/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py