litellm/tests/test_litellm/integrations/test_guardrail_logging_sync.py
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

124 lines
4.3 KiB
Python

"""
Regression test for _sync_guardrail_info_to_logging_obj.
Ensures that when the @log_guardrail_information decorator writes guardrail info
to request_data["litellm_metadata"] (as it does for /v1/messages passthrough
routes that have no "metadata" key), the helper propagates it into
logging_obj.litellm_params["metadata"] so merge_litellm_metadata surfaces it in
spend logs.
"""
import pytest
from litellm.integrations.custom_guardrail import _sync_guardrail_info_to_logging_obj
def _make_slg_entry(name: str = "headroom-test") -> dict:
return {
"guardrail_name": name,
"guardrail_response": "mask",
"guardrail_status": "success",
"duration": 0.1,
}
class _FakeLogging:
"""Minimal stand-in for litellm.litellm_core_utils.litellm_logging.Logging."""
def __init__(self, lp_metadata: dict | None = None):
self.litellm_params: dict = {"metadata": lp_metadata or {}}
self.model_call_details: dict = {"litellm_params": self.litellm_params}
def test_syncs_from_litellm_metadata_key():
"""When guardrail info is in request_data["litellm_metadata"], it is copied."""
entry = _make_slg_entry()
request_data = {
"litellm_metadata": {"standard_logging_guardrail_information": [entry]}
}
logging_obj = _FakeLogging()
_sync_guardrail_info_to_logging_obj(request_data, logging_obj)
result = logging_obj.litellm_params["metadata"].get(
"standard_logging_guardrail_information"
)
assert result == [entry]
def test_syncs_from_metadata_key():
"""When guardrail info is in request_data["metadata"], it is also copied."""
entry = _make_slg_entry()
request_data = {"metadata": {"standard_logging_guardrail_information": [entry]}}
logging_obj = _FakeLogging()
_sync_guardrail_info_to_logging_obj(request_data, logging_obj)
result = logging_obj.litellm_params["metadata"].get(
"standard_logging_guardrail_information"
)
assert result == [entry]
def test_metadata_wins_over_litellm_metadata():
"""metadata key takes precedence over litellm_metadata when both are present."""
entry_meta = _make_slg_entry("from-metadata")
entry_lm = _make_slg_entry("from-litellm_metadata")
request_data = {
"metadata": {"standard_logging_guardrail_information": [entry_meta]},
"litellm_metadata": {"standard_logging_guardrail_information": [entry_lm]},
}
logging_obj = _FakeLogging()
_sync_guardrail_info_to_logging_obj(request_data, logging_obj)
result = logging_obj.litellm_params["metadata"].get(
"standard_logging_guardrail_information"
)
assert result == [entry_meta]
def test_noop_when_no_guardrail_info():
"""Does nothing when standard_logging_guardrail_information is absent."""
request_data = {"litellm_metadata": {"other_key": "value"}}
logging_obj = _FakeLogging()
_sync_guardrail_info_to_logging_obj(request_data, logging_obj)
assert (
logging_obj.litellm_params["metadata"].get(
"standard_logging_guardrail_information"
)
is None
)
def test_noop_when_logging_obj_is_none():
"""Does nothing when logging_obj is None."""
entry = _make_slg_entry()
request_data = {
"litellm_metadata": {"standard_logging_guardrail_information": [entry]}
}
_sync_guardrail_info_to_logging_obj(request_data, None)
def test_writes_to_model_call_details_too():
"""Also writes into model_call_details["litellm_params"]["metadata"]."""
entry = _make_slg_entry()
request_data = {
"litellm_metadata": {"standard_logging_guardrail_information": [entry]}
}
logging_obj = _FakeLogging()
# Simulate litellm_params reassignment (creating a new dict) — model_call_details
# then points to the OLD dict while litellm_params points to the new one.
old_lp = logging_obj.litellm_params
logging_obj.litellm_params = {**old_lp, "extra": "added"}
logging_obj.model_call_details["litellm_params"] = old_lp # diverged
_sync_guardrail_info_to_logging_obj(request_data, logging_obj)
# Both dicts should have the info.
assert logging_obj.litellm_params["metadata"].get(
"standard_logging_guardrail_information"
) == [entry]
assert old_lp["metadata"].get("standard_logging_guardrail_information") == [entry]