From 99b1a323c13ddf2cbc798558abc035339787a071 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 26 Jun 2026 19:36:44 -0700 Subject: [PATCH] feat(guardrails): add headroom guardrail for message compression (#31407) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * 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 --- litellm/integrations/custom_guardrail.py | 45 +++ litellm/litellm_core_utils/litellm_logging.py | 1 - .../chat/guardrail_translation/handler.py | 42 ++- .../chat/guardrail_translation/handler.py | 43 ++- .../guardrail_hooks/headroom/__init__.py | 52 +++ .../guardrail_hooks/headroom/headroom.py | 210 +++++++++++ litellm/types/guardrails.py | 5 + .../guardrails/guardrail_hooks/headroom.py | 24 ++ .../test_guardrail_logging_sync.py | 124 +++++++ .../guardrail_hooks/test_headroom.py | 346 ++++++++++++++++++ .../test_structured_messages_writeback.py | 129 +++++++ 11 files changed, 998 insertions(+), 23 deletions(-) create mode 100644 litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py create mode 100644 litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/headroom.py create mode 100644 tests/test_litellm/integrations/test_guardrail_logging_sync.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 38245a2e5ba..158ef0b4356 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -1093,6 +1093,45 @@ class CustomGuardrail(CustomLogger): 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 @@ def log_guardrail_information(func): 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 @@ def log_guardrail_information(func): 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 log_guardrail_information(func): 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 log_guardrail_information(func): 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): diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index e860f48a7bd..1174907e31d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/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 diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index 7d5944b4cc2..c65c0e02cfc 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -149,6 +149,7 @@ class AnthropicMessagesHandler(BaseTranslation): 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,12 +176,23 @@ class AnthropicMessagesHandler(BaseTranslation): # 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 @@ -188,6 +200,26 @@ class AnthropicMessagesHandler(BaseTranslation): 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] = [] diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index 3b5980023c7..a62dd939a71 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -124,6 +124,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): 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 @@ class OpenAIChatCompletionsHandler(BaseTranslation): 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 diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py new file mode 100644 index 00000000000..cf1b3bc97e2 --- /dev/null +++ b/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, +} diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py new file mode 100644 index 00000000000..9d63e5d144a --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -0,0 +1,210 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Literal + +import httpx +from fastapi import HTTPException +from httpx import Response as HttpxResponse +from typing_extensions import TypeGuard + +from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.llms.custom_httpx.http_handler import ( + get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] + httpxSpecialProvider, +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.guardrails import GuardrailEventHooks, Mode +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel + +BYPASS_HEADER = "x-headroom-bypass" + + +def _is_str_object_dict(value: object) -> TypeGuard[dict[str, object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip + return isinstance(value, dict) + + +def _is_object_list(value: object) -> TypeGuard[list[object]]: # guard-ok: isinstance narrows correctly; predicate is trivially correct # fmt: skip + return isinstance(value, list) + + +class HeadroomGuardrail(CustomGuardrail): + def __init__( + self, + api_base: str | None = None, + api_key: str | None = None, + model: str | None = None, + guardrail_name: str | None = None, + event_hook: GuardrailEventHooks + | list[GuardrailEventHooks] + | Mode + | None = None, + default_on: bool = False, + ): + self.headroom_api_base = ( + api_base or get_secret_str("HEADROOM_API_BASE") or "" + ).rstrip("/") + if not self.headroom_api_base: + raise ValueError( + "Headroom guardrail requires an API base URL. " + "Set `api_base` in the guardrail config or HEADROOM_API_BASE env var." + ) + self.headroom_api_key = api_key or get_secret_str("HEADROOM_API_KEY") + self.headroom_model = model + self.async_handler = get_async_httpx_client( + llm_provider=httpxSpecialProvider.GuardrailCallback, + ) + super().__init__( # pyright: ignore[reportUnknownMemberType] + guardrail_name=guardrail_name, + event_hook=event_hook, + default_on=default_on, + ) + + def _should_bypass(self, request_data: dict) -> bool: + psr = request_data.get("proxy_server_request") + if not _is_str_object_dict(psr): + return False + headers = psr.get("headers") + if not _is_str_object_dict(headers): + return False + value = headers.get(BYPASS_HEADER) + return str(value).lower() == "true" + + async def _call_compress( + self, + messages: list[dict[str, object]], + model: str | None, + ) -> list[dict[str, object]]: + payload: dict[str, object] = {"messages": messages} + if model: + payload["model"] = model + + request_headers: dict[str, str] = {"Content-Type": "application/json"} + if self.headroom_api_key: + request_headers["Authorization"] = f"Bearer {self.headroom_api_key}" + + try: + raw_response: HttpxResponse | None = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] + url=f"{self.headroom_api_base}/v1/compress", + json=payload, + headers=request_headers, + ) + except (httpx.ConnectError, httpx.TimeoutException, httpx.TransportError) as e: + raise HTTPException( + status_code=502, + detail={ + "error": "Headroom compression service unreachable", + "detail": str(e), + }, + ) from e + if raw_response is None: + raise HTTPException( + status_code=502, + detail={"error": "Headroom compression service returned no response"}, + ) + response: HttpxResponse = raw_response + + if response.status_code != 200: + raise HTTPException( + status_code=502, + detail={ + "error": "Headroom compression service returned an error", + "status_code": response.status_code, + "body": response.text, + }, + ) + + try: + body: object = response.json() + except Exception: + raise HTTPException( + status_code=502, + detail={ + "error": "Headroom compression service returned non-JSON response", + "body": response.text[:500], + }, + ) + if not _is_str_object_dict(body): + raise HTTPException( + status_code=502, + detail={ + "error": "Headroom compression service returned unexpected response shape", + "body": response.text[:500], + }, + ) + + compressed_messages = body.get("messages") + if not _is_object_list(compressed_messages): + raise HTTPException( + status_code=502, + detail={ + "error": "Headroom compression service response missing 'messages'", + "body": response.text, + }, + ) + + filtered = [item for item in compressed_messages if _is_str_object_dict(item)] + if not filtered: + raise HTTPException( + status_code=502, + detail={ + "error": "Headroom compression service returned empty message list", + "body": response.text, + }, + ) + + verbose_proxy_logger.debug( + "Headroom: compressed %s tokens -> %s tokens (ratio %.2f)", + body.get("tokens_before", "?"), + body.get("tokens_after", "?"), + body.get("compression_ratio", 0), + ) + return filtered + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + if input_type != "request": + return inputs + + if self._should_bypass(request_data): + verbose_proxy_logger.debug( + "Headroom: %s header set; skipping compression", BYPASS_HEADER + ) + return inputs + + structured_messages = inputs.get("structured_messages") + if not _is_object_list(structured_messages) or not structured_messages: + return inputs + + messages = [m for m in structured_messages if _is_str_object_dict(m)] + if not messages: + return inputs + + model = self.headroom_model or request_data.get("model") + compressed = await self._call_compress( + messages=messages, + model=model if isinstance(model, str) else None, + ) + + return {**inputs, "structured_messages": compressed} # pyright: ignore[reportReturnType] + + @staticmethod + def get_config_model() -> type[GuardrailConfigModel[object]] | None: + from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( + HeadroomGuardrailConfigModel, + ) + + return HeadroomGuardrailConfigModel diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 6a76dd337fa..e286c6f7e0d 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -53,6 +53,9 @@ from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import ( from litellm.types.proxy.guardrails.guardrail_hooks.cisco_ai_defense import ( CiscoAIDefenseGuardrailConfigModel, ) +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( + HeadroomGuardrailConfigModel, +) """ Pydantic object defining how to set guardrails on litellm proxy @@ -119,6 +122,7 @@ class SupportedGuardrailIntegrations(Enum): RUBRIK = "rubrik" VIGIL_GUARD = "vigil_guard" REPELLOAI = "repelloai" + HEADROOM = "headroom" class Role(Enum): @@ -865,6 +869,7 @@ class LitellmParams( PresidioConfigModel, BedrockGuardrailConfigModel, LakeraV2GuardrailConfigModel, + HeadroomGuardrailConfigModel, RepelloAIGuardrailConfigModel, LassoGuardrailConfigModel, PillarGuardrailConfigModel, diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py new file mode 100644 index 00000000000..3186c9fc612 --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py @@ -0,0 +1,24 @@ +from typing import Optional + +from pydantic import BaseModel, Field + +from .base import GuardrailConfigModel + + +class HeadroomGuardrailConfigModel(GuardrailConfigModel[BaseModel]): + api_base: Optional[str] = Field( + default=None, + description="Base URL for the headroom compression service (e.g. https://api.headroom.ai). Falls back to HEADROOM_API_BASE env var.", + ) + api_key: Optional[str] = Field( + default=None, + description="API key for the headroom compression service. Falls back to HEADROOM_API_KEY env var.", + ) + model: Optional[str] = Field( + default=None, + description="Model name forwarded to the headroom /v1/compress endpoint.", + ) + + @staticmethod + def ui_friendly_name() -> str: + return "Headroom" diff --git a/tests/test_litellm/integrations/test_guardrail_logging_sync.py b/tests/test_litellm/integrations/test_guardrail_logging_sync.py new file mode 100644 index 00000000000..5dcd1114b3d --- /dev/null +++ b/tests/test_litellm/integrations/test_guardrail_logging_sync.py @@ -0,0 +1,124 @@ +""" +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] diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py new file mode 100644 index 00000000000..66395035384 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -0,0 +1,346 @@ +""" +Unit tests for the Headroom guardrail. + +Tests cover: +- apply_guardrail compresses messages via /v1/compress and returns them as structured_messages +- x-headroom-bypass: true header causes guardrail to skip compression +- missing or empty messages are passed through unchanged +- response-type input is passed through unchanged +- /v1/compress HTTP error raises HTTPException +- /v1/compress returning malformed JSON raises HTTPException +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import pytest +from fastapi import HTTPException + +from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import HeadroomGuardrail +from litellm.types.utils import GenericGuardrailAPIInputs + +FAKE_API_BASE = "https://headroom.example.com" +FAKE_API_KEY = "test-key" + +ORIGINAL_MESSAGES = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "A" * 5000}, +] +COMPRESSED_MESSAGES = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "A" * 500}, +] + + +def _make_guardrail(**kwargs) -> HeadroomGuardrail: + defaults = dict( + api_base=FAKE_API_BASE, + api_key=FAKE_API_KEY, + guardrail_name="headroom", + default_on=True, + ) + defaults.update(kwargs) + return HeadroomGuardrail(**defaults) + + +def _make_compress_response(messages: list, status: int = 200) -> MagicMock: + mock = MagicMock() + mock.status_code = status + mock.json.return_value = { + "messages": messages, + "tokens_before": 1000, + "tokens_after": 100, + "compression_ratio": 0.1, + "transforms_applied": ["router:smart_crusher:0.35"], + } + mock.text = "" + return mock + + +@pytest.fixture +def guardrail() -> HeadroomGuardrail: + return _make_guardrail() + + +@pytest.mark.asyncio +async def test_apply_guardrail_compresses_and_returns_structured_messages( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["A" * 5000], + structured_messages=ORIGINAL_MESSAGES, + ) + mock_response = _make_compress_response(COMPRESSED_MESSAGES) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + assert result.get("structured_messages") == COMPRESSED_MESSAGES + + +@pytest.mark.asyncio +async def test_apply_guardrail_bypass_header_skips_compression( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + request_data = {"proxy_server_request": {"headers": {"x-headroom-bypass": "true"}}} + + 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", + ) + mock_post.assert_not_called() + + assert result.get("structured_messages") == ORIGINAL_MESSAGES + + +@pytest.mark.asyncio +async def test_apply_guardrail_response_type_passthrough( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs( + texts=["some response text"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post: + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + ) + mock_post.assert_not_called() + + assert result is inputs + + +@pytest.mark.asyncio +async def test_apply_guardrail_empty_structured_messages_passthrough( + guardrail: HeadroomGuardrail, +): + inputs = GenericGuardrailAPIInputs(texts=["hello"]) + + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post: + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + mock_post.assert_not_called() + + assert result is inputs + + +@pytest.mark.asyncio +async def test_apply_guardrail_http_error_raises(): + guardrail = _make_guardrail() + mock_response = _make_compress_response([], status=500) + mock_response.text = "Internal Server Error" + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_apply_guardrail_transport_error_raises(): + guardrail = _make_guardrail() + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + side_effect=httpx.ConnectError("Connection refused"), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + assert "unreachable" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_apply_guardrail_missing_messages_key_raises(): + guardrail = _make_guardrail() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = {"tokens_before": 100, "tokens_after": 10} + mock_response.text = "{}" + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + + +@pytest.mark.asyncio +async def test_apply_guardrail_empty_compressed_messages_raises(): + guardrail = _make_guardrail() + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + "messages": ["not-a-dict", 42, None], + "tokens_before": 1000, + "tokens_after": 0, + "compression_ratio": 0, + } + mock_response.text = "{}" + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert exc_info.value.status_code == 502 + assert "empty message list" in str(exc_info.value.detail) + + +def test_init_raises_without_api_base(): + with pytest.raises(ValueError, match="API base URL"): + HeadroomGuardrail(api_base=None) + + +def test_bypass_header_case_insensitive(): + guardrail = _make_guardrail() + + for header_value in ("true", "True", "TRUE"): + data = { + "proxy_server_request": {"headers": {"x-headroom-bypass": header_value}} + } + assert guardrail._should_bypass(data) is True + + data = {"proxy_server_request": {"headers": {"x-headroom-bypass": "false"}}} + assert guardrail._should_bypass(data) is False + + data = {"proxy_server_request": {"headers": {}}} + assert guardrail._should_bypass(data) is False + + data = {} + assert guardrail._should_bypass(data) is False + + +@pytest.mark.asyncio +async def test_apply_guardrail_sends_model_from_config(): + guardrail = _make_guardrail(model="gpt-4o-mini") + mock_response = _make_compress_response(COMPRESSED_MESSAGES) + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_post: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + call_kwargs = mock_post.call_args + sent_payload = call_kwargs.kwargs.get("json") or call_kwargs.args[1] + assert sent_payload.get("model") == "gpt-4o-mini" + + +@pytest.mark.asyncio +async def test_apply_guardrail_sends_model_from_request_data_when_no_config_model(): + guardrail = _make_guardrail() + mock_response = _make_compress_response(COMPRESSED_MESSAGES) + + inputs = GenericGuardrailAPIInputs( + texts=["hello"], + structured_messages=ORIGINAL_MESSAGES, + ) + + with patch.object( + guardrail.async_handler, + "post", + new_callable=AsyncMock, + return_value=mock_response, + ) as mock_post: + await guardrail.apply_guardrail( + inputs=inputs, + request_data={"model": "gpt-4o"}, + input_type="request", + ) + + call_kwargs = mock_post.call_args + sent_payload = call_kwargs.kwargs.get("json") or call_kwargs.args[1] + assert sent_payload.get("model") == "gpt-4o" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py new file mode 100644 index 00000000000..642dd51b37b --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py @@ -0,0 +1,129 @@ +""" +Tests that when apply_guardrail returns structured_messages, both OpenAI and Anthropic +translation handlers write them back to data["messages"] correctly. + +For OpenAI: structured_messages (OpenAI format) written directly to data["messages"]. +For Anthropic: structured_messages (OpenAI format) converted back to Anthropic format + via anthropic_messages_pt before writing to data["messages"]. +""" + +from unittest.mock import MagicMock, patch + +import pytest + +from litellm.llms.openai.chat.guardrail_translation.handler import ( + OpenAIChatCompletionsHandler, +) +from litellm.types.utils import GenericGuardrailAPIInputs + +ORIGINAL_MESSAGES = [ + {"role": "system", "content": "Be concise."}, + {"role": "user", "content": "A" * 5000}, +] +COMPRESSED_MESSAGES = [ + {"role": "system", "content": "Be concise."}, + {"role": "user", "content": "A" * 200}, +] +COMPRESSED_MESSAGES_NON_SYSTEM = [ + {"role": "user", "content": "A" * 200}, +] + + +def _make_guardrail_returning_structured_messages(compressed: list) -> MagicMock: + guardrail = MagicMock() + guardrail.should_run_guardrail.return_value = True + guardrail.skip_system_message_in_guardrail = None + guardrail.skip_tool_message_in_guardrail = None + guardrail.experimental_use_latest_role_message_only = False + + async def apply_guardrail(inputs, request_data, input_type, logging_obj=None): + result = dict(inputs) + result["structured_messages"] = compressed + return result + + guardrail.apply_guardrail = apply_guardrail + return guardrail + + +@pytest.mark.asyncio +async def test_openai_handler_writes_structured_messages_back(): + handler = OpenAIChatCompletionsHandler() + guardrail = _make_guardrail_returning_structured_messages(COMPRESSED_MESSAGES) + + data = { + "model": "gpt-4o", + "messages": ORIGINAL_MESSAGES, + } + result = await handler.process_input_messages( + data=data, + guardrail_to_apply=guardrail, + ) + + assert result["messages"] == COMPRESSED_MESSAGES + + +@pytest.mark.asyncio +async def test_openai_handler_uses_text_patchback_when_no_structured_messages(): + handler = OpenAIChatCompletionsHandler() + + guardrail = MagicMock() + guardrail.should_run_guardrail.return_value = True + guardrail.skip_system_message_in_guardrail = None + guardrail.skip_tool_message_in_guardrail = None + guardrail.experimental_use_latest_role_message_only = False + + original_text = "hello world" + modified_text = "HELLO WORLD" + messages = [{"role": "user", "content": original_text}] + + async def apply_guardrail(inputs, request_data, input_type, logging_obj=None): + return GenericGuardrailAPIInputs(texts=[modified_text]) + + guardrail.apply_guardrail = apply_guardrail + + data = {"model": "gpt-4o", "messages": messages} + result = await handler.process_input_messages( + data=data, + guardrail_to_apply=guardrail, + ) + + assert result["messages"][0]["content"] == modified_text + + +@pytest.mark.asyncio +async def test_anthropic_handler_converts_structured_messages_to_anthropic_format(): + from litellm.llms.anthropic.chat.guardrail_translation.handler import ( + AnthropicMessagesHandler, + ) + + handler = AnthropicMessagesHandler() + guardrail = _make_guardrail_returning_structured_messages(COMPRESSED_MESSAGES) + + anthropic_messages = [ + {"role": "user", "content": [{"type": "text", "text": "A" * 5000}]} + ] + data = { + "model": "claude-3-5-sonnet-20241022", + "messages": anthropic_messages, + "max_tokens": 1024, + } + + converted_back = [ + {"role": "user", "content": [{"type": "text", "text": "A" * 200}]} + ] + + with patch( + "litellm.litellm_core_utils.prompt_templates.factory.anthropic_messages_pt", + return_value=converted_back, + ) as mock_pt: + result = await handler.process_input_messages( + data=data, + guardrail_to_apply=guardrail, + ) + + mock_pt.assert_called_once_with( + messages=COMPRESSED_MESSAGES_NON_SYSTEM, + model="claude-3-5-sonnet-20241022", + llm_provider="anthropic", + ) + assert result["messages"] == converted_back