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)
|
||||
}
|
||||
|
||||
/// 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(
|
||||
part: &Value,
|
||||
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
|
||||
.get("candidates")
|
||||
.and_then(|v| v.as_array())
|
||||
@@ -471,7 +492,7 @@ where
|
||||
|
||||
if let Some(parts) = 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(
|
||||
Role::Assistant,
|
||||
chrono::Utc::now().timestamp(),
|
||||
@@ -1072,18 +1093,16 @@ mod tests {
|
||||
async fn test_streaming_with_thought_signature() {
|
||||
use futures::StreamExt;
|
||||
|
||||
let signed_stream = concat!(
|
||||
let gemini3_stream = concat!(
|
||||
r#"data: {"candidates": [{"content": {"role": "model", "#,
|
||||
r#""parts": [{"text": "Begin", "thoughtSignature": "sig123"}]}}]}"#,
|
||||
r#""parts": [{"text": "Begin", "thoughtSignature": "sig123"}]}}], "#,
|
||||
r#""modelVersion": "gemini-3-pro"}"#,
|
||||
"\n",
|
||||
r#"data: {"candidates": [{"content": {"role": "model", "#,
|
||||
r#""parts": [{"text": " middle"}]}}]}"#,
|
||||
"\n",
|
||||
r#"data: {"candidates": [{"content": {"role": "model", "#,
|
||||
r#""parts": [{"text": " end"}]}}]}"#
|
||||
r#""parts": [{"text": " end"}]}}], "modelVersion": "gemini-3-pro"}"#
|
||||
);
|
||||
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 mut message_stream = std::pin::pin!(response_to_streaming_message(stream));
|
||||
|
||||
@@ -1094,20 +1113,70 @@ mod tests {
|
||||
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(thinking)) => {
|
||||
thinking_parts.push(thinking.thinking.clone());
|
||||
assert_eq!(thinking.signature, "sig123");
|
||||
}
|
||||
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![" 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]
|
||||
|
||||
Reference in New Issue
Block a user