From 8936aac1e2eb371f33385537c3085d69317b07e5 Mon Sep 17 00:00:00 2001 From: Geng Date: Thu, 25 Jul 2024 05:00:15 +0800 Subject: [PATCH] fix: QianFan role based chat model & unit tests (#1029) --- .../ai/qianfan/QianFanChatModel.java | 3 +- .../ai/qianfan/chat/QianFanChatModelIT.java | 94 +++++++++++++++++++ 2 files changed, 96 insertions(+), 1 deletion(-) create mode 100644 models/spring-ai-qianfan/src/test/java/org/springframework/ai/qianfan/chat/QianFanChatModelIT.java diff --git a/models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/QianFanChatModel.java b/models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/QianFanChatModel.java index ca649279a..c5be46e8b 100644 --- a/models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/QianFanChatModel.java +++ b/models/spring-ai-qianfan/src/main/java/org/springframework/ai/qianfan/QianFanChatModel.java @@ -166,6 +166,7 @@ public class QianFanChatModel implements ChatModel, StreamingChatModel { ChatCompletionMessage.Role.valueOf(m.getMessageType().name()))) .toList(); var systemMessageList = chatCompletionMessages.stream().filter(msg -> msg.role() == Role.SYSTEM).toList(); + var userMessageList = chatCompletionMessages.stream().filter(msg -> msg.role() != Role.SYSTEM).toList(); if (systemMessageList.size() > 1) { throw new IllegalArgumentException("Only one system message is allowed in the prompt"); @@ -173,7 +174,7 @@ public class QianFanChatModel implements ChatModel, StreamingChatModel { var systemMessage = systemMessageList.isEmpty() ? null : systemMessageList.get(0).content(); - var request = new ChatCompletionRequest(chatCompletionMessages, systemMessage, stream); + var request = new ChatCompletionRequest(userMessageList, systemMessage, stream); if (this.defaultOptions != null) { request = ModelOptionsUtils.merge(this.defaultOptions, request, ChatCompletionRequest.class); diff --git a/models/spring-ai-qianfan/src/test/java/org/springframework/ai/qianfan/chat/QianFanChatModelIT.java b/models/spring-ai-qianfan/src/test/java/org/springframework/ai/qianfan/chat/QianFanChatModelIT.java new file mode 100644 index 000000000..1ed3d9d1a --- /dev/null +++ b/models/spring-ai-qianfan/src/test/java/org/springframework/ai/qianfan/chat/QianFanChatModelIT.java @@ -0,0 +1,94 @@ +/* + * Copyright 2023 - 2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.qianfan.chat; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariables; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.model.ChatModel; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.model.Generation; +import org.springframework.ai.chat.model.StreamingChatModel; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.chat.prompt.SystemPromptTemplate; +import org.springframework.ai.qianfan.QianFanTestConfiguration; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.core.io.Resource; +import reactor.core.publisher.Flux; + +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Geng Rong + */ +@SpringBootTest(classes = QianFanTestConfiguration.class) +@EnabledIfEnvironmentVariables(value = { @EnabledIfEnvironmentVariable(named = "QIANFAN_API_KEY", matches = ".+"), + @EnabledIfEnvironmentVariable(named = "QIANFAN_SECRET_KEY", matches = ".+") }) +class QianFanChatModelIT { + + @Autowired + protected ChatModel chatModel; + + @Autowired + protected StreamingChatModel streamingChatModel; + + @Value("classpath:/prompts/system-message.st") + private Resource systemResource; + + @Test + void roleTest() { + UserMessage userMessage = new UserMessage( + "Tell me about three famous pirates from the Golden Age of Piracy in english, focusing on their original nicknames and what they did."); + SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemResource); + Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", "Bob", "voice", "pirate")); + Prompt prompt = new Prompt(List.of(userMessage, systemMessage)); + ChatResponse response = chatModel.call(prompt); + assertThat(response.getResults()).hasSize(1); + assertThat(response.getResults().get(0).getOutput().getContent()).contains("Blackbeard"); + } + + @Test + void streamRoleTest() { + UserMessage userMessage = new UserMessage( + "Tell me about three famous pirates from the Golden Age of Piracy in english, focusing on their original nicknames and what they did."); + SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemResource); + Message systemMessage = systemPromptTemplate.createMessage(Map.of("name", "Bob", "voice", "pirate")); + Prompt prompt = new Prompt(List.of(userMessage, systemMessage)); + Flux flux = streamingChatModel.stream(prompt); + + List responses = flux.collectList().block(); + assertThat(responses.size()).isGreaterThan(1); + + String stitchedResponseContent = responses.stream() + .map(ChatResponse::getResults) + .flatMap(List::stream) + .map(Generation::getOutput) + .map(AssistantMessage::getContent) + .collect(Collectors.joining()); + + assertThat(stitchedResponseContent).contains("Blackbeard"); + } + +} \ No newline at end of file