From 761ab1920977e1def32beea58f19d7429940be08 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 18 May 2026 18:00:18 -0700 Subject: [PATCH] fix(bedrock): sanitize batch metadata to prevent Pydantic ValidationError (#28202) * fix(bedrock): sanitize batch metadata to prevent Pydantic ValidationError Proxy guardrail hooks (Model Armor, OpenAI Moderations) and internal processing inject non-string values (dicts, floats) into the request metadata. When the Bedrock batch handler passes this metadata directly to LiteLLMBatch (which inherits OpenAI's Batch Pydantic model with metadata: Dict[str, str]), Pydantic raises a ValidationError. This causes the router retry loop to re-submit the same Bedrock job multiple times before ultimately failing. Add _get_openai_compatible_batch_metadata() that serializes non-string values to JSON strings via safe_dumps, skips None values and internal logging keys, ensuring the response object always validates. * test(bedrock): add tests for batch metadata sanitization Covers _get_openai_compatible_batch_metadata: string passthrough, dict/float serialization, None/internal key exclusion, and LiteLLMBatch compatibility. --------- Co-authored-by: Noah Nistler <60981020+noahnistler@users.noreply.github.com> --- .../llms/bedrock/batches/transformation.py | 26 +++- .../test_batch_metadata_sanitization.py | 119 ++++++++++++++++++ 2 files changed, 144 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/llms/bedrock/batches/test_batch_metadata_sanitization.py diff --git a/litellm/llms/bedrock/batches/transformation.py b/litellm/llms/bedrock/batches/transformation.py index 0602b1c2f6..620bc91732 100644 --- a/litellm/llms/bedrock/batches/transformation.py +++ b/litellm/llms/bedrock/batches/transformation.py @@ -5,6 +5,7 @@ from typing import Any, Dict, List, Literal, Optional, Union, cast from httpx import Headers, Response +from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.llms.base_llm.batches.transformation import BaseBatchesConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.secret_managers.main import get_secret_str @@ -263,9 +264,32 @@ class BedrockBatchesConfig(BaseAWSLLM, BaseBatchesConfig): cancelling_at=None, cancelled_at=None, request_counts=None, - metadata=original_request.get("metadata", {}), + metadata=self._get_openai_compatible_batch_metadata( + original_request.get("metadata", {}) + ), ) + @staticmethod + def _get_openai_compatible_batch_metadata(metadata: Any) -> Dict[str, str]: + """ + OpenAI Batch metadata only accepts string values. + """ + if not isinstance(metadata, dict): + return {} + + sanitized_metadata: Dict[str, str] = {} + for key, value in metadata.items(): + if key == "standard_logging_guardrail_information" or value is None: + continue + + str_key = str(key) + if isinstance(value, str): + sanitized_metadata[str_key] = value + else: + sanitized_metadata[str_key] = safe_dumps(value) + + return sanitized_metadata + def transform_retrieve_batch_request( self, batch_id: str, diff --git a/tests/test_litellm/llms/bedrock/batches/test_batch_metadata_sanitization.py b/tests/test_litellm/llms/bedrock/batches/test_batch_metadata_sanitization.py new file mode 100644 index 0000000000..8de4733161 --- /dev/null +++ b/tests/test_litellm/llms/bedrock/batches/test_batch_metadata_sanitization.py @@ -0,0 +1,119 @@ +""" +Test that BedrockBatchesConfig._get_openai_compatible_batch_metadata +sanitizes non-string metadata values injected by proxy guardrail hooks. + +The OpenAI Batch Pydantic model requires metadata: Dict[str, str]. +Proxy hooks (Model Armor, OpenAI Moderations, queue time tracking) inject +dicts, floats, and other non-string values that cause a ValidationError +when constructing LiteLLMBatch. This test suite verifies the sanitization +layer prevents that. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.llms.bedrock.batches.transformation import BedrockBatchesConfig + + +class TestGetOpenaiCompatibleBatchMetadata: + """Tests for _get_openai_compatible_batch_metadata.""" + + def test_string_values_pass_through_unchanged(self): + metadata = {"user_key": "user_value", "run_id": "abc123"} + result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata) + assert result == {"user_key": "user_value", "run_id": "abc123"} + + def test_dict_values_serialized_to_json_string(self): + metadata = { + "_model_armor_response": { + "sanitizationResult": {"filterMatchState": "MATCH_FOUND"} + } + } + result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata) + assert "_model_armor_response" in result + assert isinstance(result["_model_armor_response"], str) + assert "MATCH_FOUND" in result["_model_armor_response"] + + def test_float_values_serialized_to_string(self): + metadata = {"queue_time_seconds": 0.5} + result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata) + assert result == {"queue_time_seconds": "0.5"} + + def test_none_values_excluded(self): + metadata = {"key": "value", "empty": None} + result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata) + assert "empty" not in result + assert result == {"key": "value"} + + def test_standard_logging_guardrail_information_excluded(self): + metadata = { + "standard_logging_guardrail_information": {"some": "logging_data"}, + "user_key": "keep_me", + } + result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata) + assert "standard_logging_guardrail_information" not in result + assert result == {"user_key": "keep_me"} + + def test_non_dict_input_returns_empty_dict(self): + assert BedrockBatchesConfig._get_openai_compatible_batch_metadata(None) == {} + assert BedrockBatchesConfig._get_openai_compatible_batch_metadata("string") == {} + assert BedrockBatchesConfig._get_openai_compatible_batch_metadata(123) == {} + + def test_empty_dict_returns_empty_dict(self): + assert BedrockBatchesConfig._get_openai_compatible_batch_metadata({}) == {} + + def test_mixed_metadata_from_guardrails(self): + """Simulate real metadata contaminated by proxy guardrails.""" + metadata = { + "_model_armor_response": {"sanitizationResult": {"key": "val"}}, + "_model_armor_status": "success", + "_openai_moderation_response": {"id": "mod-123", "flagged": False}, + "queue_time_seconds": 1.23, + "headers": {"Authorization": "Bearer sk-xxx"}, + "standard_logging_guardrail_information": {"internal": True}, + "user_metadata_key": "user_value", + "none_field": None, + } + result = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata) + + # All values must be strings + for key, value in result.items(): + assert isinstance(value, str), f"metadata[{key!r}] is {type(value)}, not str" + + # Excluded keys + assert "standard_logging_guardrail_information" not in result + assert "none_field" not in result + + # Preserved keys + assert result["_model_armor_status"] == "success" + assert result["user_metadata_key"] == "user_value" + + def test_result_compatible_with_litellm_batch(self): + """Verify sanitized metadata can construct a LiteLLMBatch without error.""" + import time + + from litellm.types.utils import LiteLLMBatch + + metadata = { + "_model_armor_response": {"blocked": True}, + "queue_time_seconds": 0.05, + "user_key": "value", + } + sanitized = BedrockBatchesConfig._get_openai_compatible_batch_metadata(metadata) + + # This would raise ValidationError before the fix + batch = LiteLLMBatch( + id="arn:aws:bedrock:us-east-1:123:model-invocation-job/test", + object="batch", + endpoint="/v1/chat/completions", + input_file_id="file-123", + completion_window="24h", + status="validating", + created_at=int(time.time()), + metadata=sanitized, + ) + assert batch.metadata == sanitized