fix(google): Fix 400 Bad Request error with Gemini 3 thought signatures (#6035)
Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
@@ -9,10 +9,25 @@ use rmcp::model::{
|
||||
};
|
||||
use std::borrow::Cow;
|
||||
|
||||
use crate::conversation::message::{Message, MessageContent};
|
||||
use crate::conversation::message::{Message, MessageContent, ProviderMetadata};
|
||||
use serde_json::{json, Map, Value};
|
||||
use std::ops::Deref;
|
||||
|
||||
pub const THOUGHT_SIGNATURE_KEY: &str = "thoughtSignature";
|
||||
|
||||
pub fn metadata_with_signature(signature: &str) -> ProviderMetadata {
|
||||
let mut map = ProviderMetadata::new();
|
||||
map.insert(THOUGHT_SIGNATURE_KEY.to_string(), json!(signature));
|
||||
map
|
||||
}
|
||||
|
||||
pub fn get_thought_signature(metadata: &Option<ProviderMetadata>) -> Option<&str> {
|
||||
metadata
|
||||
.as_ref()
|
||||
.and_then(|m| m.get(THOUGHT_SIGNATURE_KEY))
|
||||
.and_then(|v| v.as_str())
|
||||
}
|
||||
|
||||
/// Convert internal Message format to Google's API message specification
|
||||
pub fn format_messages(messages: &[Message]) -> Vec<Value> {
|
||||
messages
|
||||
@@ -58,8 +73,8 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
|
||||
let mut part = Map::new();
|
||||
part.insert("functionCall".to_string(), json!(function_call_part));
|
||||
|
||||
if let Some(signature) = &request.thought_signature {
|
||||
part.insert("thoughtSignature".to_string(), json!(signature));
|
||||
if let Some(signature) = get_thought_signature(&request.metadata) {
|
||||
part.insert(THOUGHT_SIGNATURE_KEY.to_string(), json!(signature));
|
||||
}
|
||||
|
||||
parts.push(json!(part));
|
||||
@@ -117,15 +132,44 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
|
||||
if text.is_empty() {
|
||||
text = "Tool call is done.".to_string();
|
||||
}
|
||||
parts.push(json!({
|
||||
"functionResponse": {
|
||||
"name": response.id,
|
||||
"response": {"content": {"text": text}},
|
||||
}}
|
||||
));
|
||||
let mut part = Map::new();
|
||||
let mut function_response = Map::new();
|
||||
function_response.insert("name".to_string(), json!(response.id));
|
||||
function_response.insert(
|
||||
"response".to_string(),
|
||||
json!({"content": {"text": text}}),
|
||||
);
|
||||
part.insert(
|
||||
"functionResponse".to_string(),
|
||||
json!(function_response),
|
||||
);
|
||||
if let Some(signature) = get_thought_signature(&response.metadata) {
|
||||
part.insert(
|
||||
THOUGHT_SIGNATURE_KEY.to_string(),
|
||||
json!(signature),
|
||||
);
|
||||
}
|
||||
parts.push(json!(part));
|
||||
}
|
||||
Err(e) => {
|
||||
parts.push(json!({"text":format!("Error: {}", e)}));
|
||||
let mut part = Map::new();
|
||||
let mut function_response = Map::new();
|
||||
function_response.insert("name".to_string(), json!(response.id));
|
||||
function_response.insert(
|
||||
"response".to_string(),
|
||||
json!({"content": {"text": format!("Error: {}", e)}}),
|
||||
);
|
||||
part.insert(
|
||||
"functionResponse".to_string(),
|
||||
json!(function_response),
|
||||
);
|
||||
if let Some(signature) = get_thought_signature(&response.metadata) {
|
||||
part.insert(
|
||||
THOUGHT_SIGNATURE_KEY.to_string(),
|
||||
json!(signature),
|
||||
);
|
||||
}
|
||||
parts.push(json!(part));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -282,15 +326,30 @@ pub fn response_to_message(response: Value) -> Result<Message> {
|
||||
.and_then(|parts| parts.as_array())
|
||||
.unwrap_or(&binding);
|
||||
|
||||
// Track the last seen thought signature to use as fallback for function calls without one
|
||||
// This handles cases where Google's API returns multiple function calls but only includes
|
||||
// thoughtSignature on some of them
|
||||
let mut last_signature: Option<String> = None;
|
||||
|
||||
let has_function_calls = parts.iter().any(|p| p.get("functionCall").is_some());
|
||||
|
||||
for part in parts {
|
||||
let thought_signature = part
|
||||
.get("thoughtSignature")
|
||||
let signature = part
|
||||
.get(THOUGHT_SIGNATURE_KEY)
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
if signature.is_some() {
|
||||
last_signature = signature.clone();
|
||||
}
|
||||
|
||||
if let Some(text) = part.get("text").and_then(|v| v.as_str()) {
|
||||
if let Some(sig) = thought_signature {
|
||||
content.push(MessageContent::thinking(text.to_string(), sig));
|
||||
// Text is "thinking" only if:
|
||||
// 1. It has a signature AND
|
||||
// 2. The response also contains function calls (meaning this is reasoning before acting)
|
||||
// If there are no function calls, this is the final response and should be shown
|
||||
if let (Some(sig), true) = (&signature, has_function_calls) {
|
||||
content.push(MessageContent::thinking(text.to_string(), sig.clone()));
|
||||
} else {
|
||||
content.push(MessageContent::text(text.to_string()));
|
||||
}
|
||||
@@ -316,16 +375,17 @@ pub fn response_to_message(response: Value) -> Result<Message> {
|
||||
content.push(MessageContent::tool_request(id, Err(error)));
|
||||
} else {
|
||||
let parameters = function_call.get("args");
|
||||
if let Some(params) = parameters {
|
||||
content.push(MessageContent::tool_request_with_signature(
|
||||
id,
|
||||
Ok(CallToolRequestParam {
|
||||
name: name.into(),
|
||||
arguments: Some(object(params.clone())),
|
||||
}),
|
||||
thought_signature,
|
||||
));
|
||||
}
|
||||
let arguments = parameters.map(|params| object(params.clone()));
|
||||
let effective_signature = signature.as_deref().or(last_signature.as_deref());
|
||||
let metadata = effective_signature.map(metadata_with_signature);
|
||||
content.push(MessageContent::tool_request_with_metadata(
|
||||
id,
|
||||
Ok(CallToolRequestParam {
|
||||
name: name.into(),
|
||||
arguments,
|
||||
}),
|
||||
metadata.as_ref(),
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -919,4 +979,66 @@ mod tests {
|
||||
assert_eq!(regular_field["type"], "number");
|
||||
assert_eq!(regular_field["description"], "A regular number field");
|
||||
}
|
||||
|
||||
fn google_response(parts: Vec<Value>) -> Value {
|
||||
json!({"candidates": [{"content": {"role": "model", "parts": parts}}]})
|
||||
}
|
||||
|
||||
fn tool_result(text: &str) -> CallToolResult {
|
||||
CallToolResult {
|
||||
content: vec![Content::text(text)],
|
||||
structured_content: None,
|
||||
is_error: Some(false),
|
||||
meta: None,
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_thought_signature_roundtrip() {
|
||||
const SIG: &str = "thought_sig_abc";
|
||||
|
||||
let response_with_tools = google_response(vec![
|
||||
json!({"text": "Let me think...", "thoughtSignature": SIG}),
|
||||
json!({"functionCall": {"name": "shell", "args": {"cmd": "ls"}}, "thoughtSignature": SIG}),
|
||||
json!({"functionCall": {"name": "read", "args": {}}}),
|
||||
]);
|
||||
|
||||
let native = response_to_message(response_with_tools).unwrap();
|
||||
assert_eq!(native.content.len(), 3, "Expected thinking + 2 tool calls");
|
||||
|
||||
let thinking = native.content[0]
|
||||
.as_thinking()
|
||||
.expect("Text with function calls should be Thinking");
|
||||
assert_eq!(thinking.signature, SIG);
|
||||
|
||||
let req1 = native.content[1]
|
||||
.as_tool_request()
|
||||
.expect("Second part should be ToolRequest");
|
||||
let req2 = native.content[2]
|
||||
.as_tool_request()
|
||||
.expect("Third part should be ToolRequest");
|
||||
assert_eq!(get_thought_signature(&req1.metadata), Some(SIG));
|
||||
assert_eq!(
|
||||
get_thought_signature(&req2.metadata),
|
||||
Some(SIG),
|
||||
"Should inherit"
|
||||
);
|
||||
|
||||
let tool_response = Message::user().with_tool_response_with_metadata(
|
||||
req1.id.clone(),
|
||||
Ok(tool_result("output")),
|
||||
req1.metadata.as_ref(),
|
||||
);
|
||||
let google_out = format_messages(&[native, tool_response]);
|
||||
assert_eq!(google_out[0]["parts"][0]["thoughtSignature"], SIG);
|
||||
assert_eq!(google_out[1]["parts"][0]["thoughtSignature"], SIG);
|
||||
|
||||
let final_response =
|
||||
google_response(vec![json!({"text": "Done!", "thoughtSignature": SIG})]);
|
||||
let final_native = response_to_message(final_response).unwrap();
|
||||
assert!(
|
||||
final_native.content[0].as_text().is_some(),
|
||||
"Text-only = final answer"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user