feat(guardrails): add headroom guardrail for message compression

litellm/integrations/custom_guardrail.py

@@ -1093,6 +1093,45 @@ def get_guardrails_messages_for_call_type(
         return None
 
 
+def _append_slg_to_litellm_params(lp: object, entries: list) -> None:
+    """Merge guardrail entries into a single litellm_params dict."""
+    if not isinstance(lp, dict):
+        return
+    if lp.get("metadata") is None:
+        lp["metadata"] = {}
+    existing = lp["metadata"].setdefault("standard_logging_guardrail_information", [])
+    for entry in entries:
+        if entry not in existing:
+            existing.append(entry)
+
+
+def _sync_guardrail_info_to_logging_obj(
+    request_data: dict, logging_obj: object
+) -> None:
+    """Copy standard_logging_guardrail_information from request_data into logging_obj.
+
+    The @log_guardrail_information decorator writes guardrail info to
+    request_data["metadata"] or request_data["litellm_metadata"].  For
+    passthrough routes (/v1/messages, /v1/responses) the spend-log payload is
+    built from logging_obj.litellm_params["metadata"], which is a separate dict
+    that does not share identity with the one in request_data.  This helper
+    bridges that gap so guardrail_information is non-null in spend logs for all
+    routes, not just /v1/chat/completions.
+    """
+    if logging_obj is None:
+        return
+    meta_src = (
+        request_data.get("metadata") or request_data.get("litellm_metadata") or {}
+    )
+    slg_info = meta_src.get("standard_logging_guardrail_information")
+    if not slg_info:
+        return
+    entries: list = slg_info if isinstance(slg_info, list) else [slg_info]
+    mcd = getattr(logging_obj, "model_call_details", None) or {}
+    _append_slg_to_litellm_params(getattr(logging_obj, "litellm_params", None), entries)
+    _append_slg_to_litellm_params(mcd.get("litellm_params"), entries)
+
+
 def log_guardrail_information(func):
     """
     Decorator to add standard logging guardrail information to any function
@@ -1153,6 +1192,7 @@ async def async_wrapper(*args, **kwargs):
         if func.__name__ == "apply_guardrail" and "inputs" in kwargs:
             original_inputs = kwargs.get("inputs")
 
+        logging_obj = kwargs.get("logging_obj")
         entries_before = _count_recorded_guardrail_entries(request_data)
         try:
             response = await func(*args, **kwargs)
@@ -1178,6 +1218,8 @@ async def async_wrapper(*args, **kwargs):
                 duration=(datetime.now() - start_time).total_seconds(),
                 event_type=event_type,
             )
+        finally:
+            _sync_guardrail_info_to_logging_obj(request_data, logging_obj)
 
     @functools.wraps(func)
     def sync_wrapper(*args, **kwargs):
@@ -1191,6 +1233,7 @@ def sync_wrapper(*args, **kwargs):
         if func.__name__ == "apply_guardrail" and "inputs" in kwargs:
             original_inputs = kwargs.get("inputs")
 
+        logging_obj = kwargs.get("logging_obj")
         entries_before = _count_recorded_guardrail_entries(request_data)
         try:
             response = func(*args, **kwargs)
@@ -1212,6 +1255,8 @@ def sync_wrapper(*args, **kwargs):
                 duration=(datetime.now() - start_time).total_seconds(),
                 event_type=event_type,
             )
+        finally:
+            _sync_guardrail_info_to_logging_obj(request_data, logging_obj)
 
     @functools.wraps(func)
     def wrapper(*args, **kwargs):

litellm/litellm_core_utils/litellm_logging.py

@@ -5959,7 +5959,6 @@ def get_standard_logging_object_payload(
             ),
             standard_built_in_tools_params=standard_built_in_tools_params,
         )
-
         # emit_standard_logging_payload(payload) - Moved to success_handler to prevent double emitting
 
         return payload

litellm/llms/anthropic/chat/guardrail_translation/handler.py

@@ -149,6 +149,7 @@ async def process_input_messages(
                 inputs["images"] = images_to_check
             if tools_to_check:
                 inputs["tools"] = tools_to_check
+            original_structured_messages = structured_messages
             if structured_messages:
                 inputs["structured_messages"] = structured_messages
             # Include model information if available
@@ -175,19 +176,50 @@ async def process_input_messages(
                     # Note: MCP servers are handled separately in the main transformation
                 data["tools"] = anthropic_tools
 
-            # Step 3: Map guardrail responses back to original message structure
-            await self._apply_guardrail_responses_to_input(
-                messages=messages,
-                responses=guardrailed_texts,
-                task_mappings=task_mappings,
+            guardrailed_structured_messages = guardrailed_inputs.get(
+                "structured_messages"
             )
+            if (
+                guardrailed_structured_messages is not None
+                and guardrailed_structured_messages is not original_structured_messages
+            ):
+                self._write_back_structured_messages(
+                    data, guardrailed_structured_messages
+                )
+            else:
+                # Step 3: Map guardrail responses back to original message structure
+                await self._apply_guardrail_responses_to_input(
+                    messages=messages,
+                    responses=guardrailed_texts,
+                    task_mappings=task_mappings,
+                )
 
         verbose_proxy_logger.debug(
             "Anthropic Messages: Processed input messages: %s", messages
         )
 
         return data
 
+    @staticmethod
+    def _write_back_structured_messages(data: dict, structured_messages: list) -> None:
+        """Convert compressed structured_messages back to Anthropic format and write to data."""
+        from litellm.litellm_core_utils.prompt_templates.factory import (
+            anthropic_messages_pt,
+        )
+
+        model = str(data.get("model") or "")
+        non_system = [m for m in structured_messages if m.get("role") != "system"]
+        converted = anthropic_messages_pt(
+            messages=non_system, model=model, llm_provider="anthropic"
+        )
+        for msg in converted:
+            content = msg.get("content")
+            if isinstance(content, list):
+                for block in content:
+                    if isinstance(block, dict) and block.get("type") == "thinking":
+                        block.pop("cache_control", None)
+        data["messages"] = converted
+
     def extract_request_tool_names(self, data: dict) -> List[str]:
         """Extract tool names from Anthropic messages request (tools[].name)."""
         names: List[str] = []

litellm/llms/openai/chat/guardrail_translation/handler.py

@@ -124,6 +124,7 @@ async def process_input_messages(
             if model:
                 inputs["model"] = model
 
+            original_structured_messages = inputs.get("structured_messages")
             guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
                 inputs=inputs,
                 request_data=data,
@@ -137,26 +138,34 @@ async def process_input_messages(
             if guardrailed_tools is not None:
                 data["tools"] = guardrailed_tools
 
-            # Step 3: Map guardrail responses back to original message structure
-            if guardrailed_texts and texts_to_check:
-                await self._apply_guardrail_responses_to_input_texts(
-                    messages=messages,
-                    responses=guardrailed_texts,
-                    task_mappings=text_task_mappings,
-                )
+            guardrailed_structured_messages = guardrailed_inputs.get(
+                "structured_messages"
+            )
+            if (
+                guardrailed_structured_messages is not None
+                and guardrailed_structured_messages is not original_structured_messages
+            ):
+                data["messages"] = guardrailed_structured_messages
+            else:
+                # Step 3: Map guardrail responses back to original message structure
+                if guardrailed_texts and texts_to_check:
+                    await self._apply_guardrail_responses_to_input_texts(
+                        messages=messages,
+                        responses=guardrailed_texts,
+                        task_mappings=text_task_mappings,
+                    )
 
-            # Step 4: Apply guardrailed tool calls back to messages
-            if guardrailed_tool_calls:
-                # Note: The guardrail may modify tool_calls_to_check in place
-                # or we may need to handle returned tool calls differently
-                await self._apply_guardrail_responses_to_input_tool_calls(
-                    messages=messages,
-                    tool_calls=guardrailed_tool_calls,  # type: ignore
-                    task_mappings=tool_call_task_mappings,
-                )
+                # Step 4: Apply guardrailed tool calls back to messages
+                if guardrailed_tool_calls:
+                    await self._apply_guardrail_responses_to_input_tool_calls(
+                        messages=messages,
+                        tool_calls=guardrailed_tool_calls,  # type: ignore
+                        task_mappings=tool_call_task_mappings,
+                    )
 
         verbose_proxy_logger.debug(
-            "OpenAI Chat Completions: Processed input messages: %s", messages
+            "OpenAI Chat Completions: Processed input messages: %s",
+            data.get("messages"),
         )
 
         return data

litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py

@@ -0,0 +1,52 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from litellm.types.guardrails import (
+    GuardrailEventHooks,
+    Mode,
+    SupportedGuardrailIntegrations,
+)
+
+from .headroom import HeadroomGuardrail
+
+if TYPE_CHECKING:
+    from litellm.types.guardrails import Guardrail, LitellmParams
+
+
+def _coerce_event_hook(
+    mode: str | list[str] | Mode,
+) -> GuardrailEventHooks | list[GuardrailEventHooks] | Mode:
+    if isinstance(mode, Mode):
+        return mode
+    if isinstance(mode, list):
+        return [GuardrailEventHooks(item) for item in mode]
+    return GuardrailEventHooks(mode)
+
+
+def initialize_guardrail(
+    litellm_params: LitellmParams, guardrail: Guardrail
+) -> HeadroomGuardrail:
+    import litellm
+
+    _callback = HeadroomGuardrail(
+        api_base=litellm_params.api_base,
+        api_key=litellm_params.api_key,
+        model=litellm_params.model,
+        guardrail_name=guardrail["guardrail_name"],
+        event_hook=_coerce_event_hook(litellm_params.mode),
+        default_on=litellm_params.default_on or False,
+    )
+    litellm.logging_callback_manager.add_litellm_callback(  # pyright: ignore[reportUnknownMemberType]
+        _callback
+    )
+    return _callback
+
+
+guardrail_initializer_registry = {
+    SupportedGuardrailIntegrations.HEADROOM.value: initialize_guardrail,
+}
+
+guardrail_class_registry = {
+    SupportedGuardrailIntegrations.HEADROOM.value: HeadroomGuardrail,
+}