Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: Fixed a bug in the augmentSystemMessage method of the Prompt, where an extra system message was incorrectly added when the system message was not the first one in the message list.

Browse Source 
Signed-off-by: Sun Yuhan <1085481446@qq.com>
This commit is contained in:
Sun Yuhan
2025-05-21 13:19:35 +08:00
committed by Ilayaperumal Gopinathan
parent dd6c0a9530
commit 623e705336
2 changed files with 30 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
spring-ai-model/src/main/java/org/springframework/ai/chat/prompt/Prompt.java
View File

@@ -198,21 +198,21 @@ public class Prompt implements ModelRequest<List<Message>> {
* @return a new {@link Prompt} instance with the augmented system message.
*/
public Prompt augmentSystemMessage(Function<SystemMessage, SystemMessage> systemMessageAugmenter) {
var messagesCopy = new ArrayList<>(this.messages);
for (int i = 0; i <= this.messages.size() - 1; i++) {
boolean found = false;
for (int i = 0; i < messagesCopy.size(); i++) {
Message message = messagesCopy.get(i);
if (message instanceof SystemMessage systemMessage) {
messagesCopy.set(i, systemMessageAugmenter.apply(systemMessage));
found = true;
break;
}
if (i == 0) {
// If no system message is found, create a new one with the provided text
// and add it as the first item in the list.
messagesCopy.add(0, systemMessageAugmenter.apply(new SystemMessage("")));
}
}
if (!found) {
// If no system message is found, create a new one with the provided text
// and add it as the first item in the list.
messagesCopy.add(0, systemMessageAugmenter.apply(new SystemMessage("")));
}
return new Prompt(messagesCopy, null == this.chatOptions ? null : this.chatOptions.copy());
}

22
spring-ai-model/src/test/java/org/springframework/ai/chat/prompt/PromptTests.java
View File

@@ -239,4 +239,26 @@ class PromptTests {
assertThat(prompt.getSystemMessage().getText()).isEqualTo("");
}
@Test
void augmentSystemMessageWhenNotFirst() {
Message[] messages = { new UserMessage("Hi"), new SystemMessage("Hello") };
Prompt prompt = Prompt.builder().messages(messages).build();
assertThat(prompt.getSystemMessage()).isNotNull();
assertThat(prompt.getUserMessage()).isNotNull();
assertThat(prompt.getUserMessage().getText()).isEqualTo("Hi");
assertThat(prompt.getSystemMessage().getText()).isEqualTo("Hello");
Prompt copy = prompt.augmentSystemMessage(message -> message.mutate().text("How are you?").build());
assertThat(copy.getSystemMessage()).isNotNull();
assertThat(copy.getInstructions().size()).isEqualTo(messages.length);
assertThat(copy.getSystemMessage().getText()).isEqualTo("How are you?");
assertThat(prompt.getSystemMessage()).isNotNull();
assertThat(prompt.getUserMessage()).isNotNull();
assertThat(prompt.getUserMessage().getText()).isEqualTo("Hi");
assertThat(prompt.getSystemMessage().getText()).isEqualTo("Hello");
}
}