Fix: Ensure OpenAI requests end with user message for Bedrock/LiteLLM compatibility
When LiteLLM routes OpenAI-format requests to AWS Bedrock, the provider rejects conversations ending with an assistant message with: 'This model does not support assistant message prefill.' Add ensure_ends_with_user_message() as the final sanitization step in sanitize_messages_for_openai(). If the conversation ends with an assistant message, a minimal 'Continue.' user message is appended. Updated tests to reflect the new behavior and added dedicated tests for the new function.
This commit is contained in:
@@ -6,9 +6,11 @@ use crate::ai::provider::types::{ContentPart, ConversationMessage, MessageConten
|
||||
/// alternation and allows system messages anywhere. The main constraints are:
|
||||
/// - Tool results must reference a valid tool_call_id from a preceding assistant message
|
||||
/// - Tool calls in assistant messages must eventually have matching tool results
|
||||
/// - When routed to Bedrock via LiteLLM, the conversation must end with a user message
|
||||
pub fn sanitize_messages_for_openai(messages: &mut Vec<ConversationMessage>) {
|
||||
remove_orphaned_tool_results(messages);
|
||||
synthesize_missing_tool_results(messages);
|
||||
ensure_ends_with_user_message(messages);
|
||||
}
|
||||
|
||||
/// Removes tool_result messages that reference tool_use_ids not found in any
|
||||
@@ -218,3 +220,24 @@ fn collect_tool_result_ids(content: &MessageContent, ids: &mut std::collections:
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
/// Ensures the message array ends with a user message. Some providers (e.g. Bedrock
|
||||
/// via LiteLLM) reject requests where the conversation ends with an assistant message
|
||||
/// ("assistant message prefill"). If the last message is from the assistant, append a
|
||||
/// minimal user message to satisfy this constraint.
|
||||
fn ensure_ends_with_user_message(messages: &mut Vec<ConversationMessage>) {
|
||||
if messages.is_empty() {
|
||||
return;
|
||||
}
|
||||
if let Some(last) = messages.last() {
|
||||
if last.role == MessageRole::Assistant {
|
||||
log::info!(
|
||||
"[openai] Conversation ends with assistant message — appending user continuation message"
|
||||
);
|
||||
messages.push(ConversationMessage {
|
||||
role: MessageRole::User,
|
||||
content: MessageContent::Text("Continue.".to_string()),
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -103,8 +103,13 @@ fn test_does_not_require_user_assistant_alternation() {
|
||||
|
||||
sanitize_messages_for_openai(&mut messages);
|
||||
|
||||
// Both user messages should remain — OpenAI allows consecutive same-role
|
||||
assert_eq!(messages.len(), 3);
|
||||
// Both user messages should remain — OpenAI allows consecutive same-role.
|
||||
// A trailing user message is appended since it ends with assistant.
|
||||
assert_eq!(messages.len(), 4);
|
||||
assert_eq!(messages[0].role, MessageRole::User);
|
||||
assert_eq!(messages[1].role, MessageRole::User);
|
||||
assert_eq!(messages[2].role, MessageRole::Assistant);
|
||||
assert_eq!(messages[3].role, MessageRole::User);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -116,9 +121,12 @@ fn test_does_not_require_starting_with_user() {
|
||||
|
||||
sanitize_messages_for_openai(&mut messages);
|
||||
|
||||
// Should NOT prepend a user message (unlike Bedrock)
|
||||
assert_eq!(messages.len(), 1);
|
||||
// Should NOT prepend a user message (unlike Bedrock), but WILL append
|
||||
// a user continuation because Bedrock (via LiteLLM) requires the
|
||||
// conversation to end with a user message.
|
||||
assert_eq!(messages.len(), 2);
|
||||
assert_eq!(messages[0].role, MessageRole::Assistant);
|
||||
assert_eq!(messages[1].role, MessageRole::User);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -141,8 +149,51 @@ fn test_multipart_tool_uses_all_get_results() {
|
||||
|
||||
sanitize_messages_for_openai(&mut messages);
|
||||
|
||||
// Should synthesize results for both unanswered tool calls
|
||||
assert_eq!(messages.len(), 3);
|
||||
// Should synthesize results for both unanswered tool calls in a single
|
||||
// user message with MultiPart content
|
||||
assert_eq!(messages.len(), 2);
|
||||
assert_eq!(messages[0].role, MessageRole::Assistant);
|
||||
assert_eq!(messages[1].role, MessageRole::User);
|
||||
assert_eq!(messages[2].role, MessageRole::User);
|
||||
match &messages[1].content {
|
||||
MessageContent::MultiPart(parts) => {
|
||||
assert_eq!(parts.len(), 2);
|
||||
for part in parts {
|
||||
match part {
|
||||
ContentPart::ToolResult { is_error, .. } => assert!(is_error),
|
||||
_ => panic!("Expected ToolResult part"),
|
||||
}
|
||||
}
|
||||
}
|
||||
_ => panic!("Expected MultiPart content"),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ensure_ends_with_user_message_no_op_when_already_user() {
|
||||
let mut messages = vec![
|
||||
ConversationMessage {
|
||||
role: MessageRole::Assistant,
|
||||
content: MessageContent::Text("Hello".to_string()),
|
||||
},
|
||||
ConversationMessage {
|
||||
role: MessageRole::User,
|
||||
content: MessageContent::Text("Hi back".to_string()),
|
||||
},
|
||||
];
|
||||
|
||||
sanitize_messages_for_openai(&mut messages);
|
||||
|
||||
// Already ends with user — no extra message appended
|
||||
assert_eq!(messages.len(), 2);
|
||||
assert_eq!(messages[1].role, MessageRole::User);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_ensure_ends_with_user_message_empty_messages() {
|
||||
let mut messages: Vec<ConversationMessage> = vec![];
|
||||
|
||||
sanitize_messages_for_openai(&mut messages);
|
||||
|
||||
// Empty messages should stay empty
|
||||
assert!(messages.is_empty());
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user