mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
test: cover JSON mode preamble text preservation
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
ryan
parent
f1795e6d44
commit
e50ae35dde
3 changed files with 259 additions and 0 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue