From eed1111fc9efb9cc6d485af56c373fbb7968361c Mon Sep 17 00:00:00 2001
From: robin <850379744@qq.com>
Date: Wed, 4 Oct 2023 11:34:45 +0800
Subject: [PATCH] add knowledge support

---
 .../chat2db/server/web/api/controller/ai/ChatController.java  | 2 +-
 .../server/web/api/controller/ai/KnowledgeController.java     | 3 ++-
 .../server/web/api/controller/ai/enums/PromptType.java        | 4 ++--
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/ChatController.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/ChatController.java
index dcc56ad8..bcf47c06 100644
--- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/ChatController.java
+++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/ChatController.java
@@ -454,7 +454,7 @@ public class ChatController {
      * @return
      */
     private String buildPrompt(ChatQueryRequest queryRequest) {
-        if (PromptType.QUESTION_ANSWERING.getCode().equals(queryRequest.getPromptType())) {
+        if (PromptType.TEXT_GENERATION.getCode().equals(queryRequest.getPromptType())) {
             return queryRequest.getMessage();
         }
 
diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/KnowledgeController.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/KnowledgeController.java
index 09b6d835..5eed686f 100644
--- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/KnowledgeController.java
+++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/KnowledgeController.java
@@ -6,6 +6,7 @@ import ai.chat2db.server.tools.common.exception.ParamBusinessException;
 import ai.chat2db.server.web.api.aspect.ConnectionInfoAspect;
 import ai.chat2db.server.web.api.controller.ai.DocParser.AbstractParser;
 import ai.chat2db.server.web.api.controller.ai.DocParser.PdfParse;
+import ai.chat2db.server.web.api.controller.ai.enums.PromptType;
 import ai.chat2db.server.web.api.controller.ai.fastchat.embeddings.FastChatEmbeddingResponse;
 import ai.chat2db.server.web.api.controller.ai.request.ChatQueryRequest;
 import ai.chat2db.server.web.api.http.GatewayClientService;
@@ -102,7 +103,7 @@ public class KnowledgeController extends ChatController {
 
         // search embedding
         DataResult<KnowledgeResponse> result = gatewayClientService.knowledgeVectorSearch(contentVector);
-
+        queryRequest.setPromptType(PromptType.TEXT_GENERATION.getCode());
         String prompt = queryRequest.getMessage();
         if (CollectionUtils.isNotEmpty(result.getData().getKnowledgeList())) {
             List<String> contents = new ArrayList<>();
diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/enums/PromptType.java b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/enums/PromptType.java
index 99cf5260..9e9745c7 100644
--- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/enums/PromptType.java
+++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/enums/PromptType.java
@@ -35,9 +35,9 @@ public enum PromptType implements BaseEnum<String> {
     SQL_2_SQL("进行SQL转换"),
 
     /**
-     * knowledge qa
+     * text generation
      */
-    QUESTION_ANSWERING("问答"),
+    TEXT_GENERATION("文本生成"),
     ;
 
     final String description;