diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index 2d67c22f0aa..775f3559c55 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -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, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py index 7f412c008ca..7345a86ebdc 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -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,