mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
* 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
124 lines
4.3 KiB
Python
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]
|