查询、删除向量库文件

llm-back/ruoyi-agent/src/main/java/com/ruoyi/agent/controller/KnowLedgeController.java

@@ -13,6 +13,7 @@ import org.springframework.web.multipart.MultipartFile;
 
 import java.io.File;
 import java.io.IOException;
+import java.util.List;
 
 /**
  * cmc知识库Controller
@@ -44,7 +45,6 @@ public class KnowLedgeController extends BaseController
     public AjaxResult createKnowLedgeCollection(String collectionName, String description)
     {
         milvusService.createCollection(collectionName, description, 512);
-        milvusService.releaseCollectionName(collectionName);
         return success();
     }
 
@@ -55,7 +55,6 @@ public class KnowLedgeController extends BaseController
     public AjaxResult modifyKnowLedgeCollection(String collectionName, String newCollectionName)
     {
         milvusService.collectionRename(collectionName, newCollectionName);
-        milvusService.releaseCollectionName(collectionName);
         return success();
     }
 
@@ -69,10 +68,20 @@ public class KnowLedgeController extends BaseController
         return success();
     }
 
+    /**
+     * 查询知识库文件
+     */
+    @GetMapping("/listDocument")
+    public AjaxResult listKnowledgeDocument(String collectionName, String fileType)
+    {
+        List<String> documentList = milvusService.listDocument(collectionName, fileType);
+        return success(documentList);
+    }
+
     /**
      * 插入知识库文件
      */
-    @PostMapping("/insert")
+    @PostMapping("/insertDocument")
     public AjaxResult insertKnowledgeDocument(MultipartFile file, String collectionName) throws IOException
     {
         LangChainMilvusService langChainMilvusService = new LangChainMilvusService(
@@ -89,7 +98,16 @@ public class KnowLedgeController extends BaseController
             file.transferTo(transferFile);
         }
         langChainMilvusService.insertLangchainEmbeddingDocument(transferFile);
-        langChainMilvusService.releaseCollectionName(collectionName);
+        return success();
+    }
+
+    /**
+     * 删除知识库文件
+     */
+    @DeleteMapping("/removeDocument")
+    public AjaxResult deleteKnowledgeDocument(String fileName, String collectionName)
+    {
+        milvusService.removeDocument(collectionName, fileName);
         return success();
     }
 

llm-back/ruoyi-agent/src/main/java/com/ruoyi/agent/controller/RagController.java

@@ -37,9 +37,10 @@ public class RagController extends BaseController
 
         // 1. Milvus检索
         List<String> contexts = langChainMilvusService.retrieveFromMilvus(question, 1);
-
+        String result = langChainMilvusService.generateAnswerWithRag(question, contexts);
+        langChainMilvusService.releaseCollectionName();
         // 2. 调用本地LLM或HTTP服务
-        return success(langChainMilvusService.generateAnswerWithRag(question, contexts));
+        return success(result);
     }
 
 }

llm-back/ruoyi-agent/src/main/java/com/ruoyi/agent/service/LangChainMilvusService.java

@@ -16,8 +16,12 @@ import io.milvus.param.R;
 import io.milvus.param.RpcStatus;
 import io.milvus.param.collection.LoadCollectionParam;
 import io.milvus.param.collection.ReleaseCollectionParam;
+import io.milvus.param.dml.DeleteParam;
 import io.milvus.param.dml.InsertParam;
 import io.milvus.param.dml.SearchParam;
+import io.milvus.param.highlevel.dml.QuerySimpleParam;
+import io.milvus.param.highlevel.dml.response.QueryResponse;
+import io.milvus.response.QueryResultsWrapper;
 import io.milvus.response.SearchResultsWrapper;
 import okhttp3.*;
 
@@ -50,6 +54,9 @@ public class LangChainMilvusService {
         this.embeddingModel = embeddingModel;
     }
 
+    /**
+     * 插入知识库文件
+     */
     public void insertLangchainEmbeddingDocument(File file) throws IOException {
         // 加载文档
         InputStream fileInputStream = new FileInputStream(file);
@@ -59,6 +66,7 @@ public class LangChainMilvusService {
 
         // 提取文本和生成嵌入
         List<String> fileNames = new ArrayList<>();
+        List<String> fileTypes = new ArrayList<>();
         List<String> texts = new ArrayList<>();
         List<List<Float>> embeddings = new ArrayList<>();
 
@@ -67,6 +75,8 @@ public class LangChainMilvusService {
             if (text.trim().isEmpty())
                 continue;
             fileNames.add(file.getName());
+            String[] fileName = file.getName().split("\\.");
+            fileTypes.add(fileName[fileName.length - 1]);
             texts.add(text);
             embeddings.add(embeddingModel.embed(text).content().vectorAsList());
         }
@@ -74,6 +84,7 @@ public class LangChainMilvusService {
         // 准备插入数据
         List<InsertParam.Field> fields = new ArrayList<>();
         fields.add(new InsertParam.Field("file_name", fileNames));
+        fields.add(new InsertParam.Field("file_type", fileTypes));
         fields.add(new InsertParam.Field("content", texts));
         fields.add(new InsertParam.Field("embedding", embeddings));
 
@@ -86,7 +97,7 @@ public class LangChainMilvusService {
         R<MutationResult> insertResult = milvusClient.insert(insertParam);
 
         if (insertResult.getStatus() != R.Status.Success.getCode()) {
-            throw new RuntimeException("Failed to insert document: " + insertResult.getMessage());
+            throw new RuntimeException("插入文件失败: " + insertResult.getMessage());
         }
     }
 
@@ -95,23 +106,14 @@ public class LangChainMilvusService {
         List<List<Float>> queryVector = Arrays.asList(embeddingModel.embed(query).content().vectorAsList());
 
         //  加载集合
-        LoadCollectionParam loadParam = LoadCollectionParam.newBuilder()
-                .withCollectionName(collectionName)
-                .build();
-
-        R<RpcStatus> loadResponse = milvusClient.loadCollection(loadParam);
-        if (loadResponse.getStatus() != R.Status.Success.getCode()) {
-            System.err.println("Failed to load collection: " + loadResponse.getMessage());
-            milvusClient.close();
-            return null;
-        }
+        loadCollectionName();
 
         // 构建SearchParam
         SearchParam searchParam = SearchParam.newBuilder()
                 .withCollectionName(collectionName)
                 .withVectors(queryVector)
                 .withTopK(topK)
-                .withOutFields(Arrays.asList("content"))
+                .withOutFields(Arrays.asList("file_name", "file_type", "content"))
                 .withVectorFieldName("embedding")
                 .withMetricType(MetricType.COSINE)
                 .withParams("{\"nprobe\": 1}")
@@ -160,10 +162,25 @@ public class LangChainMilvusService {
         return sb.toString();
     }
 
+    /**
+     * 加载知识库Collection
+     */
+    public void loadCollectionName() {
+        LoadCollectionParam loadParam = LoadCollectionParam.newBuilder()
+                .withCollectionName(collectionName)
+                .build();
+
+        R<RpcStatus> loadResponse = milvusClient.loadCollection(loadParam);
+        if (loadResponse.getStatus() != R.Status.Success.getCode()) {
+            System.err.println("加载Collection失败: " + loadResponse.getMessage());
+            milvusClient.close();
+        }
+    }
+
     /**
      * 释放知识库Collection
      */
-    public void releaseCollectionName(String collectionName) {
+    public void releaseCollectionName() {
         ReleaseCollectionParam param = ReleaseCollectionParam.newBuilder()
                 .withCollectionName(collectionName)
                 .build();

llm-back/ruoyi-agent/src/main/java/com/ruoyi/agent/service/MilvusService.java

@@ -8,13 +8,20 @@ import io.milvus.param.ConnectParam;
 import io.milvus.param.IndexType;
 import io.milvus.param.MetricType;
 import io.milvus.param.collection.*;
+import io.milvus.param.dml.DeleteParam;
 import io.milvus.param.highlevel.collection.ListCollectionsParam;
 import io.milvus.param.highlevel.collection.response.ListCollectionsResponse;
+import io.milvus.param.highlevel.dml.QuerySimpleParam;
+import io.milvus.param.highlevel.dml.response.QueryResponse;
 import io.milvus.param.index.CreateIndexParam;
+import io.milvus.response.QueryResultsWrapper;
 
 import java.text.SimpleDateFormat;
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.List;
 import java.util.TimeZone;
+import java.util.stream.Collectors;
 
 public class MilvusService {
     private MilvusServiceClient milvusClient;
@@ -48,6 +55,12 @@ public class MilvusService {
                 .withMaxLength(256)
                 .build();
 
+        FieldType fieldTypeField = FieldType.newBuilder()
+                .withName("file_type")
+                .withDataType(DataType.VarChar)
+                .withMaxLength(256)
+                .build();
+
         FieldType contentField = FieldType.newBuilder()
                 .withName("content")
                 .withDataType(DataType.VarChar)
@@ -65,6 +78,7 @@ public class MilvusService {
                 .withDescription(description)
                 .addFieldType(idField)
                 .addFieldType(fileNameField)
+                .addFieldType(fieldTypeField)
                 .addFieldType(contentField)
                 .addFieldType(vectorField)
                 .build();
@@ -92,12 +106,14 @@ public class MilvusService {
         ListCollectionsResponse listResponse = milvusClient.listCollections(listParam).getData();
         if (listResponse != null) {
             List<String> collectionNames = milvusClient.listCollections(listParam).getData().collectionNames;
-            for (int i = 0; i < collectionNames.size(); i++) {
+            for (String collectionName : collectionNames) {
                 JSONObject jsonObject = new JSONObject();
-                DescribeCollectionParam describeParam = DescribeCollectionParam.newBuilder().withCollectionName(collectionNames.get(i)).build();
+                DescribeCollectionParam describeParam = DescribeCollectionParam.newBuilder()
+                                                        .withCollectionName(collectionName)
+                                                        .build();
                 DescribeCollectionResponse describeResponse = milvusClient.describeCollection(describeParam).getData();
                 jsonObject.put("collectionId", describeResponse.getCollectionID());
-                jsonObject.put("collectionName", collectionNames.get(i));
+                jsonObject.put("collectionName", collectionName);
                 SimpleDateFormat beijingFormat = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
                 beijingFormat.setTimeZone(TimeZone.getTimeZone("Asia/Shanghai"));
                 String beijingTime = beijingFormat.format(describeResponse.getCreatedUtcTimestamp());
@@ -126,10 +142,53 @@ public class MilvusService {
      */
     public void deleteCollectionName(String collectionName) {
         DropCollectionParam param = DropCollectionParam.newBuilder()
-                .withCollectionName(collectionName).build();
+                .withCollectionName(collectionName)
+                .build();
         milvusClient.dropCollection(param);
     }
 
+    /**
+     * 查询知识库文件
+     */
+    public List<String> listDocument(String collectionName, String fileType) {
+        List<String> documentList = new ArrayList<>();
+        loadCollectionName(collectionName);
+        QuerySimpleParam queryParam = QuerySimpleParam.newBuilder()
+                .withCollectionName(collectionName)
+                .withFilter(String.format("file_type == \"%s\"", fileType))
+                .withOutputFields(Arrays.asList("file_type", "file_name"))
+                .withLimit(16384L)
+                .build();
+        QueryResponse queryResults = milvusClient.query(queryParam).getData();
+        List<QueryResultsWrapper.RowRecord> rowRecordList = queryResults.getRowRecords();
+        for (QueryResultsWrapper.RowRecord rowRecord : rowRecordList) {
+            documentList.add(rowRecord.get("file_name").toString());
+        }
+        releaseCollectionName(collectionName);
+        return documentList.stream().distinct().collect(Collectors.toList());
+    }
+
+    /**
+     * 删除知识库文件
+     */
+    public void removeDocument(String collectionName, String fileName) {
+        DeleteParam deleteParam = DeleteParam.newBuilder()
+                .withCollectionName(collectionName)
+                .withExpr(String.format("file_name == \"%s\"", fileName))
+                .build();
+        milvusClient.delete(deleteParam);
+    }
+
+    /**
+     * 加载知识库Collection
+     */
+    public void loadCollectionName(String collectionName) {
+        LoadCollectionParam param = LoadCollectionParam.newBuilder()
+                .withCollectionName(collectionName)
+                .build();
+        milvusClient.loadCollection(param);
+    }
+
     /**
      * 释放知识库Collection
      */