From 7a6985dd8691ab940ac1d04000727c4110d19d3b Mon Sep 17 00:00:00 2001 From: Jast <745925668@qq.com> Date: Mon, 9 Jun 2025 03:19:02 +0800 Subject: [PATCH 1/2] feat: Add SimpleVectorStoreService for managing vector data and schema initialization --- .../ai/service/SimpleVectorStoreService.java | 189 ++++++++++++++++++ 1 file changed, 189 insertions(+) create mode 100644 spring-ai-alibaba-nl2sql/management/src/main/java/com/alibaba/cloud/ai/service/SimpleVectorStoreService.java diff --git a/spring-ai-alibaba-nl2sql/management/src/main/java/com/alibaba/cloud/ai/service/SimpleVectorStoreService.java b/spring-ai-alibaba-nl2sql/management/src/main/java/com/alibaba/cloud/ai/service/SimpleVectorStoreService.java new file mode 100644 index 0000000000..3e3fabcade --- /dev/null +++ b/spring-ai-alibaba-nl2sql/management/src/main/java/com/alibaba/cloud/ai/service/SimpleVectorStoreService.java @@ -0,0 +1,189 @@ +package com.alibaba.cloud.ai.service; + +import com.alibaba.cloud.ai.dashscope.api.DashScopeApi; +import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingModel; +import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingOptions; +import com.alibaba.cloud.ai.dbconnector.DbAccessor; +import com.alibaba.cloud.ai.dbconnector.DbConfig; +import com.alibaba.cloud.ai.dbconnector.bo.ColumnInfoBO; +import com.alibaba.cloud.ai.dbconnector.bo.DbQueryParameter; +import com.alibaba.cloud.ai.dbconnector.bo.ForeignKeyInfoBO; +import com.alibaba.cloud.ai.dbconnector.bo.TableInfoBO; +import com.alibaba.cloud.ai.request.DeleteRequest; +import com.alibaba.cloud.ai.request.SchemaInitRequest; +import com.google.gson.Gson; +import org.springframework.ai.document.Document; +import org.springframework.ai.document.MetadataMode; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.SimpleVectorStore; +import org.springframework.ai.vectorstore.filter.Filter; +import org.springframework.ai.vectorstore.filter.FilterExpressionBuilder; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.stereotype.Service; + +import java.util.*; +import java.util.stream.Collectors; + +@Service +public class SimpleVectorStoreService { + + private final SimpleVectorStore vectorStore; + + private final Gson gson; + + private final DbAccessor dbAccessor; + + private final DbConfig dbConfig; + + @Autowired + public SimpleVectorStoreService(@Value("${spring.ai.dashscope.api-key:default_api_key}") String apiKey, Gson gson, + DbAccessor dbAccessor, DbConfig dbConfig) { + this.gson = gson; + this.dbAccessor = dbAccessor; + this.dbConfig = dbConfig; + + DashScopeApi dashScopeApi = DashScopeApi.builder().apiKey(apiKey).build(); + DashScopeEmbeddingModel dashScopeEmbeddingModel = new DashScopeEmbeddingModel(dashScopeApi, MetadataMode.EMBED, + DashScopeEmbeddingOptions.builder().withModel("text-embedding-v2").build()); + this.vectorStore = SimpleVectorStore.builder(dashScopeEmbeddingModel).build(); + } + + /** + * 初始化数据库 schema 到向量库 + * @param schemaInitRequest schema 初始化请求 + * @throws Exception 如果发生错误 + */ + public Boolean schema(SchemaInitRequest schemaInitRequest) throws Exception { + DbConfig dbConfig = schemaInitRequest.getDbConfig(); + DbQueryParameter dqp = DbQueryParameter.from(dbConfig) + .setSchema(dbConfig.getSchema()) + .setTables(schemaInitRequest.getTables()); + + DeleteRequest deleteRequest = new DeleteRequest(); + deleteRequest.setVectorType("column"); + deleteDocuments(deleteRequest); + deleteRequest.setVectorType("table"); + deleteDocuments(deleteRequest); + + List foreignKeyInfoBOS = dbAccessor.showForeignKeys(dbConfig, dqp); + Map> foreignKeyMap = buildForeignKeyMap(foreignKeyInfoBOS); + + List tableInfoBOS = dbAccessor.fetchTables(dbConfig, dqp); + for (TableInfoBO tableInfoBO : tableInfoBOS) { + processTable(tableInfoBO, dqp, dbConfig); + } + + List columnDocuments = tableInfoBOS.stream().flatMap(table -> { + try { + return dbAccessor.showColumns(dbConfig, dqp).stream().map(column -> convertToDocument(table, column)); + } + catch (Exception e) { + throw new RuntimeException(e); + } + }).collect(Collectors.toList()); + + vectorStore.add(columnDocuments); + + List tableDocuments = tableInfoBOS.stream() + .map(this::convertTableToDocument) + .collect(Collectors.toList()); + + vectorStore.add(tableDocuments); + + return true; + } + + private void processTable(TableInfoBO tableInfoBO, DbQueryParameter dqp, DbConfig dbConfig) throws Exception { + dqp.setTable(tableInfoBO.getName()); + List columnInfoBOS = dbAccessor.showColumns(dbConfig, dqp); + for (ColumnInfoBO columnInfoBO : columnInfoBOS) { + dqp.setColumn(columnInfoBO.getName()); + List sampleColumn = dbAccessor.sampleColumn(dbConfig, dqp); + sampleColumn = Optional.ofNullable(sampleColumn) + .orElse(new ArrayList<>()) + .stream() + .filter(Objects::nonNull) + .distinct() + .limit(3) + .filter(s -> s.length() <= 100) + .toList(); + + columnInfoBO.setTableName(tableInfoBO.getName()); + columnInfoBO.setSamples(gson.toJson(sampleColumn)); + } + + ColumnInfoBO primaryColumnDO = columnInfoBOS.stream() + .filter(ColumnInfoBO::isPrimary) + .findFirst() + .orElse(new ColumnInfoBO()); + + tableInfoBO.setPrimaryKey(primaryColumnDO.getName()); + tableInfoBO.setForeignKey(String.join("、", buildForeignKeyList(tableInfoBO.getName()))); + } + + public Document convertToDocument(TableInfoBO tableInfoBO, ColumnInfoBO columnInfoBO) { + String text = Optional.ofNullable(columnInfoBO.getDescription()).orElse(columnInfoBO.getName()); + Map metadata = Map.of("name", columnInfoBO.getName(), "tableName", tableInfoBO.getName(), + "description", Optional.ofNullable(columnInfoBO.getDescription()).orElse(""), "type", + columnInfoBO.getType(), "primary", columnInfoBO.isPrimary(), "notnull", columnInfoBO.isNotnull(), + "vectorType", "column"); + if (columnInfoBO.getSamples() != null) { + metadata.put("samples", columnInfoBO.getSamples()); + } + return new Document(columnInfoBO.getName(), text, metadata); + } + + private List buildForeignKeyList(String tableName) { + return new ArrayList<>(); + } + + public Document convertTableToDocument(TableInfoBO tableInfoBO) { + String text = Optional.ofNullable(tableInfoBO.getDescription()).orElse(tableInfoBO.getName()); + Map metadata = Map.of("schema", Optional.ofNullable(tableInfoBO.getSchema()).orElse(""), "name", + tableInfoBO.getName(), "description", Optional.ofNullable(tableInfoBO.getDescription()).orElse(""), + "foreignKey", tableInfoBO.getForeignKey(), "primaryKey", tableInfoBO.getPrimaryKey(), "vectorType", + "table"); + return new Document(tableInfoBO.getName(), text, metadata); + } + + private Map> buildForeignKeyMap(List foreignKeyInfoBOS) { + Map> foreignKeyMap = new HashMap<>(); + for (ForeignKeyInfoBO fk : foreignKeyInfoBOS) { + String key = fk.getTable() + "." + fk.getColumn() + "=" + fk.getReferencedTable() + "." + + fk.getReferencedColumn(); + + foreignKeyMap.computeIfAbsent(fk.getTable(), k -> new ArrayList<>()).add(key); + foreignKeyMap.computeIfAbsent(fk.getReferencedTable(), k -> new ArrayList<>()).add(key); + } + return foreignKeyMap; + } + + /** + * 删除指定条件的向量数据 + * @param deleteRequest 删除请求 + * @return 是否删除成功 + */ + public Boolean deleteDocuments(DeleteRequest deleteRequest) throws Exception { + try { + if (deleteRequest.getId() != null && !deleteRequest.getId().isEmpty()) { + vectorStore.delete(Arrays.asList("comment_count")); + } + else if (deleteRequest.getVectorType() != null && !deleteRequest.getVectorType().isEmpty()) { + FilterExpressionBuilder b = new FilterExpressionBuilder(); + Filter.Expression expression = b.eq("vectorType", "column").build(); + List documents = vectorStore.similaritySearch( + SearchRequest.builder().topK(Integer.MAX_VALUE).filterExpression(expression).build()); + vectorStore.delete(documents.stream().map(Document::getId).toList()); + } + else { + throw new IllegalArgumentException("Either id or vectorType must be specified."); + } + return true; + } + catch (Exception e) { + throw new Exception("Failed to delete collection data by filterExpression: " + e.getMessage(), e); + } + } + +} From 66c6983d702e68ac86aa633c076c277e11bc7928 Mon Sep 17 00:00:00 2001 From: Jast <745925668@qq.com> Date: Mon, 9 Jun 2025 03:22:51 +0800 Subject: [PATCH 2/2] format --- .../ai/service/SimpleVectorStoreService.java | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/spring-ai-alibaba-nl2sql/management/src/main/java/com/alibaba/cloud/ai/service/SimpleVectorStoreService.java b/spring-ai-alibaba-nl2sql/management/src/main/java/com/alibaba/cloud/ai/service/SimpleVectorStoreService.java index 3e3fabcade..42691148e7 100644 --- a/spring-ai-alibaba-nl2sql/management/src/main/java/com/alibaba/cloud/ai/service/SimpleVectorStoreService.java +++ b/spring-ai-alibaba-nl2sql/management/src/main/java/com/alibaba/cloud/ai/service/SimpleVectorStoreService.java @@ -1,3 +1,18 @@ +/* + * Copyright 2024-2025 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 com.alibaba.cloud.ai.service; import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;