Skip to content

fix: preserve tool responses and thought signatures in Spring AI message conversion #1459

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Aug 26, 2026
+198 −34
Merged

fix: preserve tool responses and thought signatures in Spring AI message conversion #1459

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 58 additions & 15 deletions contrib/spring-ai/src/main/java/com/google/adk/models/springai/MessageConverter.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import com.google.adk.models.LlmResponse;
import com.google.genai.types.Content;
import com.google.genai.types.FunctionCall;
import com.google.genai.types.FunctionResponse;
import com.google.genai.types.GenerateContentResponseUsageMetadata;
import com.google.genai.types.Part;
import java.net.URI;
Expand Down Expand Up @@ -54,18 +55,23 @@
* <ul>
* <li>Text content in all message types
* <li>Tool/function calls in assistant messages
* <li>Tool/function responses as {@link ToolResponseMessage}s
* <li>System instructions and configuration options
* </ul>
*
* <p>Note: Media attachments and tool responses are currently not supported due to Spring AI 1.1.0
* API limitations (protected/private constructors). These will be added once Spring AI provides
* public APIs for these features.
*/
public class MessageConverter {

private static final TypeReference<Map<String, Object>> MAP_TYPE_REFERENCE =
new TypeReference<>() {};

/**
* Message-metadata key the Spring AI Google GenAI provider uses to carry Gemini thought
* signatures. These must be replayed on the model turn of the next request or tool calling fails,
* so they are preserved across the ADK round-trip. See
* https://ai.google.dev/gemini-api/docs/thought-signatures.
*/
private static final String THOUGHT_SIGNATURES_KEY = "thoughtSignatures";

private final ObjectMapper objectMapper;
private final ToolConverter toolConverter;
private final ConfigMapper configMapper;
Expand Down Expand Up @@ -254,17 +260,19 @@ private List<Message> toSpringAiMessages(Content content) {

private List<Message> handleUserContent(Content content) {
StringBuilder textBuilder = new StringBuilder();
List<ToolResponseMessage> toolResponseMessages = new ArrayList<>();
List<ToolResponseMessage.ToolResponse> toolResponses = new ArrayList<>();
List<Media> mediaList = new ArrayList<>();

for (Part part : content.parts().orElse(List.of())) {
if (part.text().isPresent()) {
textBuilder.append(part.text().get());
} else if (part.functionResponse().isPresent()) {
// TODO: Spring AI 1.1.0 ToolResponseMessage constructors are protected
// For now, we skip tool responses in user messages
// This will need to be addressed in a future update when Spring AI provides
// a public API for creating ToolResponseMessage
FunctionResponse functionResponse = part.functionResponse().get();
toolResponses.add(
new ToolResponseMessage.ToolResponse(
functionResponse.id().orElse(""),
functionResponse.name().orElse(""),
toJson(functionResponse.response().orElse(Map.of()))));
} else if (part.inlineData().isPresent()) {
// Handle inline media data (images, audio, video, etc.)
com.google.genai.types.Blob blob = part.inlineData().get();
Expand Down Expand Up @@ -298,15 +306,24 @@ private List<Message> handleUserContent(Content content) {
}

List<Message> messages = new ArrayList<>();
messages.add(UserMessage.builder().text(textBuilder.toString()).media(mediaList).build());
messages.addAll(toolResponseMessages);
String text = textBuilder.toString();
// Emit a UserMessage for any text/media, or for an otherwise-empty turn; but when the turn only
// carries function responses, emit just the ToolResponseMessage so the request does not end on
// the model's tool-call turn (which the backend rejects).
if (!text.isEmpty() || !mediaList.isEmpty() || toolResponses.isEmpty()) {
messages.add(UserMessage.builder().text(text).media(mediaList).build());
}
if (!toolResponses.isEmpty()) {
messages.add(ToolResponseMessage.builder().responses(toolResponses).build());
}

return messages;
}

private AssistantMessage handleAssistantContent(Content content) {
StringBuilder textBuilder = new StringBuilder();
List<AssistantMessage.ToolCall> toolCalls = new ArrayList<>();
List<byte[]> thoughtSignatures = new ArrayList<>();

for (Part part : content.parts().orElse(List.of())) {
if (part.text().isPresent()) {
Expand All @@ -323,15 +340,21 @@ private AssistantMessage handleAssistantContent(Content content) {
.name()
.orElseThrow(() -> new IllegalStateException("Function call name is missing")),
toJson(functionCall.args().orElse(Map.of()))));
part.thoughtSignature().ifPresent(thoughtSignatures::add);
}
}

String text = textBuilder.toString();
if (toolCalls.isEmpty()) {
return new AssistantMessage(text);
} else {
return AssistantMessage.builder().content(text).toolCalls(toolCalls).build();
}
Map<String, Object> properties =
thoughtSignatures.isEmpty() ? Map.of() : Map.of(THOUGHT_SIGNATURES_KEY, thoughtSignatures);
return AssistantMessage.builder()
.content(text)
.properties(properties)
.toolCalls(toolCalls)
.build();
}

private SystemMessage handleSystemContent(Content content) {
Expand Down Expand Up @@ -440,6 +463,11 @@ private Content convertAssistantMessageToContent(AssistantMessage assistantMessa
parts.add(Part.fromText(assistantMessage.getText()));
}

// Gemini thinking models return a thought signature per function-call part; keep them on the
// parts so they can be replayed on the next turn (required for tool calling).
List<byte[]> thoughtSignatures = extractThoughtSignatures(assistantMessage);
int functionCallIndex = 0;

// Add tool calls
for (AssistantMessage.ToolCall toolCall : assistantMessage.getToolCalls()) {
if ("function".equals(toolCall.type())) {
Expand All @@ -451,8 +479,13 @@ private Content convertAssistantMessageToContent(AssistantMessage assistantMessa
FunctionCall functionCall =
FunctionCall.builder().id(toolCall.id()).name(toolCall.name()).args(args).build();

// Create Part with the FunctionCall (preserves ID)
parts.add(Part.builder().functionCall(functionCall).build());
// Create Part with the FunctionCall (preserves ID), reattaching any thought signature.
Part.Builder partBuilder = Part.builder().functionCall(functionCall);
if (functionCallIndex < thoughtSignatures.size()) {
partBuilder.thoughtSignature(thoughtSignatures.get(functionCallIndex));
}
parts.add(partBuilder.build());
functionCallIndex++;
} catch (JsonProcessingException e) {
throw MessageConversionException.jsonParsingFailed("tool call arguments", e);
}
Expand All @@ -462,6 +495,16 @@ private Content convertAssistantMessageToContent(AssistantMessage assistantMessa
return Content.builder().role("model").parts(parts).build();
}

@SuppressWarnings("unchecked")
private static List<byte[]> extractThoughtSignatures(AssistantMessage assistantMessage) {
Map<String, Object> metadata = assistantMessage.getMetadata();
if (metadata == null) {
return List.of();
}
Object signatures = metadata.get(THOUGHT_SIGNATURES_KEY);
return (signatures instanceof List) ? (List<byte[]>) signatures : List.of();
}

private String toJson(Object object) {
try {
return objectMapper.writeValueAsString(object);
Expand Down
159 changes: 140 additions & 19 deletions contrib/spring-ai/src/test/java/com/google/adk/models/springai/MessageConverterTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
import com.google.genai.types.GenerateContentConfig;
import com.google.genai.types.Part;
import com.google.genai.types.Schema;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;
import java.util.Optional;
Expand All @@ -37,6 +38,7 @@
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.ToolResponseMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.metadata.ChatResponseMetadata;
import org.springframework.ai.chat.metadata.DefaultUsage;
Expand Down Expand Up @@ -154,14 +156,103 @@ void testToLlmPromptWithFunctionCall() {
}

@Test
void testToLlmPromptWithFunctionResponse() {
// TODO: This test is currently limited due to Spring AI 1.1.0 API constraints
// ToolResponseMessage constructors are protected, so function responses are skipped
// Once Spring AI provides public APIs, this test should be updated to verify:
// 1. ToolResponseMessage is created
// 2. Tool response data is properly converted
// 3. Tool call IDs are preserved
@SuppressWarnings("unchecked")
void testToLlmPromptReplaysThoughtSignatureOnFunctionCall() {
// A thought signature on the model's function-call part must be surfaced in the assistant
// message metadata so the Spring AI Google provider can replay it on the next request.
byte[] signature = "thought-signature".getBytes(StandardCharsets.UTF_8);
FunctionCall functionCall =
FunctionCall.builder()
.name("get_weather")
.args(Map.of("location", "San Francisco"))
.id("call_123")
.build();
Content assistantContent =
Content.builder()
.role("model")
.parts(Part.builder().functionCall(functionCall).thoughtSignature(signature).build())
.build();

LlmRequest request = LlmRequest.builder().contents(List.of(assistantContent)).build();

Prompt prompt = messageConverter.toLlmPrompt(request);

Message message = prompt.getInstructions().get(0);
assertThat(message).isInstanceOf(AssistantMessage.class);
AssistantMessage assistantMessage = (AssistantMessage) message;
assertThat(assistantMessage.getToolCalls()).hasSize(1);
Object signatures = assistantMessage.getMetadata().get("thoughtSignatures");
assertThat(signatures).isInstanceOf(List.class);
assertThat((List<byte[]>) signatures).containsExactly(signature);
}

@Test
void testToLlmResponsePreservesThoughtSignatureOnFunctionCall() {
// A thought signature carried in the provider's assistant-message metadata must be attached to
// the ADK function-call part so it survives the round-trip back to the model.
byte[] signature = "thought-signature".getBytes(StandardCharsets.UTF_8);
AssistantMessage assistantMessage =
AssistantMessage.builder()
.content("")
.properties(Map.of("thoughtSignatures", List.of(signature)))
.toolCalls(
List.of(
new AssistantMessage.ToolCall(
"call_123", "function", "get_weather", "{\"location\":\"San Francisco\"}")))
.build();
ChatResponse chatResponse = new ChatResponse(List.of(new Generation(assistantMessage)));

LlmResponse response = messageConverter.toLlmResponse(chatResponse);

List<Part> parts = response.content().orElseThrow().parts().orElseThrow();
Part functionCallPart =
parts.stream().filter(part -> part.functionCall().isPresent()).findFirst().orElseThrow();
assertThat(functionCallPart.thoughtSignature()).isPresent();
assertThat(functionCallPart.thoughtSignature().get()).isEqualTo(signature);
}

@Test
@SuppressWarnings("unchecked")
void testToLlmPromptReplaysThoughtSignaturesForParallelToolCalls() {
// Parallel tool calls each carry their own thought signature; order must be preserved.
byte[] signatureA = "signature-a".getBytes(StandardCharsets.UTF_8);
byte[] signatureB = "signature-b".getBytes(StandardCharsets.UTF_8);
Content assistantContent =
Content.builder()
.role("model")
.parts(
Part.builder()
.functionCall(
FunctionCall.builder()
.name("get_weather")
.args(Map.of("location", "San Francisco"))
.id("call_a")
.build())
.thoughtSignature(signatureA)
.build(),
Part.builder()
.functionCall(
FunctionCall.builder()
.name("get_time")
.args(Map.of("location", "New York"))
.id("call_b")
.build())
.thoughtSignature(signatureB)
.build())
.build();

LlmRequest request = LlmRequest.builder().contents(List.of(assistantContent)).build();

Prompt prompt = messageConverter.toLlmPrompt(request);

AssistantMessage assistantMessage = (AssistantMessage) prompt.getInstructions().get(0);
assertThat(assistantMessage.getToolCalls()).hasSize(2);
Object signatures = assistantMessage.getMetadata().get("thoughtSignatures");
assertThat((List<byte[]>) signatures).containsExactly(signatureA, signatureB);
}

@Test
void testToLlmPromptWithFunctionResponse() {
FunctionResponse functionResponse =
FunctionResponse.builder()
.name("get_weather")
Expand All @@ -174,29 +265,59 @@ void testToLlmPromptWithFunctionResponse() {
.role("user")
.parts(
Part.fromText("What's the weather?"),
Part.fromFunctionResponse(
functionResponse.name().orElse(""),
functionResponse.response().orElse(Map.of())))
Part.builder().functionResponse(functionResponse).build())
.build();

LlmRequest request = LlmRequest.builder().contents(List.of(userContent)).build();

Prompt prompt = messageConverter.toLlmPrompt(request);

// Currently only UserMessage is created (function response is skipped)
assertThat(prompt.getInstructions()).hasSize(1);
// A user text part plus a function response yield a UserMessage and a ToolResponseMessage.
assertThat(prompt.getInstructions()).hasSize(2);

Message userMessage = prompt.getInstructions().get(0);
assertThat(userMessage).isInstanceOf(UserMessage.class);
assertThat(((UserMessage) userMessage).getText()).isEqualTo("What's the weather?");

// When Spring AI provides public API for ToolResponseMessage, uncomment:
// Message toolResponseMessage = prompt.getInstructions().get(1);
// assertThat(toolResponseMessage).isInstanceOf(ToolResponseMessage.class);
// ToolResponseMessage toolResponse = (ToolResponseMessage) toolResponseMessage;
// assertThat(toolResponse.getResponses()).hasSize(1);
// ToolResponseMessage.ToolResponse response = toolResponse.getResponses().get(0);
// assertThat(response.name()).isEqualTo("get_weather");
Message toolMessage = prompt.getInstructions().get(1);
assertThat(toolMessage).isInstanceOf(ToolResponseMessage.class);
List<ToolResponseMessage.ToolResponse> responses =
((ToolResponseMessage) toolMessage).getResponses();
assertThat(responses).hasSize(1);
ToolResponseMessage.ToolResponse response = responses.get(0);
assertThat(response.id()).isEqualTo("call_123");
assertThat(response.name()).isEqualTo("get_weather");
assertThat(response.responseData()).contains("72°F").contains("sunny");
}

@Test
void testToLlmPromptWithFunctionResponseOnly() {
// A function-response-only turn must become a ToolResponseMessage, not an empty UserMessage:
// otherwise the request ends on the model's tool-call turn and the backend rejects it.
FunctionResponse functionResponse =
FunctionResponse.builder()
.name("get_weather")
.response(Map.of("temperature", "72°F"))
.id("call_123")
.build();

Content toolResponseContent =
Content.builder()
.role("user")
.parts(Part.builder().functionResponse(functionResponse).build())
.build();

LlmRequest request = LlmRequest.builder().contents(List.of(toolResponseContent)).build();

Prompt prompt = messageConverter.toLlmPrompt(request);

assertThat(prompt.getInstructions()).hasSize(1);
Message message = prompt.getInstructions().get(0);
assertThat(message).isInstanceOf(ToolResponseMessage.class);
ToolResponseMessage.ToolResponse response =
((ToolResponseMessage) message).getResponses().get(0);
assertThat(response.id()).isEqualTo("call_123");
assertThat(response.name()).isEqualTo("get_weather");
}

@Test
Expand Down
Loading