diff --git a/ruoyi-common/ruoyi-common-core/src/main/java/org/ruoyi/common/core/config/VectorStoreProperties.java b/ruoyi-common/ruoyi-common-core/src/main/java/org/ruoyi/common/core/config/VectorStoreProperties.java index 98f3ddc4..b4bb135d 100644 --- a/ruoyi-common/ruoyi-common-core/src/main/java/org/ruoyi/common/core/config/VectorStoreProperties.java +++ b/ruoyi-common/ruoyi-common-core/src/main/java/org/ruoyi/common/core/config/VectorStoreProperties.java @@ -17,7 +17,7 @@ public class VectorStoreProperties { /** * 向量库类型 */ - private String type = "weaviate"; + private String type; /** * Weaviate配置 @@ -34,17 +34,17 @@ public class VectorStoreProperties { /** * 协议 */ - private String protocol = "http"; + private String protocol; /** * 主机地址 */ - private String host = "localhost:8080"; + private String host; /** * 类名 */ - private String classname = "Document"; + private String classname; } @Data @@ -52,11 +52,11 @@ public class VectorStoreProperties { /** * 连接URL */ - private String url = "http://localhost:19530"; + private String url; /** * 集合名称 */ - private String collectionname = "knowledge_base"; + private String collectionname; } } \ No newline at end of file diff --git a/ruoyi-modules-api/ruoyi-knowledge-api/src/main/java/org/ruoyi/domain/bo/StoreEmbeddingBo.java b/ruoyi-modules-api/ruoyi-knowledge-api/src/main/java/org/ruoyi/domain/bo/StoreEmbeddingBo.java index 2b87ce05..eedfe4d6 100644 --- a/ruoyi-modules-api/ruoyi-knowledge-api/src/main/java/org/ruoyi/domain/bo/StoreEmbeddingBo.java +++ b/ruoyi-modules-api/ruoyi-knowledge-api/src/main/java/org/ruoyi/domain/bo/StoreEmbeddingBo.java @@ -32,9 +32,9 @@ public class StoreEmbeddingBo { private List fids; /** - * 向量库模型名称 + * 向量库名称 */ - private String vectorModelName; + private String vectorStoreName; /** * 向量化模型id diff --git a/ruoyi-modules-api/ruoyi-knowledge-api/src/main/java/org/ruoyi/service/strategy/impl/MilvusVectorStoreStrategy.java b/ruoyi-modules-api/ruoyi-knowledge-api/src/main/java/org/ruoyi/service/strategy/impl/MilvusVectorStoreStrategy.java index d262a7e3..9d60ce8d 100644 --- a/ruoyi-modules-api/ruoyi-knowledge-api/src/main/java/org/ruoyi/service/strategy/impl/MilvusVectorStoreStrategy.java +++ b/ruoyi-modules-api/ruoyi-knowledge-api/src/main/java/org/ruoyi/service/strategy/impl/MilvusVectorStoreStrategy.java @@ -23,12 +23,10 @@ import org.springframework.stereotype.Component; import java.util.ArrayList; import java.util.List; import java.util.stream.IntStream; +// 新增导入 +import java.util.Map; +import java.util.concurrent.ConcurrentHashMap; -/** - * Milvus向量库策略实现 - * - * @author Yzm - */ @Slf4j @Component public class MilvusVectorStoreStrategy extends AbstractVectorStoreStrategy { @@ -37,6 +35,27 @@ public class MilvusVectorStoreStrategy extends AbstractVectorStoreStrategy { super(vectorStoreProperties); } + // 缓存不同集合与 autoFlush 配置的 Milvus 连接 + private final Map> storeCache = new ConcurrentHashMap<>(); + + private EmbeddingStore getMilvusStore(String collectionName, boolean autoFlushOnInsert) { + String key = collectionName + "|" + autoFlushOnInsert; + return storeCache.computeIfAbsent(key, k -> + MilvusEmbeddingStore.builder() + .uri(vectorStoreProperties.getMilvus().getUrl()) + .collectionName(collectionName) + .dimension(2048) + .indexType(IndexType.IVF_FLAT) + .metricType(MetricType.L2) + .autoFlushOnInsert(autoFlushOnInsert) + .idFieldName("id") + .textFieldName("text") + .metadataFieldName("metadata") + .vectorFieldName("vector") + .build() + ); + } + @Override public String getVectorStoreType() { return "milvus"; @@ -44,27 +63,14 @@ public class MilvusVectorStoreStrategy extends AbstractVectorStoreStrategy { @Override public void createSchema(String vectorModelName, String kid) { - String url = vectorStoreProperties.getMilvus().getUrl(); String collectionName = vectorStoreProperties.getMilvus().getCollectionname() + kid; - MilvusEmbeddingStore store = MilvusEmbeddingStore.builder() - .uri(url) - .collectionName(collectionName) - .dimension(2048) - .indexType(IndexType.IVF_FLAT) - .metricType(MetricType.L2) - .autoFlushOnInsert(true) - .idFieldName("id") - .textFieldName("text") - .metadataFieldName("metadata") - .vectorFieldName("vector") - .build(); + // 使用缓存获取连接以确保只初始化一次 + EmbeddingStore store = getMilvusStore(collectionName, true); log.info("Milvus集合初始化完成: {}", collectionName); } @Override public void storeEmbeddings(StoreEmbeddingBo storeEmbeddingBo) { - createSchema(storeEmbeddingBo.getVectorModelName(), storeEmbeddingBo.getKid()); - EmbeddingModel embeddingModel = getEmbeddingModel(storeEmbeddingBo.getEmbeddingModelName(), storeEmbeddingBo.getApiKey(), storeEmbeddingBo.getBaseUrl()); @@ -77,57 +83,35 @@ public class MilvusVectorStoreStrategy extends AbstractVectorStoreStrategy { log.info("Milvus向量存储条数记录: {}", chunkList.size()); long startTime = System.currentTimeMillis(); - EmbeddingStore embeddingStore = MilvusEmbeddingStore.builder() - .uri(vectorStoreProperties.getMilvus().getUrl()) - .collectionName(collectionName) - .dimension(2048) - .indexType(IndexType.IVF_FLAT) - .metricType(MetricType.L2) - .autoFlushOnInsert(false) - .idFieldName("id") - .textFieldName("text") - .metadataFieldName("metadata") - .vectorFieldName("vector") - .build(); + // 复用连接，写入场景使用 autoFlush=false 以提升批量插入性能 + EmbeddingStore embeddingStore = getMilvusStore(collectionName, false); IntStream.range(0, chunkList.size()).forEach(i -> { String text = chunkList.get(i); String fid = fidList.get(i); - Embedding embedding = embeddingModel.embed(text).content(); - Metadata metadata = new Metadata() - .put("fid", fid) - .put("kid", kid) - .put("docId", docId); - TextSegment segment = TextSegment.from(text, metadata); - embeddingStore.add(embedding, segment); - }); + Metadata metadata = new Metadata(); + metadata.put("fid", fid); + metadata.put("kid", kid); + metadata.put("docId", docId); + TextSegment textSegment = TextSegment.from(text, metadata); + Embedding embedding = embeddingModel.embed(text).content(); + embeddingStore.add(embedding, textSegment); + }); long endTime = System.currentTimeMillis(); log.info("Milvus向量存储完成消耗时间：{}秒", (endTime - startTime) / 1000); } @Override public List getQueryVector(QueryVectorBo queryVectorBo) { - createSchema(queryVectorBo.getVectorModelName(), queryVectorBo.getKid()); - EmbeddingModel embeddingModel = getEmbeddingModel(queryVectorBo.getEmbeddingModelName(), queryVectorBo.getApiKey(), queryVectorBo.getBaseUrl()); Embedding queryEmbedding = embeddingModel.embed(queryVectorBo.getQuery()).content(); String collectionName = vectorStoreProperties.getMilvus().getCollectionname() + queryVectorBo.getKid(); - EmbeddingStore embeddingStore = MilvusEmbeddingStore.builder() - .uri(vectorStoreProperties.getMilvus().getUrl()) - .collectionName(collectionName) - .dimension(2048) - .indexType(IndexType.IVF_FLAT) - .metricType(MetricType.L2) - .autoFlushOnInsert(true) - .idFieldName("id") - .textFieldName("text") - .metadataFieldName("metadata") - .vectorFieldName("vector") - .build(); + // 查询复用连接，autoFlush 对查询无影响，此处保持 true + EmbeddingStore embeddingStore = getMilvusStore(collectionName, true); List resultList = new ArrayList<>(); EmbeddingSearchRequest request = EmbeddingSearchRequest.builder() @@ -147,40 +131,15 @@ public class MilvusVectorStoreStrategy extends AbstractVectorStoreStrategy { @Override @SneakyThrows public void removeById(String id, String modelName) { - String url = vectorStoreProperties.getMilvus().getUrl(); - String collectionName = vectorStoreProperties.getMilvus().getCollectionname() + id; - MilvusEmbeddingStore store = MilvusEmbeddingStore.builder() - .uri(url) - .collectionName(collectionName) - .dimension(2048) - .indexType(IndexType.IVF_FLAT) - .metricType(MetricType.L2) - .autoFlushOnInsert(true) - .idFieldName("id") - .textFieldName("text") - .metadataFieldName("metadata") - .vectorFieldName("vector") - .build(); - // 修正：MilvusEmbeddingStore 的 dropCollection 需要传入集合名 - store.dropCollection(collectionName); - log.info("Milvus集合删除成功: {}", collectionName); + // 注意：此处原逻辑使用 collectionname + id，保持现状 + EmbeddingStore embeddingStore = getMilvusStore(vectorStoreProperties.getMilvus().getCollectionname() + id, false); + embeddingStore.remove(id); } @Override public void removeByDocId(String docId, String kid) { String collectionName = vectorStoreProperties.getMilvus().getCollectionname() + kid; - EmbeddingStore embeddingStore = MilvusEmbeddingStore.builder() - .uri(vectorStoreProperties.getMilvus().getUrl()) - .collectionName(collectionName) - .dimension(2048) - .indexType(IndexType.IVF_FLAT) - .metricType(MetricType.L2) - .autoFlushOnInsert(false) - .idFieldName("id") - .textFieldName("text") - .metadataFieldName("metadata") - .vectorFieldName("vector") - .build(); + EmbeddingStore embeddingStore = getMilvusStore(collectionName, false); Filter filter = MetadataFilterBuilder.metadataKey("docId").isEqualTo(docId); embeddingStore.removeAll(filter); log.info("Milvus成功删除 docId={} 的所有向量数据", docId); @@ -189,18 +148,7 @@ public class MilvusVectorStoreStrategy extends AbstractVectorStoreStrategy { @Override public void removeByFid(String fid, String kid) { String collectionName = vectorStoreProperties.getMilvus().getCollectionname() + kid; - EmbeddingStore embeddingStore = MilvusEmbeddingStore.builder() - .uri(vectorStoreProperties.getMilvus().getUrl()) - .collectionName(collectionName) - .dimension(2048) - .indexType(IndexType.IVF_FLAT) - .metricType(MetricType.L2) - .autoFlushOnInsert(false) - .idFieldName("id") - .textFieldName("text") - .metadataFieldName("metadata") - .vectorFieldName("vector") - .build(); + EmbeddingStore embeddingStore = getMilvusStore(collectionName, false); Filter filter = MetadataFilterBuilder.metadataKey("fid").isEqualTo(fid); embeddingStore.removeAll(filter); log.info("Milvus成功删除 fid={} 的所有向量数据", fid); diff --git a/ruoyi-modules-api/ruoyi-knowledge-api/src/main/java/org/ruoyi/service/strategy/impl/WeaviateVectorStoreStrategy.java b/ruoyi-modules-api/ruoyi-knowledge-api/src/main/java/org/ruoyi/service/strategy/impl/WeaviateVectorStoreStrategy.java index 1948d2bb..e1975823 100644 --- a/ruoyi-modules-api/ruoyi-knowledge-api/src/main/java/org/ruoyi/service/strategy/impl/WeaviateVectorStoreStrategy.java +++ b/ruoyi-modules-api/ruoyi-knowledge-api/src/main/java/org/ruoyi/service/strategy/impl/WeaviateVectorStoreStrategy.java @@ -27,7 +27,7 @@ import java.util.*; /** * Weaviate向量库策略实现 * - * @author ageer + * @author Yzm */ @Slf4j @Component @@ -84,7 +84,7 @@ public class WeaviateVectorStoreStrategy extends AbstractVectorStoreStrategy { @Override public void storeEmbeddings(StoreEmbeddingBo storeEmbeddingBo) { - createSchema(storeEmbeddingBo.getVectorModelName(), storeEmbeddingBo.getKid()); + createSchema(storeEmbeddingBo.getVectorStoreName(), storeEmbeddingBo.getKid()); EmbeddingModel embeddingModel = getEmbeddingModel(storeEmbeddingBo.getEmbeddingModelName(), storeEmbeddingBo.getApiKey(), storeEmbeddingBo.getBaseUrl()); List chunkList = storeEmbeddingBo.getChunkList(); diff --git a/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/chat/service/knowledge/KnowledgeInfoServiceImpl.java b/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/chat/service/knowledge/KnowledgeInfoServiceImpl.java index d4cea2f2..e34ababf 100644 --- a/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/chat/service/knowledge/KnowledgeInfoServiceImpl.java +++ b/ruoyi-modules/ruoyi-chat/src/main/java/org/ruoyi/chat/service/knowledge/KnowledgeInfoServiceImpl.java @@ -319,8 +319,7 @@ public class KnowledgeInfoServiceImpl implements IKnowledgeInfoService { storeEmbeddingBo.setDocId(docId); storeEmbeddingBo.setFids(fids); storeEmbeddingBo.setChunkList(chunkList); - storeEmbeddingBo.setVectorModelName(knowledgeInfoVo.getVectorModelName()); - storeEmbeddingBo.setEmbeddingModelId(knowledgeInfoVo.getEmbeddingModelId()); + storeEmbeddingBo.setVectorStoreName(knowledgeInfoVo.getVectorModelName()); storeEmbeddingBo.setEmbeddingModelName(knowledgeInfoVo.getEmbeddingModelName()); storeEmbeddingBo.setApiKey(chatModelVo.getApiKey()); storeEmbeddingBo.setBaseUrl(chatModelVo.getApiHost());