Skip to content

Commit ffd0169

Browse files
JM-LabJM-Lab
committed
Refactor affected usages after AssistantMessage and ToolResponseMessage update
1 parent e46dccf commit ffd0169

File tree

17 files changed

+113
-64
lines changed

17 files changed

+113
-64
lines changed

‎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‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
import static org.assertj.core.api.Assertions.assertThat;
4747

4848
/**
49+
* @author Jemin Huh
4950
* @author Mick Semb Wever
5051
* @author Jihoon Kim
5152
* @author Enrico Rampazzo
@@ -83,8 +84,10 @@ void addAndGet() {
8384
memory.deleteByConversationId(sessionId);
8485
assertThat(memory.findByConversationId(sessionId)).isEmpty();
8586

86-
AssistantMessage assistantMessage = new AssistantMessage("test answer", Map.of(),
87-
List.of(new AssistantMessage.ToolCall("id", "type", "name", "arguments")));
87+
AssistantMessage assistantMessage = AssistantMessage.builder()
88+
.text("test answer")
89+
.toolCalls(List.of(new AssistantMessage.ToolCall("id", "type", "name", "arguments")))
90+
.build();
8891

8992
memory.saveAll(sessionId, List.of(userMessage, assistantMessage));
9093
messages = memory.findByConversationId(sessionId);
@@ -112,10 +115,11 @@ void addAndGet() {
112115
assertThat(((UserMessage) messages.get(0)).getMedia()).usingRecursiveFieldByFieldElementComparator()
113116
.isEqualTo(media);
114117
memory.deleteByConversationId(sessionId);
115-
ToolResponseMessage toolResponseMessage = new ToolResponseMessage(
116-
List.of(new ToolResponse("id", "name", "responseData"),
117-
new ToolResponse("id2", "name2", "responseData2")),
118-
Map.of("id", "id", "metadataKey", "metadata"));
118+
ToolResponseMessage toolResponseMessage = ToolResponseMessage.builder()
119+
.responses(List.of(new ToolResponse("id", "name", "responseData"),
120+
new ToolResponse("id2", "name2", "responseData2")))
121+
.metadata(Map.of("id", "id", "metadataKey", "metadata"))
122+
.build();
119123
memory.saveAll(sessionId, List.of(toolResponseMessage));
120124
messages = memory.findByConversationId(sessionId);
121125
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‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -209,14 +209,14 @@ private Message getMessage(UdtValue udt) {
209209
Map<String, Object> props = Map.of(CONVERSATION_TS, udt.getInstant(this.conf.messageUdtTimestampColumn));
210210
switch (MessageType.valueOf(udt.getString(this.conf.messageUdtTypeColumn))) {
211211
case ASSISTANT:
212-
return new AssistantMessage(content, props);
212+
return AssistantMessage.builder().text(content).metadata(props).build();
213213
case USER:
214214
return UserMessage.builder().text(content).metadata(props).build();
215215
case SYSTEM:
216216
return SystemMessage.builder().text(content).metadata(props).build();
217217
case TOOL:
218218
// todo – persist ToolResponse somehow
219-
return new ToolResponseMessage(List.of(), props);
219+
return ToolResponseMessage.builder().responses(List.of()).metadata(props).build();
220220
default:
221221
throw new IllegalStateException(
222222
String.format("unknown message type %s", udt.getString(this.conf.messageUdtTypeColumn)));

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

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -172,25 +172,28 @@ public Neo4jChatMemoryRepositoryConfig getConfig() {
172172

173173
private Message buildToolMessage(org.neo4j.driver.Record record) {
174174
Message message;
175-
message = new ToolResponseMessage(record.get("toolResponses").asList(v -> {
175+
message = ToolResponseMessage.builder().responses(record.get("toolResponses").asList(v -> {
176176
Map<String, Object> trMap = v.asMap();
177177
return new ToolResponseMessage.ToolResponse((String) trMap.get(ToolResponseAttributes.ID.getValue()),
178178
(String) trMap.get(ToolResponseAttributes.NAME.getValue()),
179179
(String) trMap.get(ToolResponseAttributes.RESPONSE_DATA.getValue()));
180-
}), record.get("metadata").asMap());
180+
})).metadata(record.get("metadata").asMap()).build();
181181
return message;
182182
}
183183

184184
private Message buildAssistantMessage(org.neo4j.driver.Record record, Map<String, Object> messageMap,
185185
List<Media> mediaList) {
186186
Message message;
187-
message = new AssistantMessage(messageMap.get(MessageAttributes.TEXT_CONTENT.getValue()).toString(),
188-
record.get("metadata").asMap(Map.of()), record.get("toolCalls").asList(v -> {
189-
var toolCallMap = v.asMap();
190-
return new AssistantMessage.ToolCall((String) toolCallMap.get("id"),
191-
(String) toolCallMap.get("type"), (String) toolCallMap.get("name"),
192-
(String) toolCallMap.get("arguments"));
193-
}), mediaList);
187+
message = AssistantMessage.builder()
188+
.text(messageMap.get(MessageAttributes.TEXT_CONTENT.getValue()).toString())
189+
.metadata(record.get("metadata").asMap(Map.of()))
190+
.toolCalls(record.get("toolCalls").asList(v -> {
191+
var toolCallMap = v.asMap();
192+
return new AssistantMessage.ToolCall((String) toolCallMap.get("id"), (String) toolCallMap.get("type"),
193+
(String) toolCallMap.get("name"), (String) toolCallMap.get("arguments"));
194+
}))
195+
.media(mediaList)
196+
.build();
194197
return message;
195198
}
196199

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

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
/**
5454
* Integration tests for {@link Neo4jChatMemoryRepository}.
5555
*
56+
* @author Jemin Huh
5657
* @author Enrico Rampazzo
5758
* @since 1.0.0
5859
*/
@@ -263,9 +264,11 @@ void handleMediaContent() {
263264
void handleAssistantMessageWithToolCalls() {
264265
var conversationId = UUID.randomUUID().toString();
265266

266-
AssistantMessage assistantMessage = new AssistantMessage("Message with tool calls", Map.of(),
267-
List.of(new AssistantMessage.ToolCall("id1", "type1", "name1", "arguments1"),
268-
new AssistantMessage.ToolCall("id2", "type2", "name2", "arguments2")));
267+
AssistantMessage assistantMessage = AssistantMessage.builder()
268+
.text("Message with tool calls")
269+
.toolCalls(List.of(new AssistantMessage.ToolCall("id1", "type1", "name1", "arguments1"),
270+
new AssistantMessage.ToolCall("id2", "type2", "name2", "arguments2")))
271+
.build();
269272

270273
this.chatMemoryRepository.saveAll(conversationId, List.<Message>of(assistantMessage));
271274

@@ -282,9 +285,11 @@ void handleAssistantMessageWithToolCalls() {
282285
void handleToolResponseMessage() {
283286
var conversationId = UUID.randomUUID().toString();
284287

285-
ToolResponseMessage toolResponseMessage = new ToolResponseMessage(List
286-
.of(new ToolResponse("id1", "name1", "responseData1"), new ToolResponse("id2", "name2", "responseData2")),
287-
Map.of("metadataKey", "metadataValue"));
288+
ToolResponseMessage toolResponseMessage = ToolResponseMessage.builder()
289+
.responses(List.of(new ToolResponse("id1", "name1", "responseData1"),
290+
new ToolResponse("id2", "name2", "responseData2")))
291+
.metadata(Map.of("metadataKey", "metadataValue"))
292+
.build();
288293

289294
this.chatMemoryRepository.saveAll(conversationId, List.<Message>of(toolResponseMessage));
290295

‎models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/AnthropicChatModel.java‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@
8282
/**
8383
* The {@link ChatModel} implementation for the Anthropic service.
8484
*
85+
* @author Jemin Huh
8586
* @author Christian Tzolov
8687
* @author luocongqiu
8788
* @author Mariusz Bernacki
@@ -302,19 +303,21 @@ private ChatResponse toChatResponse(ChatCompletionResponse chatCompletion, Usage
302303
for (ContentBlock content : chatCompletion.content()) {
303304
switch (content.type()) {
304305
case TEXT, TEXT_DELTA:
305-
generations.add(new Generation(new AssistantMessage(content.text(), Map.of()),
306+
generations.add(new Generation(new AssistantMessage(content.text()),
306307
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build()));
307308
break;
308309
case THINKING, THINKING_DELTA:
309310
Map<String, Object> thinkingProperties = new HashMap<>();
310311
thinkingProperties.put("signature", content.signature());
311-
generations.add(new Generation(new AssistantMessage(content.thinking(), thinkingProperties),
312+
generations.add(new Generation(
313+
AssistantMessage.builder().text(content.thinking()).metadata(thinkingProperties).build(),
312314
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build()));
313315
break;
314316
case REDACTED_THINKING:
315317
Map<String, Object> redactedProperties = new HashMap<>();
316318
redactedProperties.put("data", content.data());
317-
generations.add(new Generation(new AssistantMessage(null, redactedProperties),
319+
generations.add(new Generation(
320+
AssistantMessage.builder().text((String) null).metadata(redactedProperties).build(),
318321
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build()));
319322
break;
320323
case TOOL_USE:
@@ -328,13 +331,13 @@ private ChatResponse toChatResponse(ChatCompletionResponse chatCompletion, Usage
328331
}
329332

330333
if (chatCompletion.stopReason() != null && generations.isEmpty()) {
331-
Generation generation = new Generation(new AssistantMessage(null, Map.of()),
334+
Generation generation = new Generation(new AssistantMessage((String) null),
332335
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build());
333336
generations.add(generation);
334337
}
335338

336339
if (!CollectionUtils.isEmpty(toolCalls)) {
337-
AssistantMessage assistantMessage = new AssistantMessage("", Map.of(), toolCalls);
340+
AssistantMessage assistantMessage = AssistantMessage.builder().text("").toolCalls(toolCalls).build();
338341
Generation toolCallGeneration = new Generation(assistantMessage,
339342
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build());
340343
generations.add(toolCallGeneration);

‎models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiChatModel.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,7 @@
106106
* {@link ChatModel} implementation for {@literal Microsoft Azure AI} backed by
107107
* {@link OpenAIClient}.
108108
*
109+
* @author Jemin Huh
109110
* @author Mark Pollack
110111
* @author Ueibin Kim
111112
* @author John Blum
@@ -485,7 +486,7 @@ private Generation buildGeneration(ChatChoice choice, Map<String, Object> metada
485486
}
486487

487488
var content = responseMessage == null ? "" : responseMessage.getContent();
488-
var assistantMessage = new AssistantMessage(content, metadata, toolCalls);
489+
var assistantMessage = AssistantMessage.builder().text(content).metadata(metadata).toolCalls(toolCalls).build();
489490
var generationMetadata = generateChoiceMetadata(choice);
490491

491492
return new Generation(assistantMessage, generationMetadata);

‎models/spring-ai-bedrock-converse/src/main/java/org/springframework/ai/bedrock/converse/BedrockProxyChatModel.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,6 @@
2525
import java.util.ArrayList;
2626
import java.util.Base64;
2727
import java.util.List;
28-
import java.util.Map;
2928
import java.util.Set;
3029

3130
import io.micrometer.observation.Observation;
@@ -128,6 +127,7 @@
128127
* <p>
129128
* https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters.html
130129
*
130+
* @author Jemin Huh
131131
* @author Christian Tzolov
132132
* @author Wei Jiang
133133
* @author Alexandros Pappas
@@ -566,14 +566,14 @@ private ChatResponse toChatResponse(ConverseResponse response, ChatResponse perv
566566
List<Generation> generations = message.content()
567567
.stream()
568568
.filter(content -> content.type() != ContentBlock.Type.TOOL_USE)
569-
.map(content -> new Generation(new AssistantMessage(content.text(), Map.of()),
569+
.map(content -> new Generation(new AssistantMessage(content.text()),
570570
ChatGenerationMetadata.builder().finishReason(response.stopReasonAsString()).build()))
571571
.toList();
572572

573573
List<Generation> allGenerations = new ArrayList<>(generations);
574574

575575
if (response.stopReasonAsString() != null && generations.isEmpty()) {
576-
Generation generation = new Generation(new AssistantMessage(null, Map.of()),
576+
Generation generation = new Generation(new AssistantMessage((String) null),
577577
ChatGenerationMetadata.builder().finishReason(response.stopReasonAsString()).build());
578578
allGenerations.add(generation);
579579
}
@@ -597,7 +597,7 @@ private ChatResponse toChatResponse(ConverseResponse response, ChatResponse perv
597597
.add(new AssistantMessage.ToolCall(functionCallId, "function", functionName, functionArguments));
598598
}
599599

600-
AssistantMessage assistantMessage = new AssistantMessage("", Map.of(), toolCalls);
600+
AssistantMessage assistantMessage = AssistantMessage.builder().text("").toolCalls(toolCalls).build();
601601
Generation toolCallGeneration = new Generation(assistantMessage,
602602
ChatGenerationMetadata.builder().finishReason(response.stopReasonAsString()).build());
603603
allGenerations.add(toolCallGeneration);

‎models/spring-ai-bedrock-converse/src/main/java/org/springframework/ai/bedrock/converse/api/ConverseApiUtils.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
/**
6161
* Amazon Bedrock Converse API utils.
6262
*
63+
* @author Jemin Huh
6364
* @author Wei Jiang
6465
* @author Christian Tzolov
6566
* @author Alexandros Pappas
@@ -140,7 +141,7 @@ public static Flux<ChatResponse> toChatResponse(Flux<ConverseStreamOutput> respo
140141
}
141142
}
142143

143-
AssistantMessage assistantMessage = new AssistantMessage("", Map.of(), toolCalls);
144+
AssistantMessage assistantMessage = AssistantMessage.builder().text("").toolCalls(toolCalls).build();
144145
Generation toolCallGeneration = new Generation(assistantMessage,
145146
ChatGenerationMetadata.builder().finishReason("tool_use").build());
146147

@@ -175,8 +176,7 @@ else if (nextEvent instanceof ContentBlockStartEvent contentBlockStartEvent) {
175176
else if (nextEvent instanceof ContentBlockDeltaEvent contentBlockDeltaEvent) {
176177
if (contentBlockDeltaEvent.delta().type().equals(ContentBlockDelta.Type.TEXT)) {
177178

178-
var generation = new Generation(
179-
new AssistantMessage(contentBlockDeltaEvent.delta().text(), Map.of()),
179+
var generation = new Generation(new AssistantMessage(contentBlockDeltaEvent.delta().text()),
180180
ChatGenerationMetadata.builder()
181181
.finishReason(lastAggregation.metadataAggregation().stopReason())
182182
.build());

‎models/spring-ai-deepseek/src/main/java/org/springframework/ai/deepseek/DeepSeekAssistantMessage.java‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,14 @@
2323
import org.springframework.ai.chat.messages.AssistantMessage;
2424
import org.springframework.ai.content.Media;
2525

26+
/**
27+
* Represents an assistant message generated by the DeepSeek model.
28+
*
29+
* @author Jemin Huh
30+
* @author Soby Chacko
31+
* @author Mark Pollack
32+
* @since 1.0.0
33+
*/
2634
public class DeepSeekAssistantMessage extends AssistantMessage {
2735

2836
private Boolean prefix;
@@ -38,22 +46,14 @@ public DeepSeekAssistantMessage(String content, String reasoningContent) {
3846
this.reasoningContent = reasoningContent;
3947
}
4048

41-
public DeepSeekAssistantMessage(String content, Map<String, Object> properties) {
42-
super(content, properties);
43-
}
44-
45-
public DeepSeekAssistantMessage(String content, Map<String, Object> properties, List<ToolCall> toolCalls) {
46-
super(content, properties, toolCalls);
47-
}
48-
49-
public DeepSeekAssistantMessage(String content, String reasoningContent, Map<String, Object> properties,
49+
public DeepSeekAssistantMessage(String content, String reasoningContent, Map<String, Object> metadata,
5050
List<ToolCall> toolCalls) {
51-
this(content, reasoningContent, properties, toolCalls, List.of());
51+
this(content, reasoningContent, metadata, toolCalls, List.of());
5252
}
5353

54-
public DeepSeekAssistantMessage(String content, String reasoningContent, Map<String, Object> properties,
54+
public DeepSeekAssistantMessage(String content, String reasoningContent, Map<String, Object> metadata,
5555
List<ToolCall> toolCalls, List<Media> media) {
56-
super(content, properties, toolCalls, media);
56+
super(content, metadata, toolCalls, media);
5757
this.reasoningContent = reasoningContent;
5858
}
5959

@@ -102,9 +102,9 @@ public int hashCode() {
102102

103103
@Override
104104
public String toString() {
105-
return "AssistantMessage [messageType=" + this.messageType + ", toolCalls=" + super.getToolCalls()
106-
+ ", textContent=" + this.textContent + ", reasoningContent=" + this.reasoningContent + ", prefix="
107-
+ this.prefix + ", metadata=" + this.metadata + "]";
105+
return "AssistantMessage [messageType=" + this.messageType + ", toolCalls=" + this.toolCalls + ", textContent="
106+
+ this.textContent + ", reasoningContent=" + this.reasoningContent + ", prefix=" + this.prefix
107+
+ ", metadata=" + this.metadata + "]";
108108
}
109109

110110
}

‎models/spring-ai-huggingface/src/main/java/org/springframework/ai/huggingface/HuggingfaceChatModel.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
* An implementation of {@link ChatModel} that interfaces with HuggingFace Inference
4040
* Endpoints for text generation.
4141
*
42+
* @author Jemin Huh
4243
* @author Mark Pollack
4344
* @author Jihoon Kim
4445
*/
@@ -104,7 +105,8 @@ public ChatResponse call(Prompt prompt) {
104105
new TypeReference<Map<String, Object>>() {
105106

106107
});
107-
Generation generation = new Generation(new AssistantMessage(generatedText, detailsMap));
108+
Generation generation = new Generation(
109+
AssistantMessage.builder().text(generatedText).metadata(detailsMap).build());
108110
generations.add(generation);
109111
}
110112
return new ChatResponse(generations);

0 commit comments

Comments
 (0)