Commit graph

1 commit

Author SHA1 Message Date
Krrish Dholakia
99b1a323c1
feat(guardrails): add headroom guardrail for message compression (#31407)
* feat(guardrails): add headroom guardrail for message compression

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.

* fix(guardrails/headroom): add @log_guardrail_information to populate guardrail_information in spend logs

* style: fix ruff format violations

* fix(lint): replace deprecated typing aliases with builtin generics (UP006/UP037)

* fix(guardrails): only write back structured_messages when guardrail actually changed them

* fix(guardrails/headroom): raise 502 when compression returns empty message list

* fix(guardrails/headroom): catch transport errors and fix stale debug log

* fix(guardrails/anthropic): strip system messages before anthropic_messages_pt reverse-translation

* fix(guardrails/anthropic): strip cache_control from thinking blocks after write-back

* debug(headroom): add INFO logging to trace guardrail execution

* debug(headroom): use print() for immediate visibility

* debug(headroom): print request_data keys to diagnose metadata dict mismatch

* fix(guardrails/anthropic): propagate guardrail info to logging_obj.metadata for spend log

* fix: use model_call_details litellm_params metadata on Logging object

* fix(guardrails/anthropic): write guardrail info to litellm_params attr not model_call_details copy

* fix: read slg_info from litellm_metadata when metadata key absent

* fix: write slg_info to both litellm_params attr and model_call_details copy

* chore: remove debug prints; fix now verified end-to-end

* refactor(guardrails): move spend-log sync to shared helper in custom_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

* fix(lint): reduce _sync_guardrail_info_to_logging_obj complexity below C901 threshold

* fix(lint): simplify _sync_guardrail_info_to_logging_obj to reduce McCabe complexity

* fix(lint): extract _append_slg_to_litellm_params to reduce McCabe complexity

* fix(lint): extract _write_back_structured_messages to reduce process_input_messages complexity
2026-06-26 19:36:44 -07:00