Skip to content

Commit 6a63ae8

Browse files
leeyazhouleeyazhou
committed
perf: use AssistantMessage#builder()
Signed-off-by: leeyazhou <[email protected]>
1 parent f4a1c96 commit 6a63ae8

File tree

26 files changed

+233
-83
lines changed

26 files changed

+233
-83
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: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -83,8 +83,11 @@ void addAndGet() {
8383
memory.deleteByConversationId(sessionId);
8484
assertThat(memory.findByConversationId(sessionId)).isEmpty();
8585

86-
AssistantMessage assistantMessage = new AssistantMessage("test answer", Map.of(),
87-
List.of(new AssistantMessage.ToolCall("id", "type", "name", "arguments")));
86+
AssistantMessage assistantMessage = AssistantMessage.builder()
87+
.content("test answer")
88+
.properties(Map.of())
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);

‎auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-chroma/src/test/java/org/springframework/ai/vectorstore/chroma/autoconfigure/ChromaVectorStoreAutoConfigurationIT.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -95,8 +95,10 @@ public void verifyThatChromaCanHandleComplexMetadataValues() {
9595

9696
var response = ChatClientResponse.builder()
9797
.chatResponse(ChatResponse.builder()
98-
.generations(List
99-
.of(new Generation(new AssistantMessage("AssistantMessage", Map.of("annotations", List.of())))))
98+
.generations(List.of(new Generation(AssistantMessage.builder()
99+
.content("AssistantMessage")
100+
.properties(Map.of("annotations", List.of()))
101+
.build())))
100102
.build())
101103
.build();
102104
var res2 = advisor.after(response, null);

‎mcp/common/src/main/java/org/springframework/ai/mcp/McpToolUtils.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -433,7 +433,7 @@ public static List<ToolCallback> getToolCallbacksFromAsyncClients(List<McpAsyncC
433433
if (CollectionUtils.isEmpty(asyncMcpClients)) {
434434
return List.of();
435435
}
436-
return List.of((new AsyncMcpToolCallbackProvider(asyncMcpClients).getToolCallbacks()));
436+
return List.of((AsyncMcpToolCallbackProvider.builder().mcpClients(asyncMcpClients).build().getToolCallbacks()));
437437
}
438438

439439
@JsonIgnoreProperties(ignoreUnknown = true)

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -209,7 +209,7 @@ 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().content(content).properties(props).build();
213213
case USER:
214214
return UserMessage.builder().text(content).metadata(props).build();
215215
case SYSTEM:

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

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -183,14 +183,16 @@ private Message buildToolMessage(org.neo4j.driver.Record record) {
183183

184184
private Message buildAssistantMessage(org.neo4j.driver.Record record, Map<String, Object> messageMap,
185185
List<Media> mediaList) {
186-
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);
186+
Message message = AssistantMessage.builder()
187+
.content(messageMap.get(MessageAttributes.TEXT_CONTENT.getValue()).toString())
188+
.properties(record.get("metadata").asMap(Map.of()))
189+
.toolCalls(record.get("toolCalls").asList(v -> {
190+
var toolCallMap = v.asMap();
191+
return new AssistantMessage.ToolCall((String) toolCallMap.get("id"), (String) toolCallMap.get("type"),
192+
(String) toolCallMap.get("name"), (String) toolCallMap.get("arguments"));
193+
}))
194+
.media(mediaList)
195+
.build();
194196
return message;
195197
}
196198

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

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -263,9 +263,12 @@ void handleMediaContent() {
263263
void handleAssistantMessageWithToolCalls() {
264264
var conversationId = UUID.randomUUID().toString();
265265

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")));
266+
AssistantMessage assistantMessage = AssistantMessage.builder()
267+
.content("Message with tool calls")
268+
.properties(Map.of())
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

@@ -389,7 +392,8 @@ void saveAndFindMessagesWithEmptyContentOrMetadata() {
389392
assertThat(retrievedEmptyContentMsg.getMetadata().keySet()).hasSize(1); // Only
390393
// messageType
391394

392-
// Verify second message (empty metadata from input, should only have messageType
395+
// Verify second message (empty metadata from input, should only have
396+
// messageType
393397
// after retrieval)
394398
Message retrievedEmptyMetadataMsg = retrievedMessages.get(1);
395399
assertThat(retrievedEmptyMetadataMsg).isInstanceOf(UserMessage.class);

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

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -321,19 +321,24 @@ private ChatResponse toChatResponse(ChatCompletionResponse chatCompletion, Usage
321321
for (ContentBlock content : chatCompletion.content()) {
322322
switch (content.type()) {
323323
case TEXT, TEXT_DELTA:
324-
generations.add(new Generation(new AssistantMessage(content.text(), Map.of()),
324+
generations.add(new Generation(
325+
AssistantMessage.builder().content(content.text()).properties(Map.of()).build(),
325326
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build()));
326327
break;
327328
case THINKING, THINKING_DELTA:
328329
Map<String, Object> thinkingProperties = new HashMap<>();
329330
thinkingProperties.put("signature", content.signature());
330-
generations.add(new Generation(new AssistantMessage(content.thinking(), thinkingProperties),
331+
generations.add(new Generation(
332+
AssistantMessage.builder()
333+
.content(content.thinking())
334+
.properties(thinkingProperties)
335+
.build(),
331336
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build()));
332337
break;
333338
case REDACTED_THINKING:
334339
Map<String, Object> redactedProperties = new HashMap<>();
335340
redactedProperties.put("data", content.data());
336-
generations.add(new Generation(new AssistantMessage(null, redactedProperties),
341+
generations.add(new Generation(AssistantMessage.builder().properties(redactedProperties).build(),
337342
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build()));
338343
break;
339344
case TOOL_USE:
@@ -347,13 +352,17 @@ private ChatResponse toChatResponse(ChatCompletionResponse chatCompletion, Usage
347352
}
348353

349354
if (chatCompletion.stopReason() != null && generations.isEmpty()) {
350-
Generation generation = new Generation(new AssistantMessage(null, Map.of()),
355+
Generation generation = new Generation(AssistantMessage.builder().content("").properties(Map.of()).build(),
351356
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build());
352357
generations.add(generation);
353358
}
354359

355360
if (!CollectionUtils.isEmpty(toolCalls)) {
356-
AssistantMessage assistantMessage = new AssistantMessage("", Map.of(), toolCalls);
361+
AssistantMessage assistantMessage = AssistantMessage.builder()
362+
.content("")
363+
.properties(Map.of())
364+
.toolCalls(toolCalls)
365+
.build();
357366
Generation toolCallGeneration = new Generation(assistantMessage,
358367
ChatGenerationMetadata.builder().finishReason(chatCompletion.stopReason()).build());
359368
generations.add(toolCallGeneration);

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -493,7 +493,11 @@ private Generation buildGeneration(ChatChoice choice, Map<String, Object> metada
493493
}
494494

495495
var content = responseMessage == null ? "" : responseMessage.getContent();
496-
var assistantMessage = new AssistantMessage(content, metadata, toolCalls);
496+
var assistantMessage = AssistantMessage.builder()
497+
.content(content)
498+
.properties(metadata)
499+
.toolCalls(toolCalls)
500+
.build();
497501
var generationMetadata = generateChoiceMetadata(choice);
498502

499503
return new Generation(assistantMessage, generationMetadata);

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

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -572,14 +572,15 @@ private ChatResponse toChatResponse(ConverseResponse response, ChatResponse perv
572572
List<Generation> generations = message.content()
573573
.stream()
574574
.filter(content -> content.type() != ContentBlock.Type.TOOL_USE)
575-
.map(content -> new Generation(new AssistantMessage(content.text(), Map.of()),
575+
.map(content -> new Generation(
576+
AssistantMessage.builder().content(content.text()).properties(Map.of()).build(),
576577
ChatGenerationMetadata.builder().finishReason(response.stopReasonAsString()).build()))
577578
.toList();
578579

579580
List<Generation> allGenerations = new ArrayList<>(generations);
580581

581582
if (response.stopReasonAsString() != null && generations.isEmpty()) {
582-
Generation generation = new Generation(new AssistantMessage(null, Map.of()),
583+
Generation generation = new Generation(AssistantMessage.builder().properties(Map.of()).build(),
583584
ChatGenerationMetadata.builder().finishReason(response.stopReasonAsString()).build());
584585
allGenerations.add(generation);
585586
}
@@ -603,7 +604,11 @@ private ChatResponse toChatResponse(ConverseResponse response, ChatResponse perv
603604
.add(new AssistantMessage.ToolCall(functionCallId, "function", functionName, functionArguments));
604605
}
605606

606-
AssistantMessage assistantMessage = new AssistantMessage("", Map.of(), toolCalls);
607+
AssistantMessage assistantMessage = AssistantMessage.builder()
608+
.content("")
609+
.properties(Map.of())
610+
.toolCalls(toolCalls)
611+
.build();
607612
Generation toolCallGeneration = new Generation(assistantMessage,
608613
ChatGenerationMetadata.builder().finishReason(response.stopReasonAsString()).build());
609614
allGenerations.add(toolCallGeneration);

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

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,11 @@ public static Flux<ChatResponse> toChatResponse(Flux<ConverseStreamOutput> respo
140140
}
141141
}
142142

143-
AssistantMessage assistantMessage = new AssistantMessage("", Map.of(), toolCalls);
143+
AssistantMessage assistantMessage = AssistantMessage.builder()
144+
.content("")
145+
.properties(Map.of())
146+
.toolCalls(toolCalls)
147+
.build();
144148
Generation toolCallGeneration = new Generation(assistantMessage,
145149
ChatGenerationMetadata.builder().finishReason("tool_use").build());
146150

@@ -176,7 +180,10 @@ else if (nextEvent instanceof ContentBlockDeltaEvent contentBlockDeltaEvent) {
176180
if (contentBlockDeltaEvent.delta().type().equals(ContentBlockDelta.Type.TEXT)) {
177181

178182
var generation = new Generation(
179-
new AssistantMessage(contentBlockDeltaEvent.delta().text(), Map.of()),
183+
AssistantMessage.builder()
184+
.content(contentBlockDeltaEvent.delta().text())
185+
.properties(Map.of())
186+
.build(),
180187
ChatGenerationMetadata.builder()
181188
.finishReason(lastAggregation.metadataAggregation().stopReason())
182189
.build());

0 commit comments

Comments
 (0)