scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(guardrails/headroom): log skipped compression as not_run instead of success

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Krrish Dholakia 2026-07-20 22:57:15 +00:00
parent 68be053e96
commit 5eaf5b2b60
2 changed files with 133 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -31,7 +31,7 @@ from litellm.llms.custom_httpx.http_handler import (
from litellm.secret_managers.main import get_secret_str
from litellm.types.guardrails import GuardrailEventHooks, Mode
from litellm.types.integrations.custom_logger import AgenticLoopPlan, AgenticLoopRequestPatch
from litellm.types.utils import GenericGuardrailAPIInputs
from litellm.types.utils import GenericGuardrailAPIInputs, GuardrailStatus
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -459,18 +459,38 @@ class HeadroomGuardrail(CustomGuardrail):
logging_obj: LiteLLMLoggingObj | None = None,
) -> GenericGuardrailAPIInputs:
if input_type != "request":
self._record_guardrail_result(
request_data=request_data,
guardrail_status="not_run",
guardrail_json_response={"skipped": True, "reason": "unsupported_input_type"},
)
return inputs
if self._should_bypass(request_data):
verbose_proxy_logger.debug("Headroom: %s header set; skipping compression", BYPASS_HEADER)
self._record_guardrail_result(
request_data=request_data,
guardrail_status="not_run",
guardrail_json_response={"skipped": True, "reason": "bypass_header"},
)
return inputs
structured_messages = inputs.get("structured_messages")
if not _is_object_list(structured_messages) or not structured_messages:
self._record_guardrail_result(
request_data=request_data,
guardrail_status="not_run",
guardrail_json_response={"skipped": True, "reason": "no_structured_messages"},
)
return inputs
messages = [m for m in structured_messages if _is_str_object_dict(m)]
if not messages:
self._record_guardrail_result(
request_data=request_data,
guardrail_status="not_run",
guardrail_json_response={"skipped": True, "reason": "no_valid_messages"},
)
return inputs
model = self.headroom_model or request_data.get("model")
@ -482,6 +502,13 @@ class HeadroomGuardrail(CustomGuardrail):
end_time = time.time()
if not compression_succeeded:
self._record_guardrail_result(
request_data=request_data,
guardrail_status="guardrail_failed_to_respond",
guardrail_json_response={"skipped": False, "reason": "compression_failed"},
start_time=start_time,
end_time=end_time,
)
return {**inputs, "structured_messages": compressed} # pyright: ignore[reportReturnType]
self.add_standard_logging_guardrail_information_to_request_data(
@ -516,6 +543,26 @@ class HeadroomGuardrail(CustomGuardrail):
return {**inputs, "structured_messages": compressed, "tools": merged_tools} # pyright: ignore[reportReturnType]
def _record_guardrail_result(
self,
request_data: dict,
guardrail_status: GuardrailStatus,
guardrail_json_response: dict[str, object],
start_time: float | None = None,
end_time: float | None = None,
) -> None:
resolved_start_time = start_time if start_time is not None else time.time()
resolved_end_time = end_time if end_time is not None else resolved_start_time
self.add_standard_logging_guardrail_information_to_request_data(
guardrail_json_response=guardrail_json_response,
request_data=request_data,
guardrail_status=guardrail_status,
guardrail_provider=HEADROOM_GUARDRAIL_PROVIDER,
start_time=resolved_start_time,
end_time=resolved_end_time,
duration=resolved_end_time - resolved_start_time,
)
async def async_should_run_agentic_loop(
self,
response: Any,

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

@ -139,6 +139,91 @@ async def test_apply_guardrail_compresses_and_returns_structured_messages(
assert result.get("structured_messages") == COMPRESSED_MESSAGES
@pytest.mark.asyncio
async def test_apply_guardrail_skips_passthrough_without_auto_success_entry(
guardrail: HeadroomGuardrail,
):
inputs = GenericGuardrailAPIInputs(texts=["A" * 5000])
request_data: dict = {"model": "gpt-4o"}
with patch.object(
guardrail.async_handler,
"post",
new_callable=AsyncMock,
) as mock_post:
result = await guardrail.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="request",
)
entries = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(entries) == 1
assert entries[0]["guardrail_status"] == "not_run"
assert entries[0]["guardrail_response"] == {
"skipped": True,
"reason": "no_structured_messages",
}
assert result == inputs
mock_post.assert_not_called()
@pytest.mark.asyncio
async def test_apply_guardrail_records_success_for_compression(
guardrail: HeadroomGuardrail,
):
inputs = GenericGuardrailAPIInputs(
texts=["A" * 5000],
structured_messages=ORIGINAL_MESSAGES,
)
request_data: dict = {"model": "gpt-4o"}
mock_response = _make_compress_response(COMPRESSED_MESSAGES)
with patch.object(
guardrail.async_handler,
"post",
new_callable=AsyncMock,
return_value=mock_response,
):
await guardrail.apply_guardrail(
inputs=inputs,
request_data=request_data,
input_type="request",
)
entries = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(entries) == 1
assert entries[0]["guardrail_status"] == "success"
@pytest.mark.asyncio
async def test_apply_guardrail_records_failure_and_preserves_input():
guardrail = _make_guardrail(unreachable_fallback="fail_open")
inputs = GenericGuardrailAPIInputs(
texts=["A" * 5000],
structured_messages=ORIGINAL_MESSAGES,
)
request_data: dict = {"model": "gpt-4o"}
mock_response = _make_compress_response(COMPRESSED_MESSAGES, status=500)
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",
)
entries = request_data["metadata"]["standard_logging_guardrail_information"]
assert len(entries) == 1
assert entries[0]["guardrail_status"] == "guardrail_failed_to_respond"
assert result == inputs
@pytest.mark.asyncio
async def test_apply_guardrail_injects_retrieve_tool_when_hashes_present(
guardrail: HeadroomGuardrail,