Skip to content

Commit 01082b8

Browse files
sunyuhan1998ericbottard
authored andcommitted
feat: GH-4454 Added Builder pattern support for ZhiPuAiAssistantMessage and included corresponding unit tests.
Fix #4454 Signed-off-by: Sun Yuhan <[email protected]> Signed-off-by: Eric Bottard <[email protected]>
1 parent 7d0aff0 commit 01082b8

File tree

3 files changed

+196
-2
lines changed

3 files changed

+196
-2
lines changed

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

Lines changed: 53 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

@@ -33,10 +34,18 @@ public class ZhiPuAiAssistantMessage extends AssistantMessage {
3334
*/
3435
private String reasoningContent;
3536

37+
/**
38+
* @deprecated in favor of using {@link ZhiPuAiAssistantMessage.Builder}
39+
*/
40+
@Deprecated
3641
public ZhiPuAiAssistantMessage(String content) {
3742
super(content);
3843
}
3944

45+
/**
46+
* @deprecated in favor of using {@link ZhiPuAiAssistantMessage.Builder}
47+
*/
48+
@Deprecated
4049
public ZhiPuAiAssistantMessage(String content, String reasoningContent, Map<String, Object> properties,
4150
List<ToolCall> toolCalls, List<Media> media) {
4251
super(content, properties, toolCalls, media);
@@ -78,4 +87,48 @@ public String toString() {
7887
+ this.textContent + '\'' + '}';
7988
}
8089

90+
public static final class Builder {
91+
92+
private String content;
93+
94+
private Map<String, Object> properties = Map.of();
95+
96+
private List<ToolCall> toolCalls = List.of();
97+
98+
private List<Media> media = List.of();
99+
100+
private String reasoningContent;
101+
102+
public Builder content(String content) {
103+
this.content = content;
104+
return this;
105+
}
106+
107+
public Builder properties(Map<String, Object> properties) {
108+
this.properties = properties;
109+
return this;
110+
}
111+
112+
public Builder toolCalls(List<ToolCall> toolCalls) {
113+
this.toolCalls = toolCalls;
114+
return this;
115+
}
116+
117+
public Builder media(List<Media> media) {
118+
this.media = media;
119+
return this;
120+
}
121+
122+
public Builder reasoningContent(String reasoningContent) {
123+
this.reasoningContent = reasoningContent;
124+
return this;
125+
}
126+
127+
public ZhiPuAiAssistantMessage build() {
128+
return new ZhiPuAiAssistantMessage(this.content, this.reasoningContent, this.properties, this.toolCalls,
129+
this.media);
130+
}
131+
132+
}
133+
81134
}

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+
ZhiPuAiAssistantMessage.Builder builder = new ZhiPuAiAssistantMessage.Builder();
234+
var assistantMessage = builder.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);
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.Builder builder = new ZhiPuAiAssistantMessage.Builder();
121+
ZhiPuAiAssistantMessage message = builder.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+
}

0 commit comments

Comments
 (0)