mirror of
https://github.com/block/goose.git
synced 2026-06-02 06:19:33 +02:00
fix(databricks): ensure parallel tool image responses don't interleave tool results (#9241)
Signed-off-by: Steve Marshall <steve.marshall@fasthosts.com> Signed-off-by: Douwe Osinga <douwe@squareup.com> Co-authored-by: Douwe Osinga <douwe@squareup.com>
This commit is contained in:
@@ -124,6 +124,8 @@ fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<Data
|
||||
let mut content_array = Vec::new();
|
||||
let mut has_tool_calls = false;
|
||||
let mut has_multiple_content = false;
|
||||
// Deferred so all tool-role messages stay consecutive (required by Claude via Databricks).
|
||||
let mut pending_image_messages: Vec<DatabricksMessage> = Vec::new();
|
||||
|
||||
for content in &message.content {
|
||||
match content {
|
||||
@@ -190,7 +192,13 @@ fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<Data
|
||||
}
|
||||
}
|
||||
MessageContent::ToolResponse(response) => {
|
||||
result.extend(format_tool_response(response, image_format));
|
||||
for msg in format_tool_response(response, image_format) {
|
||||
if msg.role == "user" {
|
||||
pending_image_messages.push(msg);
|
||||
} else {
|
||||
result.push(msg);
|
||||
}
|
||||
}
|
||||
}
|
||||
MessageContent::Image(image) => {
|
||||
content_array.push(convert_image(image, image_format));
|
||||
@@ -212,6 +220,8 @@ fn format_messages(messages: &[Message], image_format: &ImageFormat) -> Vec<Data
|
||||
}
|
||||
}
|
||||
|
||||
result.extend(pending_image_messages);
|
||||
|
||||
if !content_array.is_empty() {
|
||||
converted.content = if content_array.len() == 1
|
||||
&& !has_multiple_content
|
||||
@@ -1683,4 +1693,70 @@ mod tests {
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_parallel_tool_responses_with_images_are_consecutive() -> anyhow::Result<()> {
|
||||
// Regression: #7449 — parallel tool calls with images must keep tool messages consecutive.
|
||||
let messages = vec![
|
||||
Message::assistant()
|
||||
.with_tool_request("id1", Ok(CallToolRequestParams::new("tool_a")))
|
||||
.with_tool_request("id2", Ok(CallToolRequestParams::new("tool_b"))),
|
||||
Message::user()
|
||||
.with_tool_response(
|
||||
"id1",
|
||||
Ok(CallToolResult::success(vec![Content::image(
|
||||
"base64data1".to_string(),
|
||||
"image/png".to_string(),
|
||||
)])),
|
||||
)
|
||||
.with_tool_response(
|
||||
"id2",
|
||||
Ok(CallToolResult::success(vec![Content::image(
|
||||
"base64data2".to_string(),
|
||||
"image/png".to_string(),
|
||||
)])),
|
||||
),
|
||||
];
|
||||
|
||||
let as_value =
|
||||
serde_json::to_value(format_messages(&messages, &ImageFormat::OpenAi)).unwrap();
|
||||
let spec = as_value.as_array().unwrap();
|
||||
let roles: Vec<&str> = spec.iter().map(|m| m["role"].as_str().unwrap()).collect();
|
||||
|
||||
// Without the fix this was ["assistant", "tool", "user", "tool", "user"].
|
||||
assert_eq!(roles, vec!["assistant", "tool", "tool", "user", "user"]);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mixed_tool_responses_image_and_text_ordering() -> anyhow::Result<()> {
|
||||
// Mixed case: only one tool response has an image.
|
||||
let messages = vec![
|
||||
Message::assistant()
|
||||
.with_tool_request("id1", Ok(CallToolRequestParams::new("tool_a")))
|
||||
.with_tool_request("id2", Ok(CallToolRequestParams::new("tool_b"))),
|
||||
Message::user()
|
||||
.with_tool_response(
|
||||
"id1",
|
||||
Ok(CallToolResult::success(vec![Content::text("text result")])),
|
||||
)
|
||||
.with_tool_response(
|
||||
"id2",
|
||||
Ok(CallToolResult::success(vec![Content::image(
|
||||
"base64data".to_string(),
|
||||
"image/png".to_string(),
|
||||
)])),
|
||||
),
|
||||
];
|
||||
|
||||
let as_value =
|
||||
serde_json::to_value(format_messages(&messages, &ImageFormat::OpenAi)).unwrap();
|
||||
let spec = as_value.as_array().unwrap();
|
||||
let roles: Vec<&str> = spec.iter().map(|m| m["role"].as_str().unwrap()).collect();
|
||||
|
||||
assert_eq!(roles, vec!["assistant", "tool", "tool", "user"]);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user