fix(google): handle thoughtSignature differently for Gemini 2.5 vs 3 (#6890)
Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
@@ -232,6 +232,25 @@ pub fn process_response_part(
|
|||||||
process_response_part_impl(part, last_signature, handling)
|
process_response_part_impl(part, last_signature, handling)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
/// Gemini 2.x includes thoughtSignature on first chunk as metadata, not actual thinking.
|
||||||
|
fn process_response_part_for_model(
|
||||||
|
part: &Value,
|
||||||
|
last_signature: &mut Option<String>,
|
||||||
|
model_version: Option<&str>,
|
||||||
|
) -> Option<MessageContent> {
|
||||||
|
let is_gemini_2 = model_version
|
||||||
|
.map(|m| m.starts_with("gemini-2"))
|
||||||
|
.unwrap_or(false);
|
||||||
|
|
||||||
|
let has_signature = part.get(THOUGHT_SIGNATURE_KEY).is_some();
|
||||||
|
let handling = if has_signature && !is_gemini_2 {
|
||||||
|
SignedTextHandling::SignedTextAsThinking
|
||||||
|
} else {
|
||||||
|
SignedTextHandling::SignedTextAsRegularText
|
||||||
|
};
|
||||||
|
process_response_part_impl(part, last_signature, handling)
|
||||||
|
}
|
||||||
|
|
||||||
fn process_response_part_non_streaming(
|
fn process_response_part_non_streaming(
|
||||||
part: &Value,
|
part: &Value,
|
||||||
last_signature: &mut Option<String>,
|
last_signature: &mut Option<String>,
|
||||||
@@ -461,6 +480,8 @@ where
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
let model_version = chunk.get("modelVersion").and_then(|v| v.as_str());
|
||||||
|
|
||||||
let parts = chunk
|
let parts = chunk
|
||||||
.get("candidates")
|
.get("candidates")
|
||||||
.and_then(|v| v.as_array())
|
.and_then(|v| v.as_array())
|
||||||
@@ -471,7 +492,7 @@ where
|
|||||||
|
|
||||||
if let Some(parts) = parts {
|
if let Some(parts) = parts {
|
||||||
for part in parts {
|
for part in parts {
|
||||||
if let Some(content) = process_response_part(part, &mut last_signature) {
|
if let Some(content) = process_response_part_for_model(part, &mut last_signature, model_version) {
|
||||||
let message = Message::new(
|
let message = Message::new(
|
||||||
Role::Assistant,
|
Role::Assistant,
|
||||||
chrono::Utc::now().timestamp(),
|
chrono::Utc::now().timestamp(),
|
||||||
@@ -1072,18 +1093,16 @@ mod tests {
|
|||||||
async fn test_streaming_with_thought_signature() {
|
async fn test_streaming_with_thought_signature() {
|
||||||
use futures::StreamExt;
|
use futures::StreamExt;
|
||||||
|
|
||||||
let signed_stream = concat!(
|
let gemini3_stream = concat!(
|
||||||
r#"data: {"candidates": [{"content": {"role": "model", "#,
|
r#"data: {"candidates": [{"content": {"role": "model", "#,
|
||||||
r#""parts": [{"text": "Begin", "thoughtSignature": "sig123"}]}}]}"#,
|
r#""parts": [{"text": "Begin", "thoughtSignature": "sig123"}]}}], "#,
|
||||||
|
r#""modelVersion": "gemini-3-pro"}"#,
|
||||||
"\n",
|
"\n",
|
||||||
r#"data: {"candidates": [{"content": {"role": "model", "#,
|
r#"data: {"candidates": [{"content": {"role": "model", "#,
|
||||||
r#""parts": [{"text": " middle"}]}}]}"#,
|
r#""parts": [{"text": " end"}]}}], "modelVersion": "gemini-3-pro"}"#
|
||||||
"\n",
|
|
||||||
r#"data: {"candidates": [{"content": {"role": "model", "#,
|
|
||||||
r#""parts": [{"text": " end"}]}}]}"#
|
|
||||||
);
|
);
|
||||||
let lines: Vec<Result<String, anyhow::Error>> =
|
let lines: Vec<Result<String, anyhow::Error>> =
|
||||||
signed_stream.lines().map(|l| Ok(l.to_string())).collect();
|
gemini3_stream.lines().map(|l| Ok(l.to_string())).collect();
|
||||||
let stream = Box::pin(futures::stream::iter(lines));
|
let stream = Box::pin(futures::stream::iter(lines));
|
||||||
let mut message_stream = std::pin::pin!(response_to_streaming_message(stream));
|
let mut message_stream = std::pin::pin!(response_to_streaming_message(stream));
|
||||||
|
|
||||||
@@ -1094,20 +1113,70 @@ mod tests {
|
|||||||
let (message, _usage) = result.unwrap();
|
let (message, _usage) = result.unwrap();
|
||||||
if let Some(msg) = message {
|
if let Some(msg) = message {
|
||||||
match msg.content.first() {
|
match msg.content.first() {
|
||||||
Some(MessageContent::Text(text)) => {
|
Some(MessageContent::Text(text)) => text_parts.push(text.text.clone()),
|
||||||
text_parts.push(text.text.clone());
|
Some(MessageContent::Thinking(t)) => thinking_parts.push(t.thinking.clone()),
|
||||||
}
|
|
||||||
Some(MessageContent::Thinking(thinking)) => {
|
|
||||||
thinking_parts.push(thinking.thinking.clone());
|
|
||||||
assert_eq!(thinking.signature, "sig123");
|
|
||||||
}
|
|
||||||
_ => {}
|
_ => {}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
assert_eq!(thinking_parts, vec!["Begin"]);
|
assert_eq!(thinking_parts, vec!["Begin"]);
|
||||||
assert_eq!(text_parts, vec![" middle", " end"]);
|
assert_eq!(text_parts, vec![" end"]);
|
||||||
|
|
||||||
|
let gemini25_stream = concat!(
|
||||||
|
r#"data: {"candidates": [{"content": {"role": "model", "#,
|
||||||
|
r#""parts": [{"text": "Begin", "thoughtSignature": "sig123"}]}}], "#,
|
||||||
|
r#""modelVersion": "gemini-2.5-pro"}"#,
|
||||||
|
"\n",
|
||||||
|
r#"data: {"candidates": [{"content": {"role": "model", "#,
|
||||||
|
r#""parts": [{"text": " end"}]}}], "modelVersion": "gemini-2.5-pro"}"#
|
||||||
|
);
|
||||||
|
let lines: Vec<Result<String, anyhow::Error>> =
|
||||||
|
gemini25_stream.lines().map(|l| Ok(l.to_string())).collect();
|
||||||
|
let stream = Box::pin(futures::stream::iter(lines));
|
||||||
|
let mut message_stream = std::pin::pin!(response_to_streaming_message(stream));
|
||||||
|
|
||||||
|
let mut text_parts = Vec::new();
|
||||||
|
|
||||||
|
while let Some(result) = message_stream.next().await {
|
||||||
|
let (message, _usage) = result.unwrap();
|
||||||
|
if let Some(msg) = message {
|
||||||
|
if let Some(MessageContent::Text(text)) = msg.content.first() {
|
||||||
|
text_parts.push(text.text.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
assert_eq!(text_parts, vec!["Begin", " end"]);
|
||||||
|
|
||||||
|
let unknown_stream = concat!(
|
||||||
|
r#"data: {"candidates": [{"content": {"role": "model", "#,
|
||||||
|
r#""parts": [{"text": "Begin", "thoughtSignature": "sig123"}]}}]}"#,
|
||||||
|
"\n",
|
||||||
|
r#"data: {"candidates": [{"content": {"role": "model", "#,
|
||||||
|
r#""parts": [{"text": " end"}]}}]}"#
|
||||||
|
);
|
||||||
|
let lines: Vec<Result<String, anyhow::Error>> =
|
||||||
|
unknown_stream.lines().map(|l| Ok(l.to_string())).collect();
|
||||||
|
let stream = Box::pin(futures::stream::iter(lines));
|
||||||
|
let mut message_stream = std::pin::pin!(response_to_streaming_message(stream));
|
||||||
|
|
||||||
|
let mut text_parts = Vec::new();
|
||||||
|
let mut thinking_parts = Vec::new();
|
||||||
|
|
||||||
|
while let Some(result) = message_stream.next().await {
|
||||||
|
let (message, _usage) = result.unwrap();
|
||||||
|
if let Some(msg) = message {
|
||||||
|
match msg.content.first() {
|
||||||
|
Some(MessageContent::Text(text)) => text_parts.push(text.text.clone()),
|
||||||
|
Some(MessageContent::Thinking(t)) => thinking_parts.push(t.thinking.clone()),
|
||||||
|
_ => {}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
assert_eq!(thinking_parts, vec!["Begin"]);
|
||||||
|
assert_eq!(text_parts, vec![" end"]);
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
|
|||||||
Reference in New Issue
Block a user