scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

feat(guardrails): add headroom guardrail for message compression (#31407)

Browse source 
* 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
This commit is contained in:
Krrish Dholakia 2026-06-26 19:36:44 -07:00 • committed by GitHub
parent b9765458ac
commit 99b1a323c1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 998 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

45
litellm/integrations/custom_guardrail.py
View file

@ -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):

1
litellm/litellm_core_utils/litellm_logging.py
View file

@ -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

42
litellm/llms/anthropic/chat/guardrail_translation/handler.py
View file

@ -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] = []

43
litellm/llms/openai/chat/guardrail_translation/handler.py
View file

@ -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

52
litellm/proxy/guardrails/guardrail_hooks/headroom/__init__.py Normal file
View file

@ -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,
}

210
litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py Normal file
View file

@ -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

5
litellm/types/guardrails.py
View file

@ -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,

24
litellm/types/proxy/guardrails/guardrail_hooks/headroom.py Normal file
View file

@ -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"

124
tests/test_litellm/integrations/test_guardrail_logging_sync.py Normal file
View file

@ -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]

346
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_headroom.py Normal file
View file

@ -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"

129
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_structured_messages_writeback.py Normal file
View file

@ -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