mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
fix(guardrails/anthropic): apply compressed structured_messages on /v1/messages and fix guardrail_information logging
Two bugs fixed for /v1/messages (Anthropic native endpoint) with headroom prompt compression guardrail: 1. guardrail_information null in spend logs: function_setup stored a .copy() of litellm_metadata as litellm_params["metadata"]. Guardrails write to litellm_metadata after this point; the copy didn't reflect those mutations. Changed to use a reference so guardrail writes are visible at logging time. 2. Compressed structured_messages not applied: AnthropicMessagesHandler.process_input_messages sent OpenAI-format structured_messages to the guardrail but never read back the compressed result. Added identity-check (is not) to detect when guardrail returns different structured_messages and convert them back to Anthropic format via anthropic_messages_pt before sending the request. Uses identity check (not equality) to distinguish pass-through guardrails (which echo the original structured_messages unchanged) from compression guardrails (which return a new list object with fewer messages).
This commit is contained in:
Krrish Dholakia
parent
9d32b4081b
commit
498b8d85cd
3 changed files with 640 additions and 1544 deletions
|
|
@ -125,9 +125,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
texts_to_check: List[str] = []
|
||||
images_to_check: List[str] = []
|
||||
tools_to_check: List[ChatCompletionToolParam] = (
|
||||
chat_completion_compatible_request.get("tools", [])
|
||||
)
|
||||
tools_to_check: List[ChatCompletionToolParam] = chat_completion_compatible_request.get("tools", [])
|
||||
task_mappings: List[Tuple[int, Optional[int]]] = []
|
||||
|
||||
# Step 1: Extract all text content and images
|
||||
|
|
@ -164,6 +162,8 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
guardrailed_texts = guardrailed_inputs.get("texts", [])
|
||||
guardrailed_tools = guardrailed_inputs.get("tools")
|
||||
guardrailed_structured_messages = guardrailed_inputs.get("structured_messages")
|
||||
|
||||
if guardrailed_tools is not None:
|
||||
# Convert tools back from OpenAI format to Anthropic format
|
||||
anthropic_config = AnthropicConfig()
|
||||
|
|
@ -172,19 +172,35 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
converted_tool, mcp_server = anthropic_config._map_tool_helper(tool)
|
||||
if converted_tool is not None:
|
||||
anthropic_tools.append(converted_tool)
|
||||
# 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,
|
||||
)
|
||||
# If guardrail returned modified structured_messages (OpenAI format), convert
|
||||
# them back to Anthropic format and replace the request messages wholesale.
|
||||
# Only do this when the returned messages differ from what we sent — a
|
||||
# pass-through guardrail echoes the original structured_messages unchanged.
|
||||
if (
|
||||
guardrailed_structured_messages is not None
|
||||
and guardrailed_structured_messages is not structured_messages
|
||||
):
|
||||
from litellm.litellm_core_utils.prompt_templates.factory import (
|
||||
anthropic_messages_pt,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Anthropic Messages: Processed input messages: %s", messages
|
||||
)
|
||||
model = data.get("model", "")
|
||||
data["messages"] = anthropic_messages_pt(
|
||||
messages=guardrailed_structured_messages,
|
||||
model=model,
|
||||
llm_provider="anthropic",
|
||||
)
|
||||
else:
|
||||
# Step 3: Map guardrail responses back to original message structure
|
||||
await self._apply_guardrail_responses_to_input(
|
||||
messages=messages,
|
||||
responses=guardrailed_texts,
|
||||
task_mappings=task_mappings,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug("Anthropic Messages: Processed input messages: %s", messages)
|
||||
|
||||
return data
|
||||
|
||||
|
|
@ -288,9 +304,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
|
||||
elif isinstance(content, list) and content_idx_optional is not None:
|
||||
# Replace specific text item in list content
|
||||
messages[msg_idx]["content"][content_idx_optional]["text"] = (
|
||||
guardrail_response
|
||||
)
|
||||
messages[msg_idx]["content"][content_idx_optional]["text"] = guardrail_response
|
||||
|
||||
async def process_output_response(
|
||||
self,
|
||||
|
|
@ -369,9 +383,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
task_mappings=task_mappings,
|
||||
)
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
"Anthropic Messages: Processed output response: %s", response
|
||||
)
|
||||
verbose_proxy_logger.debug("Anthropic Messages: Processed output response: %s", response)
|
||||
|
||||
return response
|
||||
|
||||
|
|
@ -391,20 +403,14 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
has_ended = self._check_streaming_has_ended(responses_so_far)
|
||||
if has_ended:
|
||||
# build the model response from the responses_so_far
|
||||
built_response = (
|
||||
AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
|
||||
all_chunks=responses_so_far,
|
||||
litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj),
|
||||
model="",
|
||||
)
|
||||
built_response = AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
|
||||
all_chunks=responses_so_far,
|
||||
litellm_logging_obj=cast("LiteLLMLoggingObj", litellm_logging_obj),
|
||||
model="",
|
||||
)
|
||||
|
||||
# Check if model_response is valid and has choices before accessing
|
||||
if (
|
||||
built_response is not None
|
||||
and hasattr(built_response, "choices")
|
||||
and built_response.choices
|
||||
):
|
||||
if built_response is not None and hasattr(built_response, "choices") and built_response.choices:
|
||||
model_response = cast(ModelResponse, built_response)
|
||||
first_choice = cast(Choices, model_response.choices[0])
|
||||
tool_calls_list = cast(
|
||||
|
|
@ -418,16 +424,16 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
if tool_calls_list:
|
||||
guardrail_inputs["tool_calls"] = tool_calls_list
|
||||
|
||||
_guardrailed_inputs = await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid
|
||||
inputs=guardrail_inputs,
|
||||
request_data=request_data if request_data is not None else {},
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
_guardrailed_inputs = (
|
||||
await guardrail_to_apply.apply_guardrail( # allow rejecting the response, if invalid
|
||||
inputs=guardrail_inputs,
|
||||
request_data=request_data if request_data is not None else {},
|
||||
input_type="response",
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
)
|
||||
else:
|
||||
verbose_proxy_logger.debug(
|
||||
"Skipping output guardrail - model response has no choices"
|
||||
)
|
||||
verbose_proxy_logger.debug("Skipping output guardrail - model response has no choices")
|
||||
return responses_so_far
|
||||
|
||||
string_so_far = self.get_streaming_string_so_far(responses_so_far)
|
||||
|
|
@ -454,9 +460,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
request_data[key] = response
|
||||
|
||||
if "litellm_metadata" not in request_data:
|
||||
user_metadata = self.transform_user_api_key_dict_to_metadata(
|
||||
user_api_key_dict
|
||||
)
|
||||
user_metadata = self.transform_user_api_key_dict_to_metadata(user_api_key_dict)
|
||||
if user_metadata:
|
||||
request_data["litellm_metadata"] = user_metadata
|
||||
return request_data
|
||||
|
|
@ -604,9 +608,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
if delta.get("type") == "text_delta":
|
||||
text += delta.get("text", "")
|
||||
except json.JSONDecodeError:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Failed to parse JSON from SSE data: {data_line}"
|
||||
)
|
||||
verbose_proxy_logger.warning(f"Failed to parse JSON from SSE data: {data_line}")
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"Error extracting text from SSE: {e}")
|
||||
|
|
@ -670,14 +672,10 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
if stop_reason is not None:
|
||||
return True
|
||||
except json.JSONDecodeError:
|
||||
verbose_proxy_logger.warning(
|
||||
f"Failed to parse JSON from SSE data: {data_line}"
|
||||
)
|
||||
verbose_proxy_logger.warning(f"Failed to parse JSON from SSE data: {data_line}")
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(
|
||||
f"Error checking streaming end in SSE: {e}"
|
||||
)
|
||||
verbose_proxy_logger.error(f"Error checking streaming end in SSE: {e}")
|
||||
|
||||
# Handle already-parsed dict format
|
||||
elif isinstance(response, dict):
|
||||
|
|
@ -783,10 +781,7 @@ class AnthropicMessagesHandler(BaseTranslation):
|
|||
if isinstance(content_block, dict):
|
||||
if content_block.get("type") == "text":
|
||||
cast(Dict[str, Any], content_block)["text"] = guardrail_response
|
||||
elif (
|
||||
hasattr(content_block, "type")
|
||||
and getattr(content_block, "type", None) == "text"
|
||||
):
|
||||
elif hasattr(content_block, "type") and getattr(content_block, "type", None) == "text":
|
||||
# Update Pydantic object's text attribute
|
||||
if hasattr(content_block, "text"):
|
||||
content_block.text = guardrail_response
|
||||
|
|
|
|||
1939
litellm/utils.py
1939
litellm/utils.py
File diff suppressed because it is too large
Load diff
|
|
@ -12,9 +12,7 @@ from unittest.mock import MagicMock, patch
|
|||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../../../../../../..")
|
||||
) # Adds the parent directory to the system path
|
||||
sys.path.insert(0, os.path.abspath("../../../../../../..")) # Adds the parent directory to the system path
|
||||
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.llms.anthropic.chat.guardrail_translation.handler import (
|
||||
|
|
@ -51,9 +49,7 @@ class MockDynamicGuardrail(CustomGuardrail):
|
|||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
self.dynamic_params = self.get_guardrail_dynamic_request_body_params(
|
||||
request_data
|
||||
)
|
||||
self.dynamic_params = self.get_guardrail_dynamic_request_body_params(request_data)
|
||||
return inputs
|
||||
|
||||
|
||||
|
|
@ -103,17 +99,11 @@ class TestAnthropicMessagesHandlerInputProcessing:
|
|||
data = {
|
||||
"model": "claude-3-5-sonnet-20241022",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"litellm_metadata": {
|
||||
"guardrails": [
|
||||
{"cygnal-monitor": {"extra_body": {"policy_id": "policy-123"}}}
|
||||
]
|
||||
},
|
||||
"litellm_metadata": {"guardrails": [{"cygnal-monitor": {"extra_body": {"policy_id": "policy-123"}}}]},
|
||||
}
|
||||
|
||||
with patch("litellm.proxy.proxy_server.premium_user", True):
|
||||
await handler.process_input_messages(
|
||||
data=data, guardrail_to_apply=guardrail
|
||||
)
|
||||
await handler.process_input_messages(data=data, guardrail_to_apply=guardrail)
|
||||
|
||||
assert data.get("litellm_metadata", {}).get("guardrails")
|
||||
assert guardrail.dynamic_params == {"policy_id": "policy-123"}
|
||||
|
|
@ -216,9 +206,7 @@ class TestAnthropicMessagesHandlerInputProcessing:
|
|||
# Mock _check_streaming_has_ended to return False (stream not ended)
|
||||
with (
|
||||
patch.object(handler, "_check_streaming_has_ended", return_value=False),
|
||||
patch.object(
|
||||
handler, "get_streaming_string_so_far", return_value="partial text"
|
||||
),
|
||||
patch.object(handler, "get_streaming_string_so_far", return_value="partial text"),
|
||||
):
|
||||
responses_so_far = [b"data: some chunk"]
|
||||
|
||||
|
|
@ -249,9 +237,7 @@ class TestAnthropicMessagesHandlerInputProcessing:
|
|||
|
||||
data = {
|
||||
"model": "claude-opus-4-6",
|
||||
"messages": [
|
||||
{"role": "user", "content": "What is the weather in San Francisco?"}
|
||||
],
|
||||
"messages": [{"role": "user", "content": "What is the weather in San Francisco?"}],
|
||||
"tools": [
|
||||
{
|
||||
"type": "tool_search_tool_regex_20251119",
|
||||
|
|
@ -295,6 +281,122 @@ class TestAnthropicMessagesHandlerInputProcessing:
|
|||
assert "input_schema" in tools[1]
|
||||
|
||||
|
||||
class MockStructuredMessagesGuardrail(CustomGuardrail):
|
||||
"""Mock guardrail that returns compressed structured_messages."""
|
||||
|
||||
def __init__(self, guardrail_name: str, compressed_messages: List[Any]):
|
||||
super().__init__(guardrail_name=guardrail_name)
|
||||
self.compressed_messages = compressed_messages
|
||||
|
||||
async def apply_guardrail(
|
||||
self,
|
||||
inputs: GenericGuardrailAPIInputs,
|
||||
request_data: dict,
|
||||
input_type: Literal["request", "response"],
|
||||
logging_obj: Optional[Any] = None,
|
||||
) -> GenericGuardrailAPIInputs:
|
||||
result = dict(inputs)
|
||||
result["structured_messages"] = self.compressed_messages
|
||||
return GenericGuardrailAPIInputs(**result)
|
||||
|
||||
|
||||
class TestAnthropicStructuredMessagesApplied:
|
||||
"""Test that structured_messages from guardrail response are applied back in Anthropic format."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_structured_messages_replaced_in_anthropic_format(self):
|
||||
"""When guardrail returns structured_messages (OpenAI format), they must be
|
||||
converted back to Anthropic format and written to data['messages']."""
|
||||
handler = AnthropicMessagesHandler()
|
||||
|
||||
compressed_openai_messages = [
|
||||
{"role": "user", "content": "compressed content"},
|
||||
]
|
||||
guardrail = MockStructuredMessagesGuardrail(
|
||||
guardrail_name="test-compressor",
|
||||
compressed_messages=compressed_openai_messages,
|
||||
)
|
||||
|
||||
data: dict = {
|
||||
"model": "claude-3-5-sonnet-20241022",
|
||||
"messages": [
|
||||
{"role": "user", "content": "original long content that gets compressed"},
|
||||
],
|
||||
}
|
||||
|
||||
result = await handler.process_input_messages(
|
||||
data=data, guardrail_to_apply=guardrail, litellm_logging_obj=MagicMock()
|
||||
)
|
||||
|
||||
messages = result["messages"]
|
||||
assert len(messages) == 1
|
||||
assert messages[0]["role"] == "user"
|
||||
content = messages[0]["content"]
|
||||
if isinstance(content, str):
|
||||
assert content == "compressed content"
|
||||
elif isinstance(content, list):
|
||||
assert any(block.get("text") == "compressed content" for block in content if isinstance(block, dict))
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_structured_messages_falls_back_to_text_patching(self):
|
||||
"""When guardrail returns no structured_messages, the original text-patch path runs."""
|
||||
handler = AnthropicMessagesHandler()
|
||||
guardrail = MockPassThroughGuardrail(guardrail_name="test")
|
||||
|
||||
original_content = "original content"
|
||||
data: dict = {
|
||||
"model": "claude-3-5-sonnet-20241022",
|
||||
"messages": [{"role": "user", "content": original_content}],
|
||||
}
|
||||
|
||||
result = await handler.process_input_messages(
|
||||
data=data, guardrail_to_apply=guardrail, litellm_logging_obj=MagicMock()
|
||||
)
|
||||
|
||||
messages = result["messages"]
|
||||
assert len(messages) == 1
|
||||
content = messages[0]["content"]
|
||||
if isinstance(content, str):
|
||||
assert content == original_content
|
||||
|
||||
|
||||
class TestFunctionSetupLitellmMetadataReference:
|
||||
"""Regression test: guardrail info written to litellm_metadata must appear in
|
||||
standard logging payload for endpoints that use litellm_metadata (e.g. /v1/messages)."""
|
||||
|
||||
def test_litellm_metadata_reference_not_copy(self):
|
||||
"""litellm_params['metadata'] must be the same object as kwargs['litellm_metadata']
|
||||
so mutations by guardrails after function_setup are visible at logging time."""
|
||||
import litellm
|
||||
from datetime import datetime
|
||||
|
||||
litellm_metadata_dict: dict = {"user_api_key": "test-key"}
|
||||
kwargs = {
|
||||
"model": "claude-3-5-sonnet-20241022",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"litellm_call_id": "test-call-id",
|
||||
"litellm_metadata": litellm_metadata_dict,
|
||||
}
|
||||
|
||||
logging_obj, _ = litellm.utils.function_setup(
|
||||
original_function="anthropic_messages",
|
||||
rules_obj=litellm.utils.Rules(),
|
||||
start_time=datetime.now(),
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
litellm_metadata_dict["standard_logging_guardrail_information"] = [
|
||||
{"guardrail_name": "test-guardrail", "guardrail_status": "success"}
|
||||
]
|
||||
|
||||
metadata_in_logging = logging_obj.litellm_params.get("metadata", {})
|
||||
assert "standard_logging_guardrail_information" in metadata_in_logging, (
|
||||
"guardrail info written to litellm_metadata after function_setup must be "
|
||||
"visible in litellm_params['metadata'] — check function_setup uses a "
|
||||
"reference not a copy for litellm_metadata endpoints"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Run the tests
|
||||
pytest.main([__file__, "-v"])
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue