Skip to content

Commit 2ebf81f

Browse files
committed
perf(google): avoid accumulating thoughtSignatures across conversation history
Only include signatures on the last assistant message and subsequent user messages. Reduces payload size from O(n) to O(1) for Gemini thinking models. Signed-off-by: rabi <ramishra@redhat.com>
1 parent 637049c commit 2ebf81f

1 file changed

Lines changed: 58 additions & 16 deletions

File tree

crates/goose/src/providers/formats/google.rs

Lines changed: 58 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ pub fn get_thought_signature(metadata: &Option<ProviderMetadata>) -> Option<&str
3131

3232
/// Convert internal Message format to Google's API message specification
3333
pub fn format_messages(messages: &[Message]) -> Vec<Value> {
34-
messages
34+
let filtered: Vec<_> = messages
3535
.iter()
3636
.filter(|m| m.is_agent_visible())
3737
.filter(|message| {
@@ -42,12 +42,28 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
4242
)
4343
})
4444
})
45-
.map(|message| {
45+
.collect();
46+
47+
let last_assistant_idx = filtered
48+
.iter()
49+
.enumerate()
50+
.filter(|(_, m)| m.role != Role::User)
51+
.map(|(i, _)| i)
52+
.next_back();
53+
54+
filtered
55+
.iter()
56+
.enumerate()
57+
.map(|(idx, message)| {
4658
let role = if message.role == Role::User {
4759
"user"
4860
} else {
4961
"model"
5062
};
63+
let include_signature = match last_assistant_idx {
64+
Some(last_idx) => idx >= last_idx,
65+
None => false,
66+
};
5167
let mut parts = Vec::new();
5268
for message_content in message.content.iter() {
5369
match message_content {
@@ -74,8 +90,13 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
7490
let mut part = Map::new();
7591
part.insert("functionCall".to_string(), json!(function_call_part));
7692

77-
if let Some(signature) = get_thought_signature(&request.metadata) {
78-
part.insert(THOUGHT_SIGNATURE_KEY.to_string(), json!(signature));
93+
if include_signature {
94+
if let Some(signature) = get_thought_signature(&request.metadata) {
95+
part.insert(
96+
THOUGHT_SIGNATURE_KEY.to_string(),
97+
json!(signature),
98+
);
99+
}
79100
}
80101

81102
parts.push(json!(part));
@@ -144,11 +165,15 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
144165
"functionResponse".to_string(),
145166
json!(function_response),
146167
);
147-
if let Some(signature) = get_thought_signature(&response.metadata) {
148-
part.insert(
149-
THOUGHT_SIGNATURE_KEY.to_string(),
150-
json!(signature),
151-
);
168+
if include_signature {
169+
if let Some(signature) =
170+
get_thought_signature(&response.metadata)
171+
{
172+
part.insert(
173+
THOUGHT_SIGNATURE_KEY.to_string(),
174+
json!(signature),
175+
);
176+
}
152177
}
153178
parts.push(json!(part));
154179
}
@@ -164,11 +189,15 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
164189
"functionResponse".to_string(),
165190
json!(function_response),
166191
);
167-
if let Some(signature) = get_thought_signature(&response.metadata) {
168-
part.insert(
169-
THOUGHT_SIGNATURE_KEY.to_string(),
170-
json!(signature),
171-
);
192+
if include_signature {
193+
if let Some(signature) =
194+
get_thought_signature(&response.metadata)
195+
{
196+
part.insert(
197+
THOUGHT_SIGNATURE_KEY.to_string(),
198+
json!(signature),
199+
);
200+
}
172201
}
173202
parts.push(json!(part));
174203
}
@@ -177,7 +206,9 @@ pub fn format_messages(messages: &[Message]) -> Vec<Value> {
177206
MessageContent::Thinking(thinking) => {
178207
let mut part = Map::new();
179208
part.insert("text".to_string(), json!(thinking.thinking));
180-
part.insert("thoughtSignature".to_string(), json!(thinking.signature));
209+
if include_signature {
210+
part.insert("thoughtSignature".to_string(), json!(thinking.signature));
211+
}
181212
parts.push(json!(part));
182213
}
183214

@@ -1230,10 +1261,21 @@ mod tests {
12301261
Ok(tool_result("output")),
12311262
req1.metadata.as_ref(),
12321263
);
1233-
let google_out = format_messages(&[native, tool_response]);
1264+
let google_out = format_messages(&[native.clone(), tool_response.clone()]);
12341265
assert_eq!(google_out[0]["parts"][0]["thoughtSignature"], SIG);
12351266
assert_eq!(google_out[1]["parts"][0]["thoughtSignature"], SIG);
12361267

1268+
let second_assistant =
1269+
Message::assistant().with_thinking("More thinking".to_string(), "sig_456".to_string());
1270+
let google_multi = format_messages(&[native, tool_response, second_assistant]);
1271+
assert!(google_multi[0]["parts"][0]
1272+
.get("thoughtSignature")
1273+
.is_none());
1274+
assert!(google_multi[1]["parts"][0]
1275+
.get("thoughtSignature")
1276+
.is_none());
1277+
assert_eq!(google_multi[2]["parts"][0]["thoughtSignature"], "sig_456");
1278+
12371279
// Text-only response WITH signature but WITHOUT function calls should be regular text
12381280
// (per original behavior: thinking is only when reasoning before tool calls)
12391281
let final_response_with_sig =

0 commit comments

Comments
 (0)