fix(bedrock): wrong system prompt transformation (BerriAI#10120)

* fix(bedrock): wrong system transformation * chore: add one more test case --------- Co-authored-by: Krish Dholakia <krrishdholakia@gmail.com>

Author: hewliyang

litellm/llms/bedrock/chat/converse_transformation.py

@@ -376,25 +376,27 @@ def _transform_system_message(
  system_content_blocks: List[SystemContentBlock] = []
  for idx, message in enumerate(messages):
  if message["role"] == "system":
- _system_content_block: Optional[SystemContentBlock] = None
- _cache_point_block: Optional[SystemContentBlock] = None
- if isinstance(message["content"], str) and len(message["content"]) > 0:
- _system_content_block = SystemContentBlock(text=message["content"])
- _cache_point_block = self._get_cache_point_block(
+ system_prompt_indices.append(idx)
+ if isinstance(message["content"], str) and message["content"]:
+ system_content_blocks.append(
+ SystemContentBlock(text=message["content"])
+ )
+ cache_block = self._get_cache_point_block(
  message, block_type="system"
  )
+ if cache_block:
+ system_content_blocks.append(cache_block)
  elif isinstance(message["content"], list):
  for m in message["content"]:
- if m.get("type", "") == "text" and len(m["text"]) > 0:
- _system_content_block = SystemContentBlock(text=m["text"])
- _cache_point_block = self._get_cache_point_block(
+ if m.get("type") == "text" and m.get("text"):
+ system_content_blocks.append(
+ SystemContentBlock(text=m["text"])
+ )
+ cache_block = self._get_cache_point_block(
  m, block_type="system"
  )
- if _system_content_block is not None:
- system_content_blocks.append(_system_content_block)
- if _cache_point_block is not None:
- system_content_blocks.append(_cache_point_block)
- system_prompt_indices.append(idx)
+ if cache_block:
+ system_content_blocks.append(cache_block)
  if len(system_prompt_indices) > 0:
  for idx in reversed(system_prompt_indices):
  messages.pop(idx)

tests/litellm/llms/bedrock/chat/test_converse_transformation.py

@@ -41,6 +41,85 @@ def test_transform_usage():
  assert openai_usage._cache_creation_input_tokens == usage["cacheWriteInputTokens"]
  assert openai_usage._cache_read_input_tokens == usage["cacheReadInputTokens"]
 
+def test_transform_system_message():
+ config = AmazonConverseConfig()
+
+ # Case 1:
+ # System message popped
+ # User message remains
+ messages = [
+ {"role": "system", "content": "You are a helpful assistant."},
+ {"role": "user", "content": "Hello!"},
+ ]
+ out_messages, system_blocks = config._transform_system_message(messages.copy())
+ assert len(out_messages) == 1
+ assert out_messages[0]["role"] == "user"
+ assert len(system_blocks) == 1
+ assert system_blocks[0]["text"] == "You are a helpful assistant."
+
+ # Case 2: System message with list content (type text)
+ messages = [
+ {
+ "role": "system",
+ "content": [
+ {"type": "text", "text": "System prompt 1"},
+ {"type": "text", "text": "System prompt 2"},
+ ],
+ },
+ {"role": "user", "content": "Hi!"},
+ ]
+ out_messages, system_blocks = config._transform_system_message(messages.copy())
+ assert len(out_messages) == 1
+ assert out_messages[0]["role"] == "user"
+ assert len(system_blocks) == 2
+ assert system_blocks[0]["text"] == "System prompt 1"
+ assert system_blocks[1]["text"] == "System prompt 2"
+
+ # Case 3: System message with cache_control (should add cachePoint)
+ messages = [
+ {
+ "role": "system",
+ "content": "Cache this!",
+ "cache_control": {"type": "ephemeral"},
+ },
+ {"role": "user", "content": "Hi!"},
+ ]
+ out_messages, system_blocks = config._transform_system_message(messages.copy())
+ assert len(out_messages) == 1
+ assert len(system_blocks) == 2
+ assert system_blocks[0]["text"] == "Cache this!"
+ assert "cachePoint" in system_blocks[1]
+ assert system_blocks[1]["cachePoint"]["type"] == "default"
+
+ # Case 3b: System message with two blocks, one with cache_control and one without
+ messages = [
+ {
+ "role": "system",
+ "content": [
+ {"type": "text", "text": "Cache this!", "cache_control": {"type": "ephemeral"}},
+ {"type": "text", "text": "Don't cache this!"},
+ ],
+ },
+ {"role": "user", "content": "Hi!"},
+ ]
+ out_messages, system_blocks = config._transform_system_message(messages.copy())
+ assert len(out_messages) == 1
+ assert len(system_blocks) == 3
+ assert system_blocks[0]["text"] == "Cache this!"
+ assert "cachePoint" in system_blocks[1]
+ assert system_blocks[1]["cachePoint"]["type"] == "default"
+ assert system_blocks[2]["text"] == "Don't cache this!"
+
+ # Case 4: Non-system messages are not affected
+ messages = [
+ {"role": "user", "content": "Hello!"},
+ {"role": "assistant", "content": "Hi!"},
+ ]
+ out_messages, system_blocks = config._transform_system_message(messages.copy())
+ assert len(out_messages) == 2
+ assert out_messages[0]["role"] == "user"
+ assert out_messages[1]["role"] == "assistant"
+ assert system_blocks == []
 
 def test_transform_thinking_blocks_with_redacted_content():
  thinking_blocks = [
@@ -59,3 +138,4 @@ def test_transform_thinking_blocks_with_redacted_content():
  assert len(transformed_thinking_blocks) == 2
  assert transformed_thinking_blocks[0]["type"] == "thinking"
  assert transformed_thinking_blocks[1]["type"] == "redacted_thinking"
+