From e50ae35ddec982ba51c3d3b5580caa38287ed4a3 Mon Sep 17 00:00:00 2001 From: ryan Date: Wed, 28 Oct 2026 08:06:56 +0000 Subject: [PATCH] test: cover JSON mode preamble text preservation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_bedrock_native_e2e.py | 94 ++++++++++++++++ .../test_anthropic_chat_transformation.py | 63 +++++++++++ .../chat/test_converse_transformation.py | 102 ++++++++++++++++++ 3 files changed, 259 insertions(+) diff --git a/tests/e2e/llm_translation/test_bedrock_native_e2e.py b/tests/e2e/llm_translation/test_bedrock_native_e2e.py index 23ff6ac612e..425d8b9b74a 100644 --- a/tests/e2e/llm_translation/test_bedrock_native_e2e.py +++ b/tests/e2e/llm_translation/test_bedrock_native_e2e.py @@ -9,6 +9,8 @@ from __future__ import annotations import pytest from typing import Final +from openai.types.chat import ChatCompletionMessage +from openai.types.shared_params.response_format_json_schema import ResponseFormatJSONSchema from pydantic import BaseModel, ConfigDict from e2e_config import unique_marker @@ -86,6 +88,55 @@ class InvokeResponse(BaseModel): content: list[ConverseContent] +class WarehouseAnswer(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid", strict=True) + + answer: int + + +class JsonModePreambleFields(BaseModel): + model_config = ConfigDict(frozen=True) + + json_mode_preamble_text: str | None = None + + +class JsonModeAssistantMessage(BaseModel): + model_config = ConfigDict(frozen=True) + + content: str | None = None + provider_specific_fields: JsonModePreambleFields | None = None + + +BEDROCK_JSON_SCHEMA_BACKEND: Final = "bedrock/us.anthropic.claude-opus-5-5" +WAREHOUSE_PROMPT: Final = ( + "A warehouse has 3 shelves. Shelf A holds 17 boxes of 24 items, shelf B holds 9 boxes of 36 items minus 13 " + "damaged items, and shelf C holds twice as many items as shelf B. How many usable items are there in total? " + "Work through it carefully, then answer." +) +WAREHOUSE_RESPONSE_FORMAT: Final[ResponseFormatJSONSchema] = { + "type": "json_schema", + "json_schema": { + "name": "warehouse_answer", + "strict": True, + "schema": { + "type": "object", + "properties": {"answer": {"type": "integer"}}, + "required": ["answer"], + "additionalProperties": False, + }, + }, +} + + +def _warehouse_message_has_preamble(message: ChatCompletionMessage) -> bool: + structured_message: Final = JsonModeAssistantMessage.model_validate(message.model_dump()) + assert structured_message.content is not None + WarehouseAnswer.model_validate_json(structured_message.content) + provider_fields: Final = structured_message.provider_specific_fields + preamble: Final = provider_fields.json_mode_preamble_text if provider_fields is not None else None + return preamble is not None and bool(preamble.strip()) + + def _register(proxy: ProxyClient, resources: ResourceManager) -> tuple[str, str]: model = f"e2e-bedrock-native-{unique_marker()}" model_id = proxy.create_model( @@ -144,6 +195,49 @@ class TestBedrockNative: assert parsed.events assert all(event.name and event.date and event.participants for event in parsed.events) + @pytest.mark.provider_live + @pytest.mark.covers("llm.chat_completions.bedrock_converse.structured_output.preamble.nonstream.works") + @meta( + Subject( + domain=Domain.LLM_TRANSLATION, + route=Route.CHAT_COMPLETIONS, + providers=(Provider.BEDROCK,), + models=(BEDROCK_JSON_SCHEMA_BACKEND,), + capabilities=(Capability.RESPONSE_SCHEMA,), + mode=Mode.NONSTREAM, + ) + ) + def test_bedrock_json_schema_response_exposes_text_before_json_tool_call( + self, proxy: ProxyClient, resources: ResourceManager, sdk: SdkClients + ) -> None: + model = f"e2e-bedrock-json-mode-preamble-{unique_marker()}" + model_id = proxy.create_model( + model, + LiteLLMParamsBody( + model=BEDROCK_JSON_SCHEMA_BACKEND, + aws_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + aws_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + aws_region_name="us-east-1", + ), + ) + resources.defer(lambda: proxy.delete_model(model_id)) + key: Final = resources.key() + + assert any( + _warehouse_message_has_preamble( + sdk.openai(key) + .chat.completions.create( + model=model, + messages=[{"role": "user", "content": WAREHOUSE_PROMPT}], + response_format=WAREHOUSE_RESPONSE_FORMAT, + max_tokens=1500, + ) + .choices[0] + .message + ) + for _ in range(3) + ), "No model-generated preamble appeared in three Bedrock calls" + @pytest.mark.covers("llm.bedrock_native.bedrock_converse.basic.nonstream.works") @meta( Subject( diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index ef002a73e52..1c748f0c83d 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -9716,6 +9716,69 @@ def test_json_mode_keeps_the_thinking_that_preceded_the_json_tool_call(): assert [block["signature"] for block in message.thinking_blocks] == [thinking_block["signature"]] +def test_json_mode_preserves_nonblank_text_before_json_tool_call(): + preamble: Final = "Let me calculate the usable items. " + body: Final = _json_mode_message_body("tool_use", None) + result: Final = _transform_json_mode_body( + {**body, "content": [{"type": "text", "text": preamble}, *body["content"]]} + ) + + message: Final = result.choices[0].message + assert message.content == '{"answer": "Hello"}' + assert message.provider_specific_fields["json_mode_preamble_text"] == preamble + + +def test_json_mode_preamble_preserves_refusal_provider_fields(): + preamble: Final = "Let me explain the result. " + stop_details: Final = ANTHROPIC_DOCS_REFUSAL_MESSAGE["stop_details"] + body: Final = _json_mode_message_body("refusal", stop_details) + result: Final = _transform_json_mode_body( + {**body, "content": [{"type": "text", "text": preamble}, *body["content"]]} + ) + + message: Final = result.choices[0].message + assert message.content == '{"answer": "Hello"}' + assert message.provider_specific_fields == { + "stop_reason": "refusal", + "stop_details": stop_details, + "json_mode_preamble_text": preamble, + } + + +def test_json_mode_text_without_tool_call_is_unchanged(): + body: Final = _json_mode_message_body("tool_use", None) + result: Final = _transform_json_mode_body( + {**body, "content": [{"type": "text", "text": "A plain response."}]} + ) + + message: Final = result.choices[0].message + assert message.content == "A plain response." + assert message.tool_calls is None + assert "json_mode_preamble_text" not in (message.provider_specific_fields or {}) + + +def test_json_mode_with_a_real_tool_keeps_mixed_response_behavior(): + preamble: Final = "I will look up the details. " + body: Final = _json_mode_message_body("tool_use", None) + user_tool_call: Final = { + "type": "tool_use", + "id": "toolu_user", + "name": "get_weather", + "input": {"location": "San Francisco"}, + } + result: Final = _transform_json_mode_body( + {**body, "content": [{"type": "text", "text": preamble}, *body["content"], user_tool_call]} + ) + + message: Final = result.choices[0].message + assert message.content == f'{preamble}{{"answer": "Hello"}}' + assert message.tool_calls is not None + assert len(message.tool_calls) == 1 + assert message.tool_calls[0].function.name == "get_weather" + assert message.tool_calls[0].function.arguments == '{"location": "San Francisco"}' + assert "json_mode_preamble_text" not in (message.provider_specific_fields or {}) + + def test_non_dict_upstream_body_still_transforms_without_refusal_fields() -> None: result = AnthropicConfig().transform_parsed_response( completion_response=MagicMock(), diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation.py b/tests/unit/llms/bedrock/chat/test_converse_transformation.py index 154797bab5a..f89ca376e47 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation.py @@ -6736,6 +6736,108 @@ def test_transform_response_finish_reason_stop_when_json_mode_filters_all_tools( assert result.choices[0].finish_reason == "stop" +def _transform_json_mode_content_blocks(content: list[dict[str, object]]) -> ModelResponse: + response_json: Final = { + "metrics": {"latencyMs": 100}, + "output": {"message": {"role": "assistant", "content": content}}, + "stopReason": "tool_use", + "usage": { + "inputTokens": 50, + "outputTokens": 20, + "totalTokens": 70, + "cacheReadInputTokenCount": 0, + "cacheReadInputTokens": 0, + "cacheWriteInputTokenCount": 0, + "cacheWriteInputTokens": 0, + }, + } + + class MockResponse: + def json(self) -> dict[str, object]: + return response_json + + @property + def text(self) -> str: + return json.dumps(response_json) + + return AmazonConverseConfig()._transform_response( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + response=MockResponse(), + model_response=ModelResponse(), + stream=False, + logging_obj=None, + optional_params={"json_mode": True}, + api_key=None, + data=None, + messages=[], + encoding=None, + ) + + +def test_bedrock_json_mode_preserves_nonblank_text_before_json_tool_call(): + preamble: Final = "Let me calculate the usable items. " + result: Final = _transform_json_mode_content_blocks( + [ + {"text": preamble}, + { + "toolUse": { + "toolUseId": "tooluse_json_001", + "name": "json_tool_call", + "input": {"answer": 42}, + } + }, + ] + ) + + message: Final = result.choices[0].message + assert message.content == '{"answer": 42}' + assert message.provider_specific_fields["json_mode_preamble_text"] == preamble + + +def test_bedrock_json_mode_omits_whitespace_only_preamble(): + result: Final = _transform_json_mode_content_blocks( + [ + {"text": " \n\t"}, + { + "toolUse": { + "toolUseId": "tooluse_json_001", + "name": "json_tool_call", + "input": {"answer": 42}, + } + }, + ] + ) + + message: Final = result.choices[0].message + assert message.content == '{"answer": 42}' + assert "json_mode_preamble_text" not in (message.provider_specific_fields or {}) + + +def test_bedrock_json_mode_preamble_preserves_reasoning_content_fields(): + reasoning_content: Final = { + "reasoningText": {"text": "The answer is 42.", "signature": "signature-1"} + } + preamble: Final = "Here is the result. " + result: Final = _transform_json_mode_content_blocks( + [ + {"reasoningContent": reasoning_content}, + {"text": preamble}, + { + "toolUse": { + "toolUseId": "tooluse_json_001", + "name": "json_tool_call", + "input": {"answer": 42}, + } + }, + ] + ) + + message: Final = result.choices[0].message + assert message.content == '{"answer": 42}' + assert message.provider_specific_fields["json_mode_preamble_text"] == preamble + assert message.provider_specific_fields["reasoningContentBlocks"] == [reasoning_content] + + def test_transform_response_json_mode_truncated_tool_call_keeps_length_finish_reason(): """ When json_mode filters out the synthetic json_tool_call but Bedrock