mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(bedrock): preserve cache_control when converting trailing user messages to guarded_text
This commit is contained in:
parent
668df9494a
commit
5c7bfd430c
2 changed files with 84 additions and 3 deletions
|
|
@ -179,6 +179,8 @@ class AmazonConverseConfig(BaseConfig):
|
|||
for item in content:
|
||||
if isinstance(item, dict) and item.get("type") == "text":
|
||||
new_item = {"type": "guarded_text", "text": item["text"]} # type: ignore
|
||||
if "cache_control" in item:
|
||||
new_item["cache_control"] = item["cache_control"]
|
||||
new_content.append(new_item)
|
||||
else:
|
||||
new_content.append(item)
|
||||
|
|
@ -186,9 +188,11 @@ class AmazonConverseConfig(BaseConfig):
|
|||
messages_copy[user_message_index]["content"] = new_content # type: ignore
|
||||
elif isinstance(content, str):
|
||||
# If content is a string, convert it to guarded_text
|
||||
messages_copy[user_message_index]["content"] = [ # type: ignore
|
||||
{"type": "guarded_text", "text": content} # type: ignore
|
||||
]
|
||||
guarded_text_item = {"type": "guarded_text", "text": content} # type: ignore
|
||||
message_cache_control = user_message.get("cache_control")
|
||||
if message_cache_control is not None:
|
||||
guarded_text_item["cache_control"] = message_cache_control
|
||||
messages_copy[user_message_index]["content"] = [guarded_text_item] # type: ignore
|
||||
|
||||
return messages_copy
|
||||
|
||||
|
|
|
|||
|
|
@ -2602,6 +2602,83 @@ def test_auto_convert_last_user_message_string_content():
|
|||
)
|
||||
|
||||
|
||||
def test_guarded_text_conversion_preserves_cache_control_on_content_block():
|
||||
"""Regression for #33281: converting a trailing user text block to guarded_text
|
||||
must carry over its cache_control so a cachePoint is still emitted for prompt caching."""
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Summarize the attached document",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
optional_params = {
|
||||
"guardrailConfig": {"guardrailIdentifier": "gr-abc123", "guardrailVersion": "1"}
|
||||
}
|
||||
|
||||
converted_messages = config._convert_consecutive_user_messages_to_guarded_text(
|
||||
messages, optional_params
|
||||
)
|
||||
guarded_block = converted_messages[0]["content"][0]
|
||||
assert guarded_block["type"] == "guarded_text"
|
||||
assert guarded_block["cache_control"] == {"type": "ephemeral"}
|
||||
|
||||
result = config._transform_request(
|
||||
model="us.amazon.nova-pro-v1:0",
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
user_content = result["messages"][0]["content"]
|
||||
assert any("guardContent" in block for block in user_content)
|
||||
assert any(block.get("cachePoint") == {"type": "default"} for block in user_content)
|
||||
|
||||
|
||||
def test_guarded_text_conversion_preserves_message_level_cache_control_string_content():
|
||||
"""Regression for #33281: message-level cache_control on a trailing string user
|
||||
message must survive the guarded_text conversion and still emit a cachePoint."""
|
||||
config = AmazonConverseConfig()
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Summarize the attached document",
|
||||
"cache_control": {"type": "ephemeral"},
|
||||
}
|
||||
]
|
||||
|
||||
optional_params = {
|
||||
"guardrailConfig": {"guardrailIdentifier": "gr-abc123", "guardrailVersion": "1"}
|
||||
}
|
||||
|
||||
converted_messages = config._convert_consecutive_user_messages_to_guarded_text(
|
||||
messages, optional_params
|
||||
)
|
||||
guarded_block = converted_messages[0]["content"][0]
|
||||
assert guarded_block["type"] == "guarded_text"
|
||||
assert guarded_block["cache_control"] == {"type": "ephemeral"}
|
||||
|
||||
result = config._transform_request(
|
||||
model="us.amazon.nova-pro-v1:0",
|
||||
messages=messages,
|
||||
optional_params=optional_params,
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
user_content = result["messages"][0]["content"]
|
||||
assert any("guardContent" in block for block in user_content)
|
||||
assert any(block.get("cachePoint") == {"type": "default"} for block in user_content)
|
||||
|
||||
|
||||
def test_no_conversion_when_no_guardrail_config():
|
||||
"""Test that no conversion happens when guardrailConfig is not present."""
|
||||
config = AmazonConverseConfig()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue