tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

fix(google): handle more thoughtSignature vagaries during streaming (#7204)

Browse Source 
Signed-off-by: rabi <ramishra@redhat.com>
This commit is contained in:
Rabi Mishra
2026-02-17 18:43:48 +05:30
committed by GitHub
parent 920c0099fe
commit 49e684ab43
1 changed files with 53 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/providers/formats/google.rs
+53 -107
View File
@@ -227,38 +227,6 @@ enum SignedTextHandling {
SignedTextAsRegularText, SignedTextAsRegularText,
} }
pub fn process_response_part(
part: &Value,
last_signature: &mut Option<String>,
) -> Option<MessageContent> {
let has_signature = part.get(THOUGHT_SIGNATURE_KEY).is_some();
let handling = if has_signature {
SignedTextHandling::SignedTextAsThinking
} else {
SignedTextHandling::SignedTextAsRegularText
};
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>,
@@ -488,8 +456,6 @@ 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())
@@ -500,7 +466,9 @@ 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_for_model(part, &mut last_signature, model_version) { // Always emit text as regular text during streaming — we can't
// know yet whether function calls will follow.
if let Some(content) = process_response_part_impl(part, &mut last_signature, SignedTextHandling::SignedTextAsRegularText) {
let message = Message::new( let message = Message::new(
Role::Assistant, Role::Assistant,
chrono::Utc::now().timestamp(), chrono::Utc::now().timestamp(),
@@ -1192,90 +1160,68 @@ mod tests {
async fn test_streaming_with_thought_signature() { async fn test_streaming_with_thought_signature() {
use futures::StreamExt; use futures::StreamExt;
let gemini3_stream = concat!( async fn collect_streaming_text(raw: &str) -> (String, usize) {
r#"data: {"candidates": [{"content": {"role": "model", "#, let lines: Vec<Result<String, anyhow::Error>> =
r#""parts": [{"text": "Begin", "thoughtSignature": "sig123"}]}}], "#, raw.lines().map(|l| Ok(l.to_string())).collect();
r#""modelVersion": "gemini-3-pro"}"#, let stream = Box::pin(futures::stream::iter(lines));
"\n", let mut msg_stream = std::pin::pin!(response_to_streaming_message(stream));
r#"data: {"candidates": [{"content": {"role": "model", "#, let mut text = String::new();
r#""parts": [{"text": " end"}]}}], "modelVersion": "gemini-3-pro"}"# let mut thinking = 0usize;
); while let Some(Ok((message, _))) = msg_stream.next().await {
let lines: Vec<Result<String, anyhow::Error>> = if let Some(msg) = message {
gemini3_stream.lines().map(|l| Ok(l.to_string())).collect(); for c in &msg.content {
let stream = Box::pin(futures::stream::iter(lines)); match c {
let mut message_stream = std::pin::pin!(response_to_streaming_message(stream)); MessageContent::Text(t) => text.push_str(&t.text),
MessageContent::Thinking(_) => thinking += 1,
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()),
_ => {}
} }
} }
(text, thinking)
} }
assert_eq!(thinking_parts, vec!["Begin"]); // First chunk signed
assert_eq!(text_parts, vec![" end"]); let (text, thinking) = collect_streaming_text(concat!(
let gemini25_stream = concat!(
r#"data: {"candidates": [{"content": {"role": "model", "#, r#"data: {"candidates": [{"content": {"role": "model", "#,
r#""parts": [{"text": "Begin", "thoughtSignature": "sig123"}]}}], "#, r#""parts": [{"text": "Hello", "thoughtSignature": "sig1"}]}}], "#,
r#""modelVersion": "gemini-2.5-pro"}"#, r#""modelVersion": "gemini-3-flash-preview"}"#,
"\n", "\n",
r#"data: {"candidates": [{"content": {"role": "model", "#, r#"data: {"candidates": [{"content": {"role": "model", "#,
r#""parts": [{"text": " end"}]}}], "modelVersion": "gemini-2.5-pro"}"# r#""parts": [{"text": " world"}]}}], "modelVersion": "gemini-3-flash-preview"}"#
); ))
let lines: Vec<Result<String, anyhow::Error>> = .await;
gemini25_stream.lines().map(|l| Ok(l.to_string())).collect(); assert_eq!(thinking, 0);
let stream = Box::pin(futures::stream::iter(lines)); assert_eq!(text, "Hello world");
let mut message_stream = std::pin::pin!(response_to_streaming_message(stream));
let mut text_parts = Vec::new(); // Last chunk signed (the reported truncation bug)
let (text, thinking) = collect_streaming_text(concat!(
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#"data: {"candidates": [{"content": {"role": "model", "#,
r#""parts": [{"text": "Begin", "thoughtSignature": "sig123"}]}}]}"#, r#""parts": [{"text": "SECURITY.md: Project"}]}}], "#,
r#""modelVersion": "gemini-3-flash-preview"}"#,
"\n", "\n",
r#"data: {"candidates": [{"content": {"role": "model", "#, r#"data: {"candidates": [{"content": {"role": "model", "#,
r#""parts": [{"text": " end"}]}}]}"# r#""parts": [{"text": " policies.\n\nRead it?", "thoughtSignature": "sig2"}]}}], "#,
); r#""modelVersion": "gemini-3-flash-preview"}"#
let lines: Vec<Result<String, anyhow::Error>> = ))
unknown_stream.lines().map(|l| Ok(l.to_string())).collect(); .await;
let stream = Box::pin(futures::stream::iter(lines)); assert_eq!(thinking, 0);
let mut message_stream = std::pin::pin!(response_to_streaming_message(stream)); assert_eq!(text, "SECURITY.md: Project policies.\n\nRead it?");
let mut text_parts = Vec::new(); // Intermediate chunk signed
let mut thinking_parts = Vec::new(); let (text, thinking) = collect_streaming_text(concat!(
r#"data: {"candidates": [{"content": {"role": "model", "#,
while let Some(result) = message_stream.next().await { r#""parts": [{"text": "one "}]}}], "modelVersion": "gemini-3-flash-preview"}"#,
let (message, _usage) = result.unwrap(); "\n",
if let Some(msg) = message { r#"data: {"candidates": [{"content": {"role": "model", "#,
match msg.content.first() { r#""parts": [{"text": "two ", "thoughtSignature": "sig3"}]}}], "modelVersion": "gemini-3-flash-preview"}"#,
Some(MessageContent::Text(text)) => text_parts.push(text.text.clone()), "\n",
Some(MessageContent::Thinking(t)) => thinking_parts.push(t.thinking.clone()), r#"data: {"candidates": [{"content": {"role": "model", "#,
_ => {} r#""parts": [{"text": "three"}]}}], "modelVersion": "gemini-3-flash-preview"}"#
} ))
} .await;
} assert_eq!(thinking, 0);
assert_eq!(text, "one two three");
assert_eq!(thinking_parts, vec!["Begin"]);
assert_eq!(text_parts, vec![" end"]);
} }
#[tokio::test] #[tokio::test]