feat:新增调用DALL-E模型的generation功能

ashinnotfound · Aug 10, 2023 · 1922684 · 1922684
1 parent 86e7995
commit 1922684
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 11 deletions.
diff --git a/src/main/java/com/ashin/entity/bo/ChatBO.java b/src/main/java/com/ashin/entity/bo/ChatBO.java
@@ -18,4 +18,8 @@ public class ChatBO {
      * 问题
      */
     private String prompt;
+    /**
+     * 是否ai画图功能
+     */
+    private boolean isAiDraw;
 }
diff --git a/src/main/java/com/ashin/handler/QqMessageHandler.java b/src/main/java/com/ashin/handler/QqMessageHandler.java
@@ -73,6 +73,7 @@ private void response(@NotNull MessageEvent event, ChatBO chatBO, String prompt)
             String response;
             try {
                 chatBO.setPrompt(prompt);
+                chatBO.setAiDraw(prompt.startsWith(keywordConfig.getImageGeneration()));
                 response = interactService.chat(chatBO);
             }catch (ChatException e){
                 response = e.getMessage();

diff --git a/src/main/java/com/ashin/handler/WechatMessageHandler.java b/src/main/java/com/ashin/handler/WechatMessageHandler.java
@@ -1,7 +1,6 @@
 package com.ashin.handler;
 
 import com.ashin.config.KeywordConfig;
-import com.ashin.config.WechatConfig;
 import com.ashin.entity.bo.ChatBO;
 import com.ashin.exception.ChatException;
 import com.ashin.service.InteractService;
@@ -53,6 +52,7 @@ private String textResponse(String userName, String content) {
             ChatBO chatBO = new ChatBO();
             chatBO.setPrompt(content);
             chatBO.setSessionId(userName);
+            chatBO.setAiDraw(content.startsWith(keywordConfig.getImageGeneration()));
             String response;
             try {
                 response = interactService.chat(chatBO);

diff --git a/src/main/java/com/ashin/service/impl/InteractServiceImpl.java b/src/main/java/com/ashin/service/impl/InteractServiceImpl.java
@@ -1,12 +1,14 @@
 package com.ashin.service.impl;
 
+import com.ashin.config.KeywordConfig;
 import com.ashin.entity.bo.ChatBO;
 import com.ashin.exception.ChatException;
 import com.ashin.service.InteractService;
 import com.ashin.util.BotUtil;
 import com.theokanning.openai.OpenAiHttpException;
 import com.theokanning.openai.completion.chat.ChatCompletionRequest;
 import com.theokanning.openai.completion.chat.ChatMessage;
+import com.theokanning.openai.image.CreateImageRequest;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.stereotype.Service;
 
@@ -25,34 +27,44 @@ public class InteractServiceImpl implements InteractService {
 
     @Resource
     private BotUtil botUtil;
+    @Resource
+    private KeywordConfig keywordConfig;
 
     @Override
     public String chat(ChatBO chatBO) throws ChatException {
 
         List<ChatMessage> prompt = botUtil.buildPrompt(chatBO.getSessionId(), chatBO.getPrompt());
-        if (prompt == null){
+        if (prompt == null) {
             throw new ChatException("提问失败: 提问内容过长");
         }
 
-        //向gpt提问
-        ChatCompletionRequest completionRequest = botUtil.getCompletionRequestBuilder().messages(prompt).build();
         ChatMessage answer = null;
         try {
-            answer = botUtil.getOpenAiService().createChatCompletion(completionRequest).getChoices().get(0).getMessage();
-        }catch (OpenAiHttpException e){
+            //向gpt提问
+            if (chatBO.isAiDraw()) {
+                String description = chatBO.getPrompt().replaceFirst(keywordConfig.getImageGeneration() + " ", "");
+                CreateImageRequest createImageRequest = CreateImageRequest.builder().n(1).size("1024x1024").prompt(description).build();
+                answer = new ChatMessage();
+                answer.setRole("assistant");
+                answer.setContent(botUtil.getOpenAiService().createImage(createImageRequest).getData().get(0).getUrl());
+            } else {
+                ChatCompletionRequest completionRequest = botUtil.getCompletionRequestBuilder().messages(prompt).build();
+                answer = botUtil.getOpenAiService().createChatCompletion(completionRequest).getChoices().get(0).getMessage();
+            }
+        } catch (OpenAiHttpException e) {
             log.error("向gpt提问失败，提问内容：{}，\n原因：{}\n", chatBO.getPrompt(), e.getMessage(), e);
             // https://platform.openai.com/docs/guides/error-codes/api-errors
-            if (401 == e.statusCode){
+            if (401 == e.statusCode) {
                 throw new ChatException("提问失败: 无效的apikey, 请确保apikey正确且你拥有权限");
-            }else if(429 == e.statusCode){
+            } else if (429 == e.statusCode) {
                 throw new ChatException("提问过于频繁 或者 apikey余额不足");
-            }else if(500 == e.statusCode) {
+            } else if (500 == e.statusCode) {
                 throw new ChatException("提问失败: openai服务异常, 请稍后重试");
-            }else if(503 == e.statusCode) {
+            } else if (503 == e.statusCode) {
                 throw new ChatException("提问失败: openai服务过载, 请稍后重试");
             }
         }
-        if (null == answer){
+        if (null == answer) {
             throw new ChatException("提问失败: 未知错误, 若频繁出现请提issue");
         }