scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 79

fix(model_armor): post-call response silently dropped; embeddings/images show phantom applied entries

Browse source 
Two bugs in ModelArmorGuardrail:

1. async_post_call_success_hook wrote the armor response via
   add_guardrail_response_to_standard_logging_object, which targets
   litellm_logging_obj.model_call_details["standard_logging_object"].
   The standard logging payload reads from
   metadata["standard_logging_guardrail_information"] — a different path —
   so post-call guardrail responses were silently dropped from the UI.
   Fixed by storing _model_armor_response / _model_armor_status in
   metadata (same pattern as pre-call and during-call hooks) so the
   @log_guardrail_information decorator's _process_response picks them up.

2. _process_response defaulted _model_armor_response to {} when the key
   was absent, causing it to emit a phantom "applied / success / 0 ms"
   guardrail entry for every embedding (no messages) and image (non-
   ModelResponse) request. Fixed by skipping logging when
   _model_armor_response is not present in metadata.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
RoyVivat 2026-04-06 14:59:01 -07:00
parent 50f88c8642
commit 00cd1a2be0
No known key found for this signature in database
GPG key ID: 59743472EC86530E
2 changed files with 354 additions and 29 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

52
litellm/proxy/guardrails/guardrail_hooks/model_armor/model_armor.py
View file

@ -34,7 +34,6 @@ from litellm.types.utils import (
GuardrailStatus,
ModelResponse,
ModelResponseStream,
StandardLoggingGuardrailInformation,
)
GUARDRAIL_NAME = "model_armor"
@ -341,7 +340,15 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
request_data.get("metadata", {}) if isinstance(request_data, dict) else {}
)
guardrail_response = metadata.get("_model_armor_response", {})
# Only log when an actual Model Armor API call was made. When there is
# no content to scan (e.g. embeddings have no messages, image responses
# are not text) the hooks return early *without* setting
# `_model_armor_response`, so we skip logging here to avoid emitting
# phantom "applied / success / 0 ms latency" guardrail entries.
if "_model_armor_response" not in metadata:
return response
guardrail_response = metadata["_model_armor_response"]
# Determine status – default to "success" but prefer the explicit value if present.
guardrail_status: GuardrailStatus = metadata.get(
@ -570,7 +577,6 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
):
"""Post-call hook to sanitize model responses."""
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_response_to_standard_logging_object,
add_guardrail_to_applied_guardrails_header,
)
@ -598,32 +604,22 @@ class ModelArmorGuardrail(CustomGuardrail, VertexBase):
request_data=data,
)
# Attach Model Armor response & status to this request's metadata to prevent race conditions
if isinstance(armor_response, dict):
model_armor_logged_object = {
"model_armor_response": armor_response,
"model_armor_status": (
"blocked"
if self._should_block_content(
armor_response,
allow_sanitization=self.mask_response_content,
)
else "success"
),
}
standard_logging_guardrail_information = (
StandardLoggingGuardrailInformation(
guardrail_name=self.guardrail_name,
guardrail_provider="model_armor",
guardrail_mode=GuardrailEventHooks.post_call,
guardrail_response=model_armor_logged_object,
guardrail_status="success",
start_time=data.get("start_time"),
# Store armor response in metadata so the decorator's _process_response
# picks it up via the same path as async_pre_call_hook / async_moderation_hook.
# Previously this wrote to litellm_logging_obj.model_call_details["standard_logging_object"]
# which is a different code path than what the standard logging payload reads
# (metadata["standard_logging_guardrail_information"]), causing post-call
# guardrail responses to be silently dropped from the UI / logs.
if isinstance(data, dict):
metadata = data.setdefault("metadata", {})
metadata["_model_armor_response"] = armor_response
metadata["_model_armor_status"] = (
"blocked"
if self._should_block_content(
armor_response,
allow_sanitization=self.mask_response_content,
)
)
add_guardrail_response_to_standard_logging_object(
litellm_logging_obj=data.get("litellm_logging_obj"),
guardrail_response=standard_logging_guardrail_information,
else "success"
)
# Add guardrail to applied_guardrails BEFORE potential blocking

331
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py
View file

@ -1605,4 +1605,333 @@ async def test_async_moderation_hook_api_error_fail_on_error_false():
call_type="completion"
)
assert "API Error" in str(exc_info.value)
# ---------------------------------------------------------------------------
# Regression tests for post-call logging and false-applied-entry bugs
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_post_call_stores_armor_response_in_metadata():
"""
Regression: async_post_call_success_hook must store _model_armor_response and
_model_armor_status in metadata so the @log_guardrail_information decorator's
_process_response picks them up via metadata["standard_logging_guardrail_information"].
Previously the hook used add_guardrail_response_to_standard_logging_object, which
writes to litellm_logging_obj.model_call_details["standard_logging_object"] – a
different code path that the standard logging payload never reads, causing post-call
guardrail responses to be silently dropped.
"""
mock_user_api_key_dict = UserAPIKeyAuth()
guardrail = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-post-logging",
)
armor_api_response = {
"sanitizationResult": {
"filterMatchState": "NO_MATCH_FOUND",
"filterResults": {},
}
}
mock_http_response = AsyncMock()
mock_http_response.status_code = 200
mock_http_response.json = AsyncMock(return_value=armor_api_response)
guardrail._ensure_access_token_async = AsyncMock(
return_value=("test-token", "test-project")
)
mock_llm_response = litellm.ModelResponse()
mock_llm_response.choices = [
litellm.Choices(message=litellm.Message(content="Hello there"))
]
request_data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hi"}],
"metadata": {"guardrails": ["model-armor-post-logging"]},
}
with patch.object(
guardrail.async_handler, "post", AsyncMock(return_value=mock_http_response)
):
await guardrail.async_post_call_success_hook(
data=request_data,
user_api_key_dict=mock_user_api_key_dict,
response=mock_llm_response,
)
metadata = request_data["metadata"]
# _model_armor_response must be set so _process_response can read it
assert "_model_armor_response" in metadata, (
"_model_armor_response not found in metadata – post-call logging will be dropped"
)
assert metadata["_model_armor_response"] == armor_api_response
# Status must be "success" for a clean response
assert metadata.get("_model_armor_status") == "success"
# standard_logging_guardrail_information must contain one entry for this hook
slgi = metadata.get("standard_logging_guardrail_information")
assert slgi is not None, (
"standard_logging_guardrail_information not populated – post-call entry dropped"
)
assert len(slgi) >= 1
entry = slgi[0]
assert entry["guardrail_name"] == "model-armor-post-logging"
assert entry["guardrail_status"] == "success"
@pytest.mark.asyncio
async def test_post_call_blocked_response_sets_metadata_status_blocked():
"""
Regression: when a post-call armor response blocks content, _model_armor_status
in metadata must be 'blocked' so the logging decorator records it correctly.
"""
mock_user_api_key_dict = UserAPIKeyAuth()
guardrail = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-post-blocked",
)
armor_api_response = {
"sanitizationResult": {
"filterMatchState": "MATCH_FOUND",
"filterResults": {
"rai": {
"raiFilterResult": {
"matchState": "MATCH_FOUND",
}
}
},
}
}
mock_http_response = AsyncMock()
mock_http_response.status_code = 200
mock_http_response.json = AsyncMock(return_value=armor_api_response)
guardrail._ensure_access_token_async = AsyncMock(
return_value=("test-token", "test-project")
)
mock_llm_response = litellm.ModelResponse()
mock_llm_response.choices = [
litellm.Choices(message=litellm.Message(content="Some harmful content"))
]
request_data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Prompt"}],
"metadata": {"guardrails": ["model-armor-post-blocked"]},
}
with patch.object(
guardrail.async_handler, "post", AsyncMock(return_value=mock_http_response)
):
with pytest.raises(HTTPException) as exc_info:
await guardrail.async_post_call_success_hook(
data=request_data,
user_api_key_dict=mock_user_api_key_dict,
response=mock_llm_response,
)
assert exc_info.value.status_code == 400
# Even though the request was blocked, the metadata must still be populated
metadata = request_data["metadata"]
assert metadata.get("_model_armor_status") == "blocked"
assert metadata.get("_model_armor_response") == armor_api_response
# The decorator logs with guardrail_intervened status on HTTPException
slgi = metadata.get("standard_logging_guardrail_information")
assert slgi is not None
assert len(slgi) >= 1
@pytest.mark.asyncio
async def test_pre_call_embedding_request_produces_no_guardrail_log_entry():
"""
Regression: when async_pre_call_hook is called for an embedding request (which has
no 'messages' field), it should return early WITHOUT emitting a
standard_logging_guardrail_information entry.
Previously the @log_guardrail_information decorator called _process_response even
after early returns, and _process_response defaulted _model_armor_response to {}
and _model_armor_status to "success", producing a phantom "applied / 0 ms" entry
in the UI for every embedding call.
"""
mock_user_api_key_dict = UserAPIKeyAuth()
mock_cache = MagicMock(spec=DualCache)
guardrail = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-embed",
)
guardrail._ensure_access_token_async = AsyncMock(
return_value=("test-token", "test-project")
)
# Embedding requests have 'input', not 'messages'
request_data = {
"model": "text-embedding-ada-002",
"input": "Hello world",
"metadata": {"guardrails": ["model-armor-embed"]},
}
with patch.object(
guardrail.async_handler, "post", AsyncMock()
) as mock_post:
await guardrail.async_pre_call_hook(
user_api_key_dict=mock_user_api_key_dict,
cache=mock_cache,
data=request_data,
call_type="embedding",
)
# No Model Armor API call must have been made
mock_post.assert_not_called()
# No guardrail log entry must be emitted – this is the phantom entry bug
slgi = request_data.get("metadata", {}).get("standard_logging_guardrail_information")
assert slgi is None, (
f"Phantom guardrail log entry emitted for embedding pre-call: {slgi}"
)
@pytest.mark.asyncio
async def test_post_call_image_response_produces_no_guardrail_log_entry():
"""
Regression: when async_post_call_success_hook is called with an ImageResponse
(non-ModelResponse), it should return early WITHOUT emitting a
standard_logging_guardrail_information entry.
Previously the decorator would emit an "applied / 0 ms / empty response" entry
for every image generation call, making the UI report all 6 guardrails as applied
with no actual Model Armor response data.
"""
mock_user_api_key_dict = UserAPIKeyAuth()
guardrail = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-image",
)
guardrail._ensure_access_token_async = AsyncMock(
return_value=("test-token", "test-project")
)
# ImageResponse is not a litellm.ModelResponse – the hook must skip it
from litellm.types.utils import ImageResponse as LiteLLMImageResponse
from openai.types import ImagesResponse
from openai.types.image import Image
image_response = LiteLLMImageResponse(
created=1234567890,
data=[Image(url="https://example.com/image.png")],
)
request_data = {
"model": "dall-e-3",
"metadata": {"guardrails": ["model-armor-image"]},
}
with patch.object(
guardrail.async_handler, "post", AsyncMock()
) as mock_post:
await guardrail.async_post_call_success_hook(
data=request_data,
user_api_key_dict=mock_user_api_key_dict,
response=image_response,
)
# No Model Armor API call must have been made for image responses
mock_post.assert_not_called()
# No guardrail log entry must be emitted
slgi = request_data.get("metadata", {}).get("standard_logging_guardrail_information")
assert slgi is None, (
f"Phantom guardrail log entry emitted for image post-call: {slgi}"
)
def test_process_response_skips_logging_when_no_armor_api_call():
"""
Unit test for _process_response: it must return without emitting
standard_logging_guardrail_information when _model_armor_response is absent
from metadata (i.e. no Model Armor API call was made).
"""
guardrail = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-unit",
)
request_data_without_call = {
"model": "gpt-4",
"metadata": {}, # _model_armor_response intentionally absent
}
guardrail._process_response(
response=None,
request_data=request_data_without_call,
duration=0.001,
event_type=GuardrailEventHooks.pre_call,
)
# No entry should be added
assert "standard_logging_guardrail_information" not in request_data_without_call["metadata"]
def test_process_response_logs_when_armor_api_call_made():
"""
Unit test for _process_response: it must emit standard_logging_guardrail_information
when _model_armor_response IS present in metadata (an actual API call was made).
"""
guardrail = ModelArmorGuardrail(
template_id="test-template",
project_id="test-project",
location="us-central1",
guardrail_name="model-armor-unit",
)
armor_response = {"sanitizationResult": {"filterMatchState": "NO_MATCH_FOUND"}}
request_data_with_call = {
"model": "gpt-4",
"metadata": {
"_model_armor_response": armor_response,
"_model_armor_status": "success",
},
}
guardrail._process_response(
response=None,
request_data=request_data_with_call,
duration=0.142,
event_type=GuardrailEventHooks.pre_call,
)
slgi = request_data_with_call["metadata"].get("standard_logging_guardrail_information")
assert slgi is not None, "_process_response did not emit guardrail log when API call was made"
assert len(slgi) == 1
entry = slgi[0]
assert entry["guardrail_name"] == "model-armor-unit"
assert entry["guardrail_status"] == "success"
assert entry["guardrail_response"] == armor_response