Add builder to ToolResponseMessage

Fixes #3243

Signed-off-by: Eric Bottard <eric.bottard@broadcom.com>

Author: ericbottard

auto-configurations/models/chat/memory/repository/spring-ai-autoconfigure-model-chat-memory-repository-neo4j/src/test/java/org/springframework/ai/model/chat/memory/repository/neo4j/autoconfigure/Neo4jChatMemoryRepositoryAutoConfigurationIT.java

@@ -115,10 +115,11 @@ void addAndGet() {
 			assertThat(((UserMessage) messages.get(0)).getMedia()).usingRecursiveFieldByFieldElementComparator()
 				.isEqualTo(media);
 			memory.deleteByConversationId(sessionId);
-			ToolResponseMessage toolResponseMessage = new ToolResponseMessage(
-					List.of(new ToolResponse("id", "name", "responseData"),
-							new ToolResponse("id2", "name2", "responseData2")),
-					Map.of("id", "id", "metadataKey", "metadata"));
+			ToolResponseMessage toolResponseMessage = ToolResponseMessage.builder()
+				.responses(List.of(new ToolResponse("id", "name", "responseData"),
+						new ToolResponse("id2", "name2", "responseData2")))
+				.metadata(Map.of("id", "id", "metadataKey", "metadata"))
+				.build();
 			memory.saveAll(sessionId, List.of(toolResponseMessage));
 			messages = memory.findByConversationId(sessionId);
 			assertThat(messages.size()).isEqualTo(1);

memory/repository/spring-ai-model-chat-memory-repository-cassandra/src/main/java/org/springframework/ai/chat/memory/repository/cassandra/CassandraChatMemoryRepository.java

@@ -216,7 +216,7 @@ private Message getMessage(UdtValue udt) {
 				return SystemMessage.builder().text(content).metadata(props).build();
 			case TOOL:
 				// todo – persist ToolResponse somehow
-				return new ToolResponseMessage(List.of(), props);
+				return ToolResponseMessage.builder().responses(List.of()).metadata(props).build();
 			default:
 				throw new IllegalStateException(
 						String.format("unknown message type %s", udt.getString(this.conf.messageUdtTypeColumn)));

memory/repository/spring-ai-model-chat-memory-repository-cosmos-db/src/main/java/org/springframework/ai/chat/memory/repository/cosmosdb/CosmosDBChatMemoryRepository.java

@@ -235,7 +235,7 @@ private Message mapToMessage(Map<String, Object> doc) {
 			case ASSISTANT -> new AssistantMessage(content, metadata);
 			case USER -> UserMessage.builder().text(content).metadata(metadata).build();
 			case SYSTEM -> SystemMessage.builder().text(content).metadata(metadata).build();
-			case TOOL -> new ToolResponseMessage(List.of(), metadata);
+			case TOOL -> ToolResponseMessage.builder().responses(List.of()).metadata(metadata).build();
 			default -> throw new IllegalStateException(String.format("Unknown message type: %s", messageTypeStr));
 		};
 	}

memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java

@@ -148,7 +148,7 @@ public Message mapRow(ResultSet rs, int i) throws SQLException {
 				// The content is always stored empty for ToolResponseMessages.
 				// If we want to capture the actual content, we need to extend
 				// AddBatchPreparedStatement to support it.
-				case TOOL -> new ToolResponseMessage(List.of());
+				case TOOL -> ToolResponseMessage.builder().responses(List.of()).build();
 			};
 		}
 

memory/repository/spring-ai-model-chat-memory-repository-neo4j/src/main/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepository.java

@@ -35,6 +35,7 @@
 import org.springframework.ai.chat.messages.MessageType;
 import org.springframework.ai.chat.messages.SystemMessage;
 import org.springframework.ai.chat.messages.ToolResponseMessage;
+import org.springframework.ai.chat.messages.ToolResponseMessage.ToolResponse;
 import org.springframework.ai.chat.messages.UserMessage;
 import org.springframework.ai.content.Media;
 import org.springframework.ai.content.MediaContent;
@@ -172,12 +173,12 @@ public Neo4jChatMemoryRepositoryConfig getConfig() {
 
 	private Message buildToolMessage(org.neo4j.driver.Record record) {
 		Message message;
-		message = new ToolResponseMessage(record.get("toolResponses").asList(v -> {
+		message = ToolResponseMessage.builder().responses(record.get("toolResponses").asList(v -> {
 			Map<String, Object> trMap = v.asMap();
-			return new ToolResponseMessage.ToolResponse((String) trMap.get(ToolResponseAttributes.ID.getValue()),
+			return new ToolResponse((String) trMap.get(ToolResponseAttributes.ID.getValue()),
 					(String) trMap.get(ToolResponseAttributes.NAME.getValue()),
 					(String) trMap.get(ToolResponseAttributes.RESPONSE_DATA.getValue()));
-		}), record.get("metadata").asMap());
+		})).metadata(record.get("metadata").asMap()).build();
 		return message;
 	}
 

memory/repository/spring-ai-model-chat-memory-repository-neo4j/src/test/java/org/springframework/ai/chat/memory/repository/neo4j/Neo4jChatMemoryRepositoryIT.java

@@ -130,7 +130,9 @@ void saveAndFindMultipleMessages() {
 		List<Message> messages = List.of(new AssistantMessage("Message from assistant - " + conversationId),
 				new UserMessage("Message from user - " + conversationId),
 				new SystemMessage("Message from system - " + conversationId),
-				new ToolResponseMessage(List.of(new ToolResponse("id", "name", "responseData"))));
+				ToolResponseMessage.builder()
+					.responses(List.of(new ToolResponse("id", "name", "responseData")))
+					.build());
 
 		this.chatMemoryRepository.saveAll(conversationId, messages);
 		List<Message> retrievedMessages = this.chatMemoryRepository.findByConversationId(conversationId);
@@ -285,9 +287,11 @@ void handleAssistantMessageWithToolCalls() {
 	void handleToolResponseMessage() {
 		var conversationId = UUID.randomUUID().toString();
 
-		ToolResponseMessage toolResponseMessage = new ToolResponseMessage(List
-			.of(new ToolResponse("id1", "name1", "responseData1"), new ToolResponse("id2", "name2", "responseData2")),
-				Map.of("metadataKey", "metadataValue"));
+		ToolResponseMessage toolResponseMessage = ToolResponseMessage.builder()
+			.responses(List.of(new ToolResponse("id1", "name1", "responseData1"),
+					new ToolResponse("id2", "name2", "responseData2")))
+			.metadata(Map.of("metadataKey", "metadataValue"))
+			.build();
 
 		this.chatMemoryRepository.saveAll(conversationId, List.<Message>of(toolResponseMessage));
 
@@ -408,7 +412,9 @@ private Message createMessageByType(String content, MessageType messageType) {
 			case ASSISTANT -> new AssistantMessage(content);
 			case USER -> new UserMessage(content);
 			case SYSTEM -> new SystemMessage(content);
-			case TOOL -> new ToolResponseMessage(List.of(new ToolResponse("id", "name", "responseData")));
+			case TOOL -> ToolResponseMessage.builder()
+				.responses(List.of(new ToolResponse("id", "name", "responseData")))
+				.build();
 		};
 	}
 

models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiChatCompletionRequestTests.java

@@ -199,7 +199,9 @@ void createChatCompletionMessagesWithToolResponseMessage() {
 		var toolResponse1 = createToolResponse(1);
 		var toolResponse2 = createToolResponse(2);
 		var toolResponse3 = createToolResponse(3);
-		var toolResponseMessage = new ToolResponseMessage(List.of(toolResponse1, toolResponse2, toolResponse3));
+		var toolResponseMessage = ToolResponseMessage.builder()
+			.responses(List.of(toolResponse1, toolResponse2, toolResponse3))
+			.build();
 		var prompt = createPrompt(toolResponseMessage);
 		var chatCompletionRequest = this.chatModel.createRequest(prompt, false);
 		var chatCompletionMessages = chatCompletionRequest.messages();
@@ -212,7 +214,7 @@ void createChatCompletionMessagesWithToolResponseMessage() {
 	@Test
 	void createChatCompletionMessagesWithInvalidToolResponseMessage() {
 		var toolResponse = new ToolResponseMessage.ToolResponse(null, null, null);
-		var toolResponseMessage = new ToolResponseMessage(List.of(toolResponse));
+		var toolResponseMessage = ToolResponseMessage.builder().responses(List.of(toolResponse)).build();
 		var prompt = createPrompt(toolResponseMessage);
 		assertThatThrownBy(() -> this.chatModel.createRequest(prompt, false))
 			.isInstanceOf(IllegalArgumentException.class)

models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/OllamaChatRequestTests.java

@@ -25,6 +25,7 @@
 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.ToolResponseMessage.ToolResponse;
 import org.springframework.ai.chat.messages.UserMessage;
 import org.springframework.ai.chat.prompt.ChatOptions;
 import org.springframework.ai.chat.prompt.Prompt;
@@ -256,11 +257,10 @@ private static List<Message> createMessagesWithAllMessageTypes() {
 		var systemMessage = new SystemMessage("Test system message");
 		var userMessage = new UserMessage("Test user message");
 		// @formatter:off
-		var toolResponseMessage = new ToolResponseMessage(List.of(
-				new ToolResponseMessage.ToolResponse("tool1", "Tool 1", "Test tool response 1"),
-				new ToolResponseMessage.ToolResponse("tool2", "Tool 2", "Test tool response 2"),
-				new ToolResponseMessage.ToolResponse("tool3", "Tool 3", "Test tool response 3"))
-		);
+		var toolResponseMessage = ToolResponseMessage.builder().responses(List.of(
+				new ToolResponse("tool1", "Tool 1", "Test tool response 1"),
+				new ToolResponse("tool2", "Tool 2", "Test tool response 2"),
+				new ToolResponse("tool3", "Tool 3", "Test tool response 3"))).build();
 		// @formatter:on
 		var assistantMessage = new AssistantMessage("Test assistant message");
 

spring-ai-model/src/main/java/org/springframework/ai/chat/messages/ToolResponseMessage.java

@@ -25,21 +25,34 @@
  * application.
  *
  * @author Christian Tzolov
+ * @author Eric Bottard
  * @since 1.0.0
  */
 public class ToolResponseMessage extends AbstractMessage {
 
 	protected final List<ToolResponse> responses;
 
+	/**
+	 * @deprecated in favor of using {@link ToolResponseMessage.Builder}
+	 */
+	@Deprecated
 	public ToolResponseMessage(List<ToolResponse> responses) {
 		this(responses, Map.of());
 	}
 
+	/**
+	 * @deprecated in favor of using {@link ToolResponseMessage.Builder}
+	 */
+	@Deprecated
 	public ToolResponseMessage(List<ToolResponse> responses, Map<String, Object> metadata) {
 		super(MessageType.TOOL, "", metadata);
 		this.responses = responses;
 	}
 
+	public static Builder builder() {
+		return new Builder();
+	}
+
 	public List<ToolResponse> getResponses() {
 		return this.responses;
 	}
@@ -73,4 +86,29 @@ public record ToolResponse(String id, String name, String responseData) {
 
 	}
 
+	public static final class Builder {
+
+		private List<ToolResponse> responses;
+
+		private Map<String, Object> metadata = Map.of();
+
+		private Builder() {
+		}
+
+		public Builder responses(List<ToolResponse> responses) {
+			this.responses = responses;
+			return this;
+		}
+
+		public Builder metadata(Map<String, Object> metadata) {
+			this.metadata = metadata;
+			return this;
+		}
+
+		public ToolResponseMessage build() {
+			return new ToolResponseMessage(this.responses, this.metadata);
+		}
+
+	}
+
 }

spring-ai-model/src/main/java/org/springframework/ai/chat/prompt/Prompt.java

@@ -184,8 +184,10 @@ else if (message instanceof AssistantMessage assistantMessage) {
 					.build());
 			}
 			else if (message instanceof ToolResponseMessage toolResponseMessage) {
-				messagesCopy.add(new ToolResponseMessage(new ArrayList<>(toolResponseMessage.getResponses()),
-						new HashMap<>(toolResponseMessage.getMetadata())));
+				messagesCopy.add(ToolResponseMessage.builder()
+					.responses(new ArrayList<>(toolResponseMessage.getResponses()))
+					.metadata(new HashMap<>(toolResponseMessage.getMetadata()))
+					.build());
 			}
 			else {
 				throw new IllegalArgumentException("Unsupported message type: " + message.getClass().getName());