Skip to content

Commit 39c39fd

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

File tree

15 files changed

+90
-50
lines changed

15 files changed

+90
-50
lines changed

‎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

‎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: 13 additions & 13 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,7 +102,7 @@ public int hashCode() {
102102

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

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

‎models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxChatModel.java‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,7 @@
7676
* {@link ChatModel} and {@link StreamingChatModel} implementation for {@literal MiniMax}
7777
* backed by {@link MiniMaxApi}.
7878
*
79+
* @author Jemin Huh
7980
* @author Geng Rong
8081
* @author Alexandros Pappas
8182
* @author Ilayaperumal Gopinathan
@@ -225,7 +226,11 @@ private static Generation buildGeneration(Choice choice, Map<String, Object> met
225226
acc1.addAll(acc2);
226227
return acc1;
227228
});
228-
var assistantMessage = new AssistantMessage(choice.message().content(), metadata, toolCalls);
229+
var assistantMessage = AssistantMessage.builder()
230+
.text(choice.message().content())
231+
.metadata(metadata)
232+
.toolCalls(toolCalls)
233+
.build();
229234
String finishReason = (choice.finishReason() != null ? choice.finishReason().name() : "");
230235
var generationMetadata = ChatGenerationMetadata.builder().finishReason(finishReason).build();
231236
return new Generation(assistantMessage, generationMetadata);
@@ -424,7 +429,11 @@ private Generation buildGeneration(ChatCompletionMessage message, ChatCompletion
424429
toolCall.function().name(), toolCall.function().arguments()))
425430
.toList();
426431

427-
var assistantMessage = new AssistantMessage(message.content(), metadata, toolCalls);
432+
var assistantMessage = AssistantMessage.builder()
433+
.text(message.content())
434+
.metadata(metadata)
435+
.toolCalls(toolCalls)
436+
.build();
428437
String finishReason = (completionFinishReason != null ? completionFinishReason.name() : "");
429438
var generationMetadata = ChatGenerationMetadata.builder().finishReason(finishReason).build();
430439
return new Generation(assistantMessage, generationMetadata);

‎models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatModel.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -354,7 +354,11 @@ private Generation buildGeneration(Choice choice, Map<String, Object> metadata)
354354
toolCall.function().name(), toolCall.function().arguments()))
355355
.toList();
356356

357-
var assistantMessage = new AssistantMessage(choice.message().content(), metadata, toolCalls);
357+
var assistantMessage = AssistantMessage.builder()
358+
.text(choice.message().content())
359+
.metadata(metadata)
360+
.toolCalls(toolCalls)
361+
.build();
358362
String finishReason = (choice.finishReason() != null ? choice.finishReason().name() : "");
359363
var generationMetadata = ChatGenerationMetadata.builder().finishReason(finishReason).build();
360364
return new Generation(assistantMessage, generationMetadata);

0 commit comments

Comments
 (0)