|
@@ -32,13 +32,12 @@ import reactor.core.publisher.Flux;
|
|
|
|
|
|
import java.util.ArrayList;
|
|
|
import java.util.List;
|
|
|
-import java.util.Objects;
|
|
|
|
|
|
import static cn.iocoder.yudao.framework.common.pojo.CommonResult.error;
|
|
|
import static cn.iocoder.yudao.framework.common.pojo.CommonResult.success;
|
|
|
|
|
|
/**
|
|
|
- * AI 写作 Service 实现类
|
|
|
+ * AI 思维导图 Service 实现类
|
|
|
*
|
|
|
* @author xiaoxin
|
|
|
*/
|
|
@@ -58,30 +57,28 @@ public class AiMindMapServiceImpl implements AiMindMapService {
|
|
|
|
|
|
@Override
|
|
|
public Flux<CommonResult<String>> generateMindMap(AiMindMapGenerateReqVO generateReqVO, Long userId) {
|
|
|
- // 1 获取脑图模型 尝试获取思维导图助手角色,如果没有则使用默认模型
|
|
|
- AiChatRoleDO mindMapRole = CollUtil.getFirst(chatRoleService.getChatRoleListByName(AiChatRoleEnum.AI_MIND_MAP_ROLE.getName()));
|
|
|
+ // 1. 获取脑图模型。尝试获取思维导图助手角色,如果没有则使用默认模型
|
|
|
+ AiChatRoleDO role = CollUtil.getFirst(
|
|
|
+ chatRoleService.getChatRoleListByName(AiChatRoleEnum.AI_MIND_MAP_ROLE.getName()));
|
|
|
// 1.1 获取脑图执行模型
|
|
|
- AiChatModelDO model = getModel(mindMapRole);
|
|
|
+ AiChatModelDO model = getModel(role);
|
|
|
// 1.2 获取角色设定消息
|
|
|
- String systemMessage = Objects.nonNull(mindMapRole) && StrUtil.isNotBlank(mindMapRole.getSystemMessage())
|
|
|
- ? mindMapRole.getSystemMessage() : AiChatRoleEnum.AI_MIND_MAP_ROLE.getSystemMessage();
|
|
|
+ String systemMessage = role != null && StrUtil.isNotBlank(role.getSystemMessage())
|
|
|
+ ? role.getSystemMessage() : AiChatRoleEnum.AI_MIND_MAP_ROLE.getSystemMessage();
|
|
|
// 1.3 校验平台
|
|
|
AiPlatformEnum platform = AiPlatformEnum.validatePlatform(model.getPlatform());
|
|
|
ChatModel chatModel = apiKeyService.getChatModel(model.getKeyId());
|
|
|
|
|
|
- // 2 插入思维导图信息
|
|
|
+ // 2. 插入思维导图信息
|
|
|
AiMindMapDO mindMapDO = BeanUtils.toBean(generateReqVO, AiMindMapDO.class,
|
|
|
mindMap -> mindMap.setUserId(userId).setModel(model.getModel()).setPlatform(platform.getPlatform()));
|
|
|
mindMapMapper.insert(mindMapDO);
|
|
|
|
|
|
- ChatOptions chatOptions = AiUtils.buildChatOptions(platform, model.getModel(), model.getTemperature(), model.getMaxTokens());
|
|
|
- // 3.1 角色设定
|
|
|
- List<Message> chatMessages = buildMessages(generateReqVO, systemMessage);
|
|
|
- // 3.3 构建提示词
|
|
|
- Prompt prompt = new Prompt(chatMessages, chatOptions);
|
|
|
-
|
|
|
+ // 3.1 构建 Prompt,并进行调用
|
|
|
+ Prompt prompt = buildPrompt(generateReqVO, model, systemMessage);
|
|
|
Flux<ChatResponse> streamResponse = chatModel.stream(prompt);
|
|
|
- // 3.4 流式返回
|
|
|
+
|
|
|
+ // 3.2 流式返回
|
|
|
StringBuffer contentBuffer = new StringBuffer();
|
|
|
return streamResponse.map(chunk -> {
|
|
|
String newContent = chunk.getResult() != null ? chunk.getResult().getOutput().getContent() : null;
|
|
@@ -102,24 +99,32 @@ public class AiMindMapServiceImpl implements AiMindMapService {
|
|
|
|
|
|
}
|
|
|
|
|
|
+ private Prompt buildPrompt(AiMindMapGenerateReqVO generateReqVO, AiChatModelDO model, String systemMessage) {
|
|
|
+ // 1. 构建 message 列表
|
|
|
+ List<Message> chatMessages = buildMessages(generateReqVO, systemMessage);
|
|
|
+ // 2. 构建 options 对象
|
|
|
+ AiPlatformEnum platform = AiPlatformEnum.validatePlatform(model.getPlatform());
|
|
|
+ ChatOptions options = AiUtils.buildChatOptions(platform, model.getModel(), model.getTemperature(), model.getMaxTokens());
|
|
|
+ return new Prompt(chatMessages, options);
|
|
|
+ }
|
|
|
+
|
|
|
private static List<Message> buildMessages(AiMindMapGenerateReqVO generateReqVO, String systemMessage) {
|
|
|
List<Message> chatMessages = new ArrayList<>();
|
|
|
+ // 1. 角色设定
|
|
|
if (StrUtil.isNotBlank(systemMessage)) {
|
|
|
- // 1.1 角色设定
|
|
|
chatMessages.add(new SystemMessage(systemMessage));
|
|
|
}
|
|
|
- // 1.2 用户输入
|
|
|
+ // 2. 用户输入
|
|
|
chatMessages.add(new UserMessage(generateReqVO.getPrompt()));
|
|
|
return chatMessages;
|
|
|
}
|
|
|
|
|
|
- // TODO 芋艿:这里脑图、写作都用到了,是不是可以抽哪里去
|
|
|
- private AiChatModelDO getModel(AiChatRoleDO chatRoleDO) {
|
|
|
+ private AiChatModelDO getModel(AiChatRoleDO role) {
|
|
|
AiChatModelDO model = null;
|
|
|
- if (Objects.nonNull(chatRoleDO) && Objects.nonNull(chatRoleDO.getModelId())) {
|
|
|
- model = chatModalService.getChatModel(chatRoleDO.getModelId());
|
|
|
+ if (role != null && role.getModelId() != null) {
|
|
|
+ model = chatModalService.getChatModel(role.getModelId());
|
|
|
}
|
|
|
- if (Objects.isNull(model)) {
|
|
|
+ if (model != null) {
|
|
|
model = chatModalService.getRequiredDefaultChatModel();
|
|
|
}
|
|
|
Assert.notNull(model, "[AI] 获取不到模型");
|