Skip to content

Commit f47b272

Browse files
sunyuhan1998sunyuhan1998
committed
feat: GH-4454 Added Builder pattern support for ZhiPuAiAssistantMessage and included corresponding unit tests.
Signed-off-by: Sun Yuhan <[email protected]>
1 parent 84efb6a commit f47b272

File tree

4 files changed

+194
-8
lines changed

4 files changed

+194
-8
lines changed

‎models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiAssistantMessage.java‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@
2525

2626
/**
2727
* @author YunKui Lu
28+
* @author Sun Yuhan
2829
*/
2930
public class ZhiPuAiAssistantMessage extends AssistantMessage {
3031

@@ -78,4 +79,48 @@ public String toString() {
7879
+ this.textContent + '\'' + '}';
7980
}
8081

82+
public static Builder builder() {
83+
return new Builder();
84+
}
85+
86+
public static final class Builder extends AssistantMessage.Builder {
87+
88+
private String reasoningContent;
89+
90+
@Override
91+
public Builder content(String content) {
92+
this.content = content;
93+
return this;
94+
}
95+
96+
@Override
97+
public Builder properties(Map<String, Object> properties) {
98+
this.properties = properties;
99+
return this;
100+
}
101+
102+
@Override
103+
public Builder toolCalls(List<ToolCall> toolCalls) {
104+
this.toolCalls = toolCalls;
105+
return this;
106+
}
107+
108+
@Override
109+
public Builder media(List<Media> media) {
110+
this.media = media;
111+
return this;
112+
}
113+
114+
public Builder reasoningContent(String reasoningContent) {
115+
this.reasoningContent = reasoningContent;
116+
return this;
117+
}
118+
119+
public ZhiPuAiAssistantMessage build() {
120+
return new ZhiPuAiAssistantMessage(this.content, this.reasoningContent, this.properties, this.toolCalls,
121+
this.media);
122+
}
123+
124+
}
125+
81126
}

‎models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiChatModel.java‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -230,8 +230,14 @@ private static Generation buildGeneration(Choice choice, Map<String, Object> met
230230
String textContent = choice.message().content();
231231
String reasoningContent = choice.message().reasoningContent();
232232

233-
var assistantMessage = new ZhiPuAiAssistantMessage(textContent, reasoningContent, metadata, toolCalls,
234-
List.of());
233+
var assistantMessage = ZhiPuAiAssistantMessage.builder()
234+
.content(textContent)
235+
.reasoningContent(reasoningContent)
236+
.properties(metadata)
237+
.toolCalls(toolCalls)
238+
.media(List.of())
239+
.build();
240+
235241
String finishReason = (choice.finishReason() != null ? choice.finishReason().name() : "");
236242
var generationMetadata = ChatGenerationMetadata.builder().finishReason(finishReason).build();
237243
return new Generation(assistantMessage, generationMetadata);

‎models/spring-ai-zhipuai/src/test/java/org/springframework/ai/zhipuai/ZhiPuAiAssistantMessageTests.java‎

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
/*
2+
* Copyright 2025-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package org.springframework.ai.zhipuai;
18+
19+
import java.util.HashMap;
20+
import java.util.List;
21+
import java.util.Map;
22+
23+
import org.junit.jupiter.api.Test;
24+
25+
import org.springframework.ai.chat.messages.AssistantMessage.ToolCall;
26+
import org.springframework.ai.content.Media;
27+
28+
import static org.assertj.core.api.Assertions.assertThat;
29+
import static org.assertj.core.api.Assertions.assertThatNoException;
30+
31+
/**
32+
* Tests for {@link ZhiPuAiAssistantMessage}.
33+
*
34+
* @author Sun Yuhan
35+
*/
36+
class ZhiPuAiAssistantMessageTests {
37+
38+
@Test
39+
public void testConstructorWithContentOnly() {
40+
String content = "Hello, world!";
41+
ZhiPuAiAssistantMessage message = new ZhiPuAiAssistantMessage(content);
42+
43+
assertThat(message.getText()).isEqualTo(content);
44+
assertThat(message.getReasoningContent()).isNull();
45+
}
46+
47+
@Test
48+
public void testConstructorWithAllParameters() {
49+
String content = "Hello, world!";
50+
String reasoningContent = "This is my reasoning";
51+
Map<String, Object> properties = new HashMap<>();
52+
properties.put("key1", "value1");
53+
List<ToolCall> toolCalls = List.of(new ToolCall("1", "function", "myFunction", "{}"));
54+
List<Media> media = List.of();
55+
56+
ZhiPuAiAssistantMessage message = new ZhiPuAiAssistantMessage(content, reasoningContent, properties, toolCalls,
57+
media);
58+
59+
assertThat(message.getText()).isEqualTo(content);
60+
assertThat(message.getReasoningContent()).isEqualTo(reasoningContent);
61+
assertThat(message.getMetadata()).containsAllEntriesOf(properties);
62+
assertThat(message.getToolCalls()).isEqualTo(toolCalls);
63+
}
64+
65+
@Test
66+
public void testSettersAndGetters() {
67+
ZhiPuAiAssistantMessage message = new ZhiPuAiAssistantMessage("test");
68+
69+
String reasoningContent = "New reasoning content";
70+
71+
message.setReasoningContent(reasoningContent);
72+
73+
assertThat(message.getReasoningContent()).isEqualTo(reasoningContent);
74+
}
75+
76+
@Test
77+
public void testEqualsAndHashCode() {
78+
String content = "Hello, world!";
79+
String reasoningContent = "This is my reasoning";
80+
Map<String, Object> properties = new HashMap<>();
81+
properties.put("key1", "value1");
82+
List<ToolCall> toolCalls = List.of(new ToolCall("1", "function", "myFunction", "{}"));
83+
List<Media> media = List.of();
84+
85+
ZhiPuAiAssistantMessage message1 = new ZhiPuAiAssistantMessage(content, reasoningContent, properties, toolCalls,
86+
media);
87+
ZhiPuAiAssistantMessage message2 = new ZhiPuAiAssistantMessage(content, reasoningContent, properties, toolCalls,
88+
media);
89+
90+
assertThat(message1).isEqualTo(message2);
91+
assertThat(message1.hashCode()).isEqualTo(message2.hashCode());
92+
93+
ZhiPuAiAssistantMessage message3 = new ZhiPuAiAssistantMessage(content, "different reasoning", properties,
94+
toolCalls, media);
95+
assertThat(message1).isNotEqualTo(message3);
96+
}
97+
98+
@Test
99+
public void testToString() {
100+
String content = "Hello, world!";
101+
String reasoningContent = "This is my reasoning";
102+
Map<String, Object> properties = new HashMap<>();
103+
properties.put("key1", "value1");
104+
List<ToolCall> toolCalls = List.of(new ToolCall("1", "function", "myFunction", "{}"));
105+
List<Media> media = List.of();
106+
107+
ZhiPuAiAssistantMessage message = new ZhiPuAiAssistantMessage(content, reasoningContent, properties, toolCalls,
108+
media);
109+
110+
assertThatNoException().isThrownBy(message::toString);
111+
assertThat(message.toString()).contains(content, reasoningContent);
112+
}
113+
114+
@Test
115+
public void testBuilderComplete() {
116+
Map<String, Object> properties = Map.of("key", "value");
117+
List<ToolCall> toolCalls = List.of(new ToolCall("1", "function", "testFunction", "{}"));
118+
List<Media> media = List.of();
119+
120+
ZhiPuAiAssistantMessage message = ZhiPuAiAssistantMessage.builder()
121+
.content("content")
122+
.reasoningContent("reasoning")
123+
.properties(properties)
124+
.toolCalls(toolCalls)
125+
.media(media)
126+
.build();
127+
128+
assertThat(message.getText()).isEqualTo("content");
129+
assertThat(message.getReasoningContent()).isEqualTo("reasoning");
130+
assertThat(message.getMetadata()).containsAllEntriesOf(properties);
131+
assertThat(message.getToolCalls()).isEqualTo(toolCalls);
132+
assertThat(message.getMedia()).isEqualTo(media);
133+
}
134+
135+
}

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -121,17 +121,17 @@ public record ToolCall(String id, String type, String name, String arguments) {
121121

122122
}
123123

124-
public static final class Builder {
124+
public static class Builder {
125125

126-
private String content;
126+
protected String content;
127127

128-
private Map<String, Object> properties = Map.of();
128+
protected Map<String, Object> properties = Map.of();
129129

130-
private List<ToolCall> toolCalls = List.of();
130+
protected List<ToolCall> toolCalls = List.of();
131131

132-
private List<Media> media = List.of();
132+
protected List<Media> media = List.of();
133133

134-
private Builder() {
134+
protected Builder() {
135135
}
136136

137137
public Builder content(String content) {

0 commit comments

Comments
 (0)