Skip to content

Commit 9bdf182

Browse files
committed
feat(google-genai): Support thought signatures for Gemini 3 Pro function calling
Add thought signature support required by Gemini 3 Pro when using function calling with includeThoughts enabled. Thought signatures preserve reasoning context during the internal tool execution loop. Changes from original implementation: - Attach thought signatures to functionCall parts per Google specification (not to text parts or separate empty parts) - Clarify in documentation that validation applies only to the current turn's function calling loop, not to historical conversation messages - Add integration tests validating function calls with signatures Key behaviors: - Signatures are extracted from model responses and stored in message metadata - During internal tool execution, signatures are correctly attached to functionCall parts when sending back function responses - Handles both parallel and sequential function calls correctly: - Parallel: only first functionCall gets the signature (per API spec) - Sequential: each step's first functionCall gets its signature - Previous turn signatures in conversation history are not validated by the API See: https://ai.google.dev/gemini-api/docs/thought-signatures Signed-off-by: Dan Dobrin <dan.dobrin@broadcom.com> Signed-off-by: Mark Pollack <mark.pollack@broadcom.com>
1 parent 2c7b10e commit 9bdf182

File tree

10 files changed

+776
-36
lines changed

10 files changed

+776
-36
lines changed

auto-configurations/models/spring-ai-autoconfigure-model-google-genai/src/test/java/org/springframework/ai/model/google/genai/autoconfigure/chat/GoogleGenAiPropertiesTests.java

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,25 @@ void extendedUsageMetadataDefaultBinding() {
131131
});
132132
}
133133

134+
@Test
135+
void includeThoughtsPropertiesBinding() {
136+
this.contextRunner.withPropertyValues("spring.ai.google.genai.chat.options.include-thoughts=true")
137+
.run(context -> {
138+
GoogleGenAiChatProperties chatProperties = context.getBean(GoogleGenAiChatProperties.class);
139+
assertThat(chatProperties.getOptions().getIncludeThoughts()).isTrue();
140+
});
141+
}
142+
143+
@Test
144+
void includeThoughtsDefaultBinding() {
145+
// Test that defaults are applied when not specified
146+
this.contextRunner.run(context -> {
147+
GoogleGenAiChatProperties chatProperties = context.getBean(GoogleGenAiChatProperties.class);
148+
// Should be null when not set
149+
assertThat(chatProperties.getOptions().getIncludeThoughts()).isNull();
150+
});
151+
}
152+
134153
@Configuration
135154
@EnableConfigurationProperties({ GoogleGenAiConnectionProperties.class, GoogleGenAiChatProperties.class,
136155
GoogleGenAiEmbeddingConnectionProperties.class })

models/spring-ai-google-genai/src/main/java/org/springframework/ai/google/genai/GoogleGenAiChatModel.java

Lines changed: 66 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -281,20 +281,48 @@ else if (message instanceof UserMessage userMessage) {
281281
}
282282
else if (message instanceof AssistantMessage assistantMessage) {
283283
List<Part> parts = new ArrayList<>();
284+
285+
// Check if there are thought signatures to restore.
286+
// Per Google's documentation, thought signatures must be attached to the
287+
// first functionCall part in each step of the current turn.
288+
// See: https://ai.google.dev/gemini-api/docs/thought-signatures
289+
List<byte[]> thoughtSignatures = null;
290+
if (assistantMessage.getMetadata() != null
291+
&& assistantMessage.getMetadata().containsKey("thoughtSignatures")) {
292+
Object signaturesObj = assistantMessage.getMetadata().get("thoughtSignatures");
293+
if (signaturesObj instanceof List) {
294+
thoughtSignatures = new ArrayList<>((List<byte[]>) signaturesObj);
295+
}
296+
}
297+
298+
// Add text part (without thought signature - signatures go on functionCall
299+
// parts)
284300
if (StringUtils.hasText(assistantMessage.getText())) {
285-
parts.add(Part.fromText(assistantMessage.getText()));
301+
parts.add(Part.builder().text(assistantMessage.getText()).build());
286302
}
303+
304+
// Add function call parts with thought signatures attached.
305+
// Per Google's docs: "The first functionCall part in each step of the
306+
// current turn must include its thought_signature."
287307
if (!CollectionUtils.isEmpty(assistantMessage.getToolCalls())) {
288-
parts.addAll(assistantMessage.getToolCalls()
289-
.stream()
290-
.map(toolCall -> Part.builder()
308+
List<AssistantMessage.ToolCall> toolCalls = assistantMessage.getToolCalls();
309+
for (int i = 0; i < toolCalls.size(); i++) {
310+
AssistantMessage.ToolCall toolCall = toolCalls.get(i);
311+
Part.Builder partBuilder = Part.builder()
291312
.functionCall(FunctionCall.builder()
292313
.name(toolCall.name())
293314
.args(parseJsonToMap(toolCall.arguments()))
294-
.build())
295-
.build())
296-
.toList());
315+
.build());
316+
317+
// Attach thought signature to function call part if available
318+
if (thoughtSignatures != null && !thoughtSignatures.isEmpty()) {
319+
partBuilder.thoughtSignature(thoughtSignatures.remove(0));
320+
}
321+
322+
parts.add(partBuilder.build());
323+
}
297324
}
325+
298326
return parts;
299327
}
300328
else if (message instanceof ToolResponseMessage toolResponseMessage) {
@@ -598,8 +626,22 @@ protected List<Generation> responseCandidateToGeneration(Candidate candidate) {
598626
int candidateIndex = candidate.index().orElse(0);
599627
FinishReason candidateFinishReason = candidate.finishReason().orElse(new FinishReason(FinishReason.Known.STOP));
600628

601-
Map<String, Object> messageMetadata = Map.of("candidateIndex", candidateIndex, "finishReason",
602-
candidateFinishReason);
629+
Map<String, Object> messageMetadata = new HashMap<>();
630+
messageMetadata.put("candidateIndex", candidateIndex);
631+
messageMetadata.put("finishReason", candidateFinishReason);
632+
633+
// Extract thought signatures from response parts if present
634+
if (candidate.content().isPresent() && candidate.content().get().parts().isPresent()) {
635+
List<Part> parts = candidate.content().get().parts().get();
636+
List<byte[]> thoughtSignatures = parts.stream()
637+
.filter(part -> part.thoughtSignature().isPresent())
638+
.map(part -> part.thoughtSignature().get())
639+
.toList();
640+
641+
if (!thoughtSignatures.isEmpty()) {
642+
messageMetadata.put("thoughtSignatures", thoughtSignatures);
643+
}
644+
}
603645

604646
ChatGenerationMetadata chatGenerationMetadata = ChatGenerationMetadata.builder()
605647
.finishReason(candidateFinishReason.toString())
@@ -710,10 +752,19 @@ GeminiRequest createGeminiRequest(Prompt prompt) {
710752
if (requestOptions.getPresencePenalty() != null) {
711753
configBuilder.presencePenalty(requestOptions.getPresencePenalty().floatValue());
712754
}
713-
if (requestOptions.getThinkingBudget() != null) {
714-
configBuilder
715-
.thinkingConfig(ThinkingConfig.builder().thinkingBudget(requestOptions.getThinkingBudget()).build());
755+
756+
// Build thinking config if either thinkingBudget or includeThoughts is set
757+
if (requestOptions.getThinkingBudget() != null || requestOptions.getIncludeThoughts() != null) {
758+
ThinkingConfig.Builder thinkingBuilder = ThinkingConfig.builder();
759+
if (requestOptions.getThinkingBudget() != null) {
760+
thinkingBuilder.thinkingBudget(requestOptions.getThinkingBudget());
761+
}
762+
if (requestOptions.getIncludeThoughts() != null) {
763+
thinkingBuilder.includeThoughts(requestOptions.getIncludeThoughts());
764+
}
765+
configBuilder.thinkingConfig(thinkingBuilder.build());
716766
}
767+
717768
if (requestOptions.getLabels() != null && !requestOptions.getLabels().isEmpty()) {
718769
configBuilder.labels(requestOptions.getLabels());
719770
}
@@ -1062,7 +1113,9 @@ public enum ChatModel implements ChatModelDescription {
10621113
* See: <a href=
10631114
* "https://cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/2-5-flash-lite">gemini-2.5-flash-lite</a>
10641115
*/
1065-
GEMINI_2_5_FLASH_LIGHT("gemini-2.5-flash-lite");
1116+
GEMINI_2_5_FLASH_LIGHT("gemini-2.5-flash-lite"),
1117+
1118+
GEMINI_3_PRO_PREVIEW("gemini-3-pro-preview");
10661119

10671120
public final String value;
10681121

models/spring-ai-google-genai/src/main/java/org/springframework/ai/google/genai/GoogleGenAiChatOptions.java

Lines changed: 37 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,19 @@ public class GoogleGenAiChatOptions implements ToolCallingChatOptions {
113113
*/
114114
private @JsonProperty("thinkingBudget") Integer thinkingBudget;
115115

116+
/**
117+
* Optional. Whether to include thoughts in the response.
118+
* When true, thoughts are returned if the model supports them and thoughts are available.
119+
*
120+
* <p><strong>IMPORTANT:</strong> For Gemini 3 Pro with function calling,
121+
* this MUST be set to true to avoid validation errors. Thought signatures
122+
* are automatically propagated in multi-turn conversations to maintain context.
123+
*
124+
* <p>Note: Enabling thoughts increases token usage and API costs.
125+
* This is part of the thinkingConfig in GenerationConfig.
126+
*/
127+
private @JsonProperty("includeThoughts") Boolean includeThoughts;
128+
116129
/**
117130
* Optional. Whether to include extended usage metadata in responses.
118131
* When true, includes thinking tokens, cached content, tool-use tokens, and modality details.
@@ -206,6 +219,7 @@ public static GoogleGenAiChatOptions fromOptions(GoogleGenAiChatOptions fromOpti
206219
options.setInternalToolExecutionEnabled(fromOptions.getInternalToolExecutionEnabled());
207220
options.setToolContext(fromOptions.getToolContext());
208221
options.setThinkingBudget(fromOptions.getThinkingBudget());
222+
options.setIncludeThoughts(fromOptions.getIncludeThoughts());
209223
options.setLabels(fromOptions.getLabels());
210224
options.setIncludeExtendedUsageMetadata(fromOptions.getIncludeExtendedUsageMetadata());
211225
options.setCachedContentName(fromOptions.getCachedContentName());
@@ -357,6 +371,14 @@ public void setThinkingBudget(Integer thinkingBudget) {
357371
this.thinkingBudget = thinkingBudget;
358372
}
359373

374+
public Boolean getIncludeThoughts() {
375+
return this.includeThoughts;
376+
}
377+
378+
public void setIncludeThoughts(Boolean includeThoughts) {
379+
this.includeThoughts = includeThoughts;
380+
}
381+
360382
public Boolean getIncludeExtendedUsageMetadata() {
361383
return this.includeExtendedUsageMetadata;
362384
}
@@ -448,6 +470,7 @@ public boolean equals(Object o) {
448470
&& Objects.equals(this.frequencyPenalty, that.frequencyPenalty)
449471
&& Objects.equals(this.presencePenalty, that.presencePenalty)
450472
&& Objects.equals(this.thinkingBudget, that.thinkingBudget)
473+
&& Objects.equals(this.includeThoughts, that.includeThoughts)
451474
&& Objects.equals(this.maxOutputTokens, that.maxOutputTokens) && Objects.equals(this.model, that.model)
452475
&& Objects.equals(this.responseMimeType, that.responseMimeType)
453476
&& Objects.equals(this.toolCallbacks, that.toolCallbacks)
@@ -460,21 +483,22 @@ public boolean equals(Object o) {
460483
@Override
461484
public int hashCode() {
462485
return Objects.hash(this.stopSequences, this.temperature, this.topP, this.topK, this.candidateCount,
463-
this.frequencyPenalty, this.presencePenalty, this.thinkingBudget, this.maxOutputTokens, this.model,
464-
this.responseMimeType, this.toolCallbacks, this.toolNames, this.googleSearchRetrieval,
465-
this.safetySettings, this.internalToolExecutionEnabled, this.toolContext, this.labels);
486+
this.frequencyPenalty, this.presencePenalty, this.thinkingBudget, this.includeThoughts,
487+
this.maxOutputTokens, this.model, this.responseMimeType, this.toolCallbacks, this.toolNames,
488+
this.googleSearchRetrieval, this.safetySettings, this.internalToolExecutionEnabled, this.toolContext,
489+
this.labels);
466490
}
467491

468492
@Override
469493
public String toString() {
470494
return "GoogleGenAiChatOptions{" + "stopSequences=" + this.stopSequences + ", temperature=" + this.temperature
471495
+ ", topP=" + this.topP + ", topK=" + this.topK + ", frequencyPenalty=" + this.frequencyPenalty
472496
+ ", presencePenalty=" + this.presencePenalty + ", thinkingBudget=" + this.thinkingBudget
473-
+ ", candidateCount=" + this.candidateCount + ", maxOutputTokens=" + this.maxOutputTokens + ", model='"
474-
+ this.model + '\'' + ", responseMimeType='" + this.responseMimeType + '\'' + ", toolCallbacks="
475-
+ this.toolCallbacks + ", toolNames=" + this.toolNames + ", googleSearchRetrieval="
476-
+ this.googleSearchRetrieval + ", safetySettings=" + this.safetySettings + ", labels=" + this.labels
477-
+ '}';
497+
+ ", includeThoughts=" + this.includeThoughts + ", candidateCount=" + this.candidateCount
498+
+ ", maxOutputTokens=" + this.maxOutputTokens + ", model='" + this.model + '\'' + ", responseMimeType='"
499+
+ this.responseMimeType + '\'' + ", toolCallbacks=" + this.toolCallbacks + ", toolNames="
500+
+ this.toolNames + ", googleSearchRetrieval=" + this.googleSearchRetrieval + ", safetySettings="
501+
+ this.safetySettings + ", labels=" + this.labels + '}';
478502
}
479503

480504
@Override
@@ -602,6 +626,11 @@ public Builder thinkingBudget(Integer thinkingBudget) {
602626
return this;
603627
}
604628

629+
public Builder includeThoughts(Boolean includeThoughts) {
630+
this.options.setIncludeThoughts(includeThoughts);
631+
return this;
632+
}
633+
605634
public Builder includeExtendedUsageMetadata(Boolean includeExtendedUsageMetadata) {
606635
this.options.setIncludeExtendedUsageMetadata(includeExtendedUsageMetadata);
607636
return this;

models/spring-ai-google-genai/src/test/java/org/springframework/ai/google/genai/GoogleGenAiChatModelIT.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ void googleSearchToolPro() {
101101
GoogleGenAiChatOptions.builder().model(ChatModel.GEMINI_2_5_PRO).googleSearchRetrieval(true).build());
102102
ChatResponse response = this.chatModel.call(prompt);
103103
assertThat(response.getResult().getOutput().getText()).containsAnyOf("Blackbeard", "Bartholomew", "Calico Jack",
104-
"Anne Bonny");
104+
"Bob", "Anne Bonny");
105105
}
106106

107107
@Test

models/spring-ai-google-genai/src/test/java/org/springframework/ai/google/genai/GoogleGenAiChatModelObservationApiKeyIT.java

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,7 @@ void beforeEach() {
6464
void observationForChatOperation() {
6565

6666
var options = GoogleGenAiChatOptions.builder()
67-
.model(GoogleGenAiChatModel.ChatModel.GEMINI_2_0_FLASH.getValue())
67+
.model(GoogleGenAiChatModel.ChatModel.GEMINI_3_PRO_PREVIEW.getValue())
6868
.temperature(0.7)
6969
.stopSequences(List.of("this-is-the-end"))
7070
.maxOutputTokens(2048)
@@ -86,7 +86,7 @@ void observationForChatOperation() {
8686
void observationForStreamingOperation() {
8787

8888
var options = GoogleGenAiChatOptions.builder()
89-
.model(GoogleGenAiChatModel.ChatModel.GEMINI_2_0_FLASH.getValue())
89+
.model(GoogleGenAiChatModel.ChatModel.GEMINI_3_PRO_PREVIEW.getValue())
9090
.temperature(0.7)
9191
.stopSequences(List.of("this-is-the-end"))
9292
.maxOutputTokens(2048)
@@ -126,7 +126,7 @@ private void validate(ChatResponseMetadata responseMetadata) {
126126
AiProvider.GOOGLE_GENAI_AI.value())
127127
.hasLowCardinalityKeyValue(
128128
ChatModelObservationDocumentation.LowCardinalityKeyNames.REQUEST_MODEL.asString(),
129-
GoogleGenAiChatModel.ChatModel.GEMINI_2_0_FLASH.getValue())
129+
GoogleGenAiChatModel.ChatModel.GEMINI_3_PRO_PREVIEW.getValue())
130130
.hasHighCardinalityKeyValue(
131131
ChatModelObservationDocumentation.HighCardinalityKeyNames.REQUEST_MAX_TOKENS.asString(), "2048")
132132
.hasHighCardinalityKeyValue(
@@ -174,8 +174,9 @@ public GoogleGenAiChatModel vertexAiEmbedding(Client genAiClient, TestObservatio
174174
return GoogleGenAiChatModel.builder()
175175
.genAiClient(genAiClient)
176176
.observationRegistry(observationRegistry)
177-
.defaultOptions(
178-
GoogleGenAiChatOptions.builder().model(GoogleGenAiChatModel.ChatModel.GEMINI_2_0_FLASH).build())
177+
.defaultOptions(GoogleGenAiChatOptions.builder()
178+
.model(GoogleGenAiChatModel.ChatModel.GEMINI_3_PRO_PREVIEW)
179+
.build())
179180
.build();
180181
}
181182

models/spring-ai-google-genai/src/test/java/org/springframework/ai/google/genai/GoogleGenAiRetryTests.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ public void setUp() {
6161
GoogleGenAiChatOptions.builder()
6262
.temperature(0.7)
6363
.topP(1.0)
64-
.model(GoogleGenAiChatModel.ChatModel.GEMINI_2_0_FLASH.getValue())
64+
.model(GoogleGenAiChatModel.ChatModel.GEMINI_3_PRO_PREVIEW.getValue())
6565
.build(),
6666
this.retryTemplate);
6767

0 commit comments

Comments
 (0)