Merge pull request #652 from chat2db/dev

Dev
CodePhiliaX · Oct 26, 2023 · 6460041 · 6460041
2 parents d747536 + 8351f2b
commit 6460041
Show file tree

Hide file tree

Showing 14 changed files with 396,465 additions and 66 deletions.
diff --git a/.gitignore b/.gitignore
@@ -11,6 +11,7 @@
 *.iws
 *.iml
 *.ipr
+.out
 
 ### Visual Studio Code ###
 package-lock.json
@@ -27,3 +28,4 @@ package-lock.json
 /chat2db-server/ali-dbhub-server-domain/ali-dbhub-server-domain-support/src/main/resources/lib/*
 /chat2db-server/ali-dbhub-server-domain/ali-dbhub-server-domain-support/lib/*
 /lib
+/out/*
diff --git a/chat2db-client/src/components/Console/index.tsx b/chat2db-client/src/components/Console/index.tsx
@@ -491,12 +491,12 @@ function Console(props: IProps, ref: ForwardedRef<IConsoleRef>) {
             }}
             selectedTables={selectedTables}
             onSelectTables={(tables: string[]) => {
-              if (tables.length > 8) {
-                message.warning({
-                  content: i18n('chat.input.tableSelect.error.TooManyTable'),
-                });
-                return;
-              }
+              // if (tables.length > 8) {
+              //   message.warning({
+              //     content: i18n('chat.input.tableSelect.error.TooManyTable'),
+              //   });
+              //   return;
+              // }
               setSelectedTables(tables);
             }}
             onClickRemainBtn={handleClickRemainBtn}

diff --git a/chat2db-server/chat2db-server-web/chat2db-server-web-api/pom.xml b/chat2db-server/chat2db-server-web/chat2db-server-web-api/pom.xml
@@ -82,6 +82,13 @@
             <groupId>ai.chat2db</groupId>
             <artifactId>chat2db-spi</artifactId>
         </dependency>
+
+
+        <dependency>
+            <groupId>org.ansj</groupId>
+            <artifactId>ansj_seg</artifactId>
+        </dependency>
+
     </dependencies>
 
 </project>
diff --git a/...-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/ChatController.java b/...-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/ChatController.java
@@ -1,15 +1,5 @@
 package ai.chat2db.server.web.api.controller.ai;
 
-import java.io.IOException;
-import java.math.BigDecimal;
-import java.time.Duration;
-import java.time.LocalDateTime;
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-import java.util.stream.Collectors;
-
 import ai.chat2db.server.domain.api.enums.AiSqlSourceEnum;
 import ai.chat2db.server.domain.api.model.Config;
 import ai.chat2db.server.domain.api.model.DataSource;
@@ -18,68 +8,68 @@
 import ai.chat2db.server.domain.api.service.ConfigService;
 import ai.chat2db.server.domain.api.service.DataSourceService;
 import ai.chat2db.server.domain.api.service.TableService;
+import ai.chat2db.server.tools.base.enums.WhiteListTypeEnum;
 import ai.chat2db.server.tools.base.wrapper.result.DataResult;
 import ai.chat2db.server.tools.common.exception.ParamBusinessException;
 import ai.chat2db.server.tools.common.util.EasyEnumUtils;
 import ai.chat2db.server.web.api.aspect.ConnectionInfoAspect;
 import ai.chat2db.server.web.api.controller.ai.azure.client.AzureOpenAIClient;
+import ai.chat2db.server.web.api.controller.ai.azure.listener.AzureOpenAIEventSourceListener;
 import ai.chat2db.server.web.api.controller.ai.azure.model.AzureChatMessage;
 import ai.chat2db.server.web.api.controller.ai.azure.model.AzureChatRole;
 import ai.chat2db.server.web.api.controller.ai.chat2db.client.Chat2dbAIClient;
 import ai.chat2db.server.web.api.controller.ai.claude.client.ClaudeAIClient;
+import ai.chat2db.server.web.api.controller.ai.claude.listener.ClaudeAIEventSourceListener;
 import ai.chat2db.server.web.api.controller.ai.claude.model.ClaudeChatCompletionsOptions;
 import ai.chat2db.server.web.api.controller.ai.claude.model.ClaudeChatMessage;
 import ai.chat2db.server.web.api.controller.ai.config.LocalCache;
 import ai.chat2db.server.web.api.controller.ai.converter.ChatConverter;
 import ai.chat2db.server.web.api.controller.ai.enums.PromptType;
-import ai.chat2db.server.web.api.controller.ai.azure.listener.AzureOpenAIEventSourceListener;
-import ai.chat2db.server.web.api.controller.ai.claude.listener.ClaudeAIEventSourceListener;
 import ai.chat2db.server.web.api.controller.ai.fastchat.client.FastChatAIClient;
 import ai.chat2db.server.web.api.controller.ai.fastchat.embeddings.FastChatEmbeddingResponse;
 import ai.chat2db.server.web.api.controller.ai.fastchat.listener.FastChatAIEventSourceListener;
 import ai.chat2db.server.web.api.controller.ai.fastchat.model.FastChatMessage;
 import ai.chat2db.server.web.api.controller.ai.fastchat.model.FastChatRole;
+import ai.chat2db.server.web.api.controller.ai.openai.client.OpenAIClient;
 import ai.chat2db.server.web.api.controller.ai.openai.listener.OpenAIEventSourceListener;
-import ai.chat2db.server.web.api.controller.ai.rest.listener.RestAIEventSourceListener;
 import ai.chat2db.server.web.api.controller.ai.request.ChatQueryRequest;
 import ai.chat2db.server.web.api.controller.ai.request.ChatRequest;
 import ai.chat2db.server.web.api.controller.ai.rest.client.RestAIClient;
+import ai.chat2db.server.web.api.controller.ai.rest.listener.RestAIEventSourceListener;
 import ai.chat2db.server.web.api.http.GatewayClientService;
 import ai.chat2db.server.web.api.http.model.EsTableSchema;
 import ai.chat2db.server.web.api.http.model.TableSchema;
 import ai.chat2db.server.web.api.http.request.EsTableSchemaRequest;
 import ai.chat2db.server.web.api.http.request.TableSchemaRequest;
+import ai.chat2db.server.web.api.http.request.WhiteListRequest;
 import ai.chat2db.server.web.api.http.response.EsTableSchemaResponse;
 import ai.chat2db.server.web.api.http.response.TableSchemaResponse;
 import ai.chat2db.server.web.api.util.ApplicationContextUtil;
-import ai.chat2db.server.web.api.controller.ai.openai.client.OpenAIClient;
-import ai.chat2db.spi.model.TableColumn;
+import ai.chat2db.server.web.api.util.SegmentUtils;
 import cn.hutool.core.util.StrUtil;
 import cn.hutool.json.JSONUtil;
 import com.alibaba.fastjson2.JSON;
 import com.google.common.collect.Lists;
-import com.google.common.collect.Maps;
-import com.unfbx.chatgpt.OpenAiApi;
 import com.unfbx.chatgpt.entity.chat.Message;
-import com.unfbx.chatgpt.entity.embeddings.Embedding;
-import com.unfbx.chatgpt.entity.embeddings.EmbeddingResponse;
-import io.reactivex.Single;
 import jakarta.annotation.Resource;
-import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.collections4.CollectionUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Value;
-import org.springframework.web.bind.annotation.CrossOrigin;
-import org.springframework.web.bind.annotation.GetMapping;
-import org.springframework.web.bind.annotation.PostMapping;
-import org.springframework.web.bind.annotation.RequestBody;
-import org.springframework.web.bind.annotation.RequestHeader;
-import org.springframework.web.bind.annotation.RequestMapping;
-import org.springframework.web.bind.annotation.RestController;
+import org.springframework.web.bind.annotation.*;
 import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
 
+import java.io.IOException;
+import java.math.BigDecimal;
+import java.time.Duration;
+import java.time.LocalDateTime;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.stream.Collectors;
+
 /**
  * 描述：
  *
@@ -501,7 +491,7 @@ private String buildPrompt(ChatQueryRequest queryRequest) {
             TableQueryParam queryParam = chatConverter.chat2tableQuery(queryRequest);
             properties = buildTableColumn(queryParam, queryRequest.getTableNames());
         } else {
-            properties = querySchemaByEs(queryRequest);
+            properties = mappingDatabaseSchema(queryRequest);
         }
         String prompt = queryRequest.getMessage();
         String promptType = StringUtils.isBlank(queryRequest.getPromptType()) ? PromptType.NL_2_SQL.getCode()
@@ -524,6 +514,26 @@ private String buildPrompt(ChatQueryRequest queryRequest) {
         return schemaProperty;
     }
 
+    /**
+     * query chat2db apikey
+     *
+     * @return
+     */
+    public String getApiKey() {
+        ConfigService configService = ApplicationContextUtil.getBean(ConfigService.class);
+        Config config = configService.find(RestAIClient.AI_SQL_SOURCE).getData();
+        String aiSqlSource = AiSqlSourceEnum.CHAT2DBAI.getCode();
+        // only sync for chat2db ai
+        if (Objects.isNull(config) || !aiSqlSource.equals(config.getContent())) {
+            return null;
+        }
+        Config keyConfig = configService.find(Chat2dbAIClient.CHAT2DB_OPENAI_KEY).getData();
+        if (Objects.isNull(keyConfig) || StringUtils.isBlank(keyConfig.getContent())) {
+            return null;
+        }
+        return keyConfig.getContent();
+    }
+
     /**
      * query database type
      *
@@ -540,6 +550,19 @@ public String queryDatabaseType(ChatQueryRequest queryRequest) {
         return dataSourceType;
     }
 
+    public String mappingDatabaseSchema(ChatQueryRequest queryRequest) {
+        String properties = "";
+        String apiKey = getApiKey();
+        if (StringUtils.isNotBlank(apiKey)) {
+            boolean res = gatewayClientService.checkInWhite(new WhiteListRequest(apiKey, WhiteListTypeEnum.VECTOR.getCode())).getData();
+            if (res) {
+                properties = queryDatabaseSchema(queryRequest) + querySchemaByEs(queryRequest);
+            }
+        } else {
+            properties = querySchemaByEs(queryRequest);
+        }
+        return properties;
+    }
 
     /**
      * query database schema
@@ -550,7 +573,9 @@ public String queryDatabaseType(ChatQueryRequest queryRequest) {
      */
     public String queryDatabaseSchema(ChatQueryRequest queryRequest) {
         // request embedding
-        FastChatEmbeddingResponse response = distributeAIEmbedding(queryRequest.getMessage());
+        String input = SegmentUtils.baseAnalysis(queryRequest.getMessage());
+        log.info("search message:{}", input);
+        FastChatEmbeddingResponse response = distributeAIEmbedding(input);
         List<List<BigDecimal>> contentVector = new ArrayList<>();
         contentVector.add(response.getData().get(0).getEmbedding());
 
@@ -574,7 +599,12 @@ public String queryDatabaseSchema(ChatQueryRequest queryRequest) {
                     schemas.add(data.getTableSchema());
                 }
             }
-            return JSON.toJSONString(schemas);
+            if (CollectionUtils.isEmpty(schemas)) {
+                return "";
+            }
+            String res = JSON.toJSONString(schemas);
+            log.info("search vector result:{}", res);
+            return res;
         } catch (Exception exception) {
             log.error("query table error, do nothing");
             return "";
@@ -609,7 +639,12 @@ public String querySchemaByEs(ChatQueryRequest queryRequest) {
                     schemas.add(data.getTableSchemaContent());
                 }
             }
-            return JSON.toJSONString(schemas);
+            if (CollectionUtils.isEmpty(schemas)) {
+                return "";
+            }
+            String res = JSON.toJSONString(schemas);
+            log.info("search es result:{}", res);
+            return res;
         } catch (Exception exception) {
             log.error("query es table error, do nothing");
             return "";

diff --git a/...er-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/EmbeddingController.java b/...er-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/EmbeddingController.java
@@ -242,20 +242,12 @@ public void syncTableVector(TableBriefQueryRequest param) throws Exception {
             return;
         }
 
-        ConfigService configService = ApplicationContextUtil.getBean(ConfigService.class);
-        Config config = configService.find(RestAIClient.AI_SQL_SOURCE).getData();
-        String aiSqlSource = AiSqlSourceEnum.CHAT2DBAI.getCode();
-        // only sync for chat2db ai
-        if (Objects.isNull(config) || !aiSqlSource.equals(config.getContent())) {
-            return;
-        }
-        Config keyConfig = configService.find(Chat2dbAIClient.CHAT2DB_OPENAI_KEY).getData();
-        if (Objects.isNull(keyConfig) || StringUtils.isBlank(keyConfig.getContent())) {
+        String apiKey = getApiKey();
+        if (StringUtils.isBlank(apiKey)) {
             return;
         }
 
         TableMilvusQueryRequest request = rdbWebConverter.request2request(param);
-        String apiKey = keyConfig.getContent();
         request.setApikey(apiKey);
 
         vectorParam.setApiKey(apiKey);
@@ -318,29 +310,13 @@ public void syncTableEs(TableBriefQueryRequest param) throws Exception {
             return;
         }
 
-        ConfigService configService = ApplicationContextUtil.getBean(ConfigService.class);
-        Config config = configService.find(RestAIClient.AI_SQL_SOURCE).getData();
-        String aiSqlSource = AiSqlSourceEnum.CHAT2DBAI.getCode();
-        // only sync for chat2db ai
-        if (Objects.isNull(config) || !aiSqlSource.equals(config.getContent())) {
-            return;
-        }
-        Config keyConfig = configService.find(Chat2dbAIClient.CHAT2DB_OPENAI_KEY).getData();
-        if (Objects.isNull(keyConfig) || StringUtils.isBlank(keyConfig.getContent())) {
-            return;
-        }
-
-        String apiKey = keyConfig.getContent();
-        TableVectorParam vectorParam = rdbWebConverter.param2param(param);
-        vectorParam.setApiKey(apiKey);
-        DataResult<Boolean> result = tableService.checkTableVector(vectorParam);
-        if (result.getData()) {
+        String apiKey = getApiKey();
+        if (StringUtils.isBlank(apiKey)) {
             return;
         }
 
         esParam.setApiKey(apiKey);
         es(esParam);
-        tableService.saveTableVector(vectorParam);
     }
 
     /**

diff --git a/...rver-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/RdbDdlController.java b/...rver-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/RdbDdlController.java
@@ -77,6 +77,7 @@ public WebPageResult<TableVO> list(@Valid TableBriefQueryRequest request) {
         singleThreadExecutor.submit(() -> {
             try {
                 Chat2DBContext.putContext(connectInfo);
+                syncTableVector(request);
                 syncTableEs(request);
             } catch (Exception e) {
                 log.error("sync table vector error", e);

diff --git a/...erver-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/TableController.java b/...erver-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/TableController.java
@@ -68,6 +68,7 @@ public WebPageResult<TableVO> list(@Valid TableBriefQueryRequest request) {
         singleThreadExecutor.submit(() -> {
             try {
                 Chat2DBContext.putContext(connectInfo);
+                syncTableVector(request);
                 syncTableEs(request);
             } catch (Exception e) {
                 log.error("sync table vector error", e);