소스 검색

查询、删除向量库文件

lamphua 1 주 전
부모
커밋
b29d75b195

+ 22
- 4
llm-back/ruoyi-agent/src/main/java/com/ruoyi/agent/controller/KnowLedgeController.java 파일 보기

13
 
13
 
14
 import java.io.File;
14
 import java.io.File;
15
 import java.io.IOException;
15
 import java.io.IOException;
16
+import java.util.List;
16
 
17
 
17
 /**
18
 /**
18
  * cmc知识库Controller
19
  * cmc知识库Controller
44
     public AjaxResult createKnowLedgeCollection(String collectionName, String description)
45
     public AjaxResult createKnowLedgeCollection(String collectionName, String description)
45
     {
46
     {
46
         milvusService.createCollection(collectionName, description, 512);
47
         milvusService.createCollection(collectionName, description, 512);
47
-        milvusService.releaseCollectionName(collectionName);
48
         return success();
48
         return success();
49
     }
49
     }
50
 
50
 
55
     public AjaxResult modifyKnowLedgeCollection(String collectionName, String newCollectionName)
55
     public AjaxResult modifyKnowLedgeCollection(String collectionName, String newCollectionName)
56
     {
56
     {
57
         milvusService.collectionRename(collectionName, newCollectionName);
57
         milvusService.collectionRename(collectionName, newCollectionName);
58
-        milvusService.releaseCollectionName(collectionName);
59
         return success();
58
         return success();
60
     }
59
     }
61
 
60
 
69
         return success();
68
         return success();
70
     }
69
     }
71
 
70
 
71
+    /**
72
+     * 查询知识库文件
73
+     */
74
+    @GetMapping("/listDocument")
75
+    public AjaxResult listKnowledgeDocument(String collectionName, String fileType)
76
+    {
77
+        List<String> documentList = milvusService.listDocument(collectionName, fileType);
78
+        return success(documentList);
79
+    }
80
+
72
     /**
81
     /**
73
      * 插入知识库文件
82
      * 插入知识库文件
74
      */
83
      */
75
-    @PostMapping("/insert")
84
+    @PostMapping("/insertDocument")
76
     public AjaxResult insertKnowledgeDocument(MultipartFile file, String collectionName) throws IOException
85
     public AjaxResult insertKnowledgeDocument(MultipartFile file, String collectionName) throws IOException
77
     {
86
     {
78
         LangChainMilvusService langChainMilvusService = new LangChainMilvusService(
87
         LangChainMilvusService langChainMilvusService = new LangChainMilvusService(
89
             file.transferTo(transferFile);
98
             file.transferTo(transferFile);
90
         }
99
         }
91
         langChainMilvusService.insertLangchainEmbeddingDocument(transferFile);
100
         langChainMilvusService.insertLangchainEmbeddingDocument(transferFile);
92
-        langChainMilvusService.releaseCollectionName(collectionName);
101
+        return success();
102
+    }
103
+
104
+    /**
105
+     * 删除知识库文件
106
+     */
107
+    @DeleteMapping("/removeDocument")
108
+    public AjaxResult deleteKnowledgeDocument(String fileName, String collectionName)
109
+    {
110
+        milvusService.removeDocument(collectionName, fileName);
93
         return success();
111
         return success();
94
     }
112
     }
95
 
113
 

+ 3
- 2
llm-back/ruoyi-agent/src/main/java/com/ruoyi/agent/controller/RagController.java 파일 보기

37
 
37
 
38
         // 1. Milvus检索
38
         // 1. Milvus检索
39
         List<String> contexts = langChainMilvusService.retrieveFromMilvus(question, 1);
39
         List<String> contexts = langChainMilvusService.retrieveFromMilvus(question, 1);
40
-
40
+        String result = langChainMilvusService.generateAnswerWithRag(question, contexts);
41
+        langChainMilvusService.releaseCollectionName();
41
         // 2. 调用本地LLM或HTTP服务
42
         // 2. 调用本地LLM或HTTP服务
42
-        return success(langChainMilvusService.generateAnswerWithRag(question, contexts));
43
+        return success(result);
43
     }
44
     }
44
 
45
 
45
 }
46
 }

+ 30
- 13
llm-back/ruoyi-agent/src/main/java/com/ruoyi/agent/service/LangChainMilvusService.java 파일 보기

16
 import io.milvus.param.RpcStatus;
16
 import io.milvus.param.RpcStatus;
17
 import io.milvus.param.collection.LoadCollectionParam;
17
 import io.milvus.param.collection.LoadCollectionParam;
18
 import io.milvus.param.collection.ReleaseCollectionParam;
18
 import io.milvus.param.collection.ReleaseCollectionParam;
19
+import io.milvus.param.dml.DeleteParam;
19
 import io.milvus.param.dml.InsertParam;
20
 import io.milvus.param.dml.InsertParam;
20
 import io.milvus.param.dml.SearchParam;
21
 import io.milvus.param.dml.SearchParam;
22
+import io.milvus.param.highlevel.dml.QuerySimpleParam;
23
+import io.milvus.param.highlevel.dml.response.QueryResponse;
24
+import io.milvus.response.QueryResultsWrapper;
21
 import io.milvus.response.SearchResultsWrapper;
25
 import io.milvus.response.SearchResultsWrapper;
22
 import okhttp3.*;
26
 import okhttp3.*;
23
 
27
 
50
         this.embeddingModel = embeddingModel;
54
         this.embeddingModel = embeddingModel;
51
     }
55
     }
52
 
56
 
57
+    /**
58
+     * 插入知识库文件
59
+     */
53
     public void insertLangchainEmbeddingDocument(File file) throws IOException {
60
     public void insertLangchainEmbeddingDocument(File file) throws IOException {
54
         // 加载文档
61
         // 加载文档
55
         InputStream fileInputStream = new FileInputStream(file);
62
         InputStream fileInputStream = new FileInputStream(file);
59
 
66
 
60
         // 提取文本和生成嵌入
67
         // 提取文本和生成嵌入
61
         List<String> fileNames = new ArrayList<>();
68
         List<String> fileNames = new ArrayList<>();
69
+        List<String> fileTypes = new ArrayList<>();
62
         List<String> texts = new ArrayList<>();
70
         List<String> texts = new ArrayList<>();
63
         List<List<Float>> embeddings = new ArrayList<>();
71
         List<List<Float>> embeddings = new ArrayList<>();
64
 
72
 
67
             if (text.trim().isEmpty())
75
             if (text.trim().isEmpty())
68
                 continue;
76
                 continue;
69
             fileNames.add(file.getName());
77
             fileNames.add(file.getName());
78
+            String[] fileName = file.getName().split("\\.");
79
+            fileTypes.add(fileName[fileName.length - 1]);
70
             texts.add(text);
80
             texts.add(text);
71
             embeddings.add(embeddingModel.embed(text).content().vectorAsList());
81
             embeddings.add(embeddingModel.embed(text).content().vectorAsList());
72
         }
82
         }
74
         // 准备插入数据
84
         // 准备插入数据
75
         List<InsertParam.Field> fields = new ArrayList<>();
85
         List<InsertParam.Field> fields = new ArrayList<>();
76
         fields.add(new InsertParam.Field("file_name", fileNames));
86
         fields.add(new InsertParam.Field("file_name", fileNames));
87
+        fields.add(new InsertParam.Field("file_type", fileTypes));
77
         fields.add(new InsertParam.Field("content", texts));
88
         fields.add(new InsertParam.Field("content", texts));
78
         fields.add(new InsertParam.Field("embedding", embeddings));
89
         fields.add(new InsertParam.Field("embedding", embeddings));
79
 
90
 
86
         R<MutationResult> insertResult = milvusClient.insert(insertParam);
97
         R<MutationResult> insertResult = milvusClient.insert(insertParam);
87
 
98
 
88
         if (insertResult.getStatus() != R.Status.Success.getCode()) {
99
         if (insertResult.getStatus() != R.Status.Success.getCode()) {
89
-            throw new RuntimeException("Failed to insert document: " + insertResult.getMessage());
100
+            throw new RuntimeException("插入文件失败: " + insertResult.getMessage());
90
         }
101
         }
91
     }
102
     }
92
 
103
 
95
         List<List<Float>> queryVector = Arrays.asList(embeddingModel.embed(query).content().vectorAsList());
106
         List<List<Float>> queryVector = Arrays.asList(embeddingModel.embed(query).content().vectorAsList());
96
 
107
 
97
         //  加载集合
108
         //  加载集合
98
-        LoadCollectionParam loadParam = LoadCollectionParam.newBuilder()
99
-                .withCollectionName(collectionName)
100
-                .build();
101
-
102
-        R<RpcStatus> loadResponse = milvusClient.loadCollection(loadParam);
103
-        if (loadResponse.getStatus() != R.Status.Success.getCode()) {
104
-            System.err.println("Failed to load collection: " + loadResponse.getMessage());
105
-            milvusClient.close();
106
-            return null;
107
-        }
109
+        loadCollectionName();
108
 
110
 
109
         // 构建SearchParam
111
         // 构建SearchParam
110
         SearchParam searchParam = SearchParam.newBuilder()
112
         SearchParam searchParam = SearchParam.newBuilder()
111
                 .withCollectionName(collectionName)
113
                 .withCollectionName(collectionName)
112
                 .withVectors(queryVector)
114
                 .withVectors(queryVector)
113
                 .withTopK(topK)
115
                 .withTopK(topK)
114
-                .withOutFields(Arrays.asList("content"))
116
+                .withOutFields(Arrays.asList("file_name", "file_type", "content"))
115
                 .withVectorFieldName("embedding")
117
                 .withVectorFieldName("embedding")
116
                 .withMetricType(MetricType.COSINE)
118
                 .withMetricType(MetricType.COSINE)
117
                 .withParams("{\"nprobe\": 1}")
119
                 .withParams("{\"nprobe\": 1}")
160
         return sb.toString();
162
         return sb.toString();
161
     }
163
     }
162
 
164
 
165
+    /**
166
+     * 加载知识库Collection
167
+     */
168
+    public void loadCollectionName() {
169
+        LoadCollectionParam loadParam = LoadCollectionParam.newBuilder()
170
+                .withCollectionName(collectionName)
171
+                .build();
172
+
173
+        R<RpcStatus> loadResponse = milvusClient.loadCollection(loadParam);
174
+        if (loadResponse.getStatus() != R.Status.Success.getCode()) {
175
+            System.err.println("加载Collection失败: " + loadResponse.getMessage());
176
+            milvusClient.close();
177
+        }
178
+    }
179
+
163
     /**
180
     /**
164
      * 释放知识库Collection
181
      * 释放知识库Collection
165
      */
182
      */
166
-    public void releaseCollectionName(String collectionName) {
183
+    public void releaseCollectionName() {
167
         ReleaseCollectionParam param = ReleaseCollectionParam.newBuilder()
184
         ReleaseCollectionParam param = ReleaseCollectionParam.newBuilder()
168
                 .withCollectionName(collectionName)
185
                 .withCollectionName(collectionName)
169
                 .build();
186
                 .build();

+ 63
- 4
llm-back/ruoyi-agent/src/main/java/com/ruoyi/agent/service/MilvusService.java 파일 보기

8
 import io.milvus.param.IndexType;
8
 import io.milvus.param.IndexType;
9
 import io.milvus.param.MetricType;
9
 import io.milvus.param.MetricType;
10
 import io.milvus.param.collection.*;
10
 import io.milvus.param.collection.*;
11
+import io.milvus.param.dml.DeleteParam;
11
 import io.milvus.param.highlevel.collection.ListCollectionsParam;
12
 import io.milvus.param.highlevel.collection.ListCollectionsParam;
12
 import io.milvus.param.highlevel.collection.response.ListCollectionsResponse;
13
 import io.milvus.param.highlevel.collection.response.ListCollectionsResponse;
14
+import io.milvus.param.highlevel.dml.QuerySimpleParam;
15
+import io.milvus.param.highlevel.dml.response.QueryResponse;
13
 import io.milvus.param.index.CreateIndexParam;
16
 import io.milvus.param.index.CreateIndexParam;
17
+import io.milvus.response.QueryResultsWrapper;
14
 
18
 
15
 import java.text.SimpleDateFormat;
19
 import java.text.SimpleDateFormat;
20
+import java.util.ArrayList;
21
+import java.util.Arrays;
16
 import java.util.List;
22
 import java.util.List;
17
 import java.util.TimeZone;
23
 import java.util.TimeZone;
24
+import java.util.stream.Collectors;
18
 
25
 
19
 public class MilvusService {
26
 public class MilvusService {
20
     private MilvusServiceClient milvusClient;
27
     private MilvusServiceClient milvusClient;
48
                 .withMaxLength(256)
55
                 .withMaxLength(256)
49
                 .build();
56
                 .build();
50
 
57
 
58
+        FieldType fieldTypeField = FieldType.newBuilder()
59
+                .withName("file_type")
60
+                .withDataType(DataType.VarChar)
61
+                .withMaxLength(256)
62
+                .build();
63
+
51
         FieldType contentField = FieldType.newBuilder()
64
         FieldType contentField = FieldType.newBuilder()
52
                 .withName("content")
65
                 .withName("content")
53
                 .withDataType(DataType.VarChar)
66
                 .withDataType(DataType.VarChar)
65
                 .withDescription(description)
78
                 .withDescription(description)
66
                 .addFieldType(idField)
79
                 .addFieldType(idField)
67
                 .addFieldType(fileNameField)
80
                 .addFieldType(fileNameField)
81
+                .addFieldType(fieldTypeField)
68
                 .addFieldType(contentField)
82
                 .addFieldType(contentField)
69
                 .addFieldType(vectorField)
83
                 .addFieldType(vectorField)
70
                 .build();
84
                 .build();
92
         ListCollectionsResponse listResponse = milvusClient.listCollections(listParam).getData();
106
         ListCollectionsResponse listResponse = milvusClient.listCollections(listParam).getData();
93
         if (listResponse != null) {
107
         if (listResponse != null) {
94
             List<String> collectionNames = milvusClient.listCollections(listParam).getData().collectionNames;
108
             List<String> collectionNames = milvusClient.listCollections(listParam).getData().collectionNames;
95
-            for (int i = 0; i < collectionNames.size(); i++) {
109
+            for (String collectionName : collectionNames) {
96
                 JSONObject jsonObject = new JSONObject();
110
                 JSONObject jsonObject = new JSONObject();
97
-                DescribeCollectionParam describeParam = DescribeCollectionParam.newBuilder().withCollectionName(collectionNames.get(i)).build();
111
+                DescribeCollectionParam describeParam = DescribeCollectionParam.newBuilder()
112
+                                                        .withCollectionName(collectionName)
113
+                                                        .build();
98
                 DescribeCollectionResponse describeResponse = milvusClient.describeCollection(describeParam).getData();
114
                 DescribeCollectionResponse describeResponse = milvusClient.describeCollection(describeParam).getData();
99
                 jsonObject.put("collectionId", describeResponse.getCollectionID());
115
                 jsonObject.put("collectionId", describeResponse.getCollectionID());
100
-                jsonObject.put("collectionName", collectionNames.get(i));
116
+                jsonObject.put("collectionName", collectionName);
101
                 SimpleDateFormat beijingFormat = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
117
                 SimpleDateFormat beijingFormat = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
102
                 beijingFormat.setTimeZone(TimeZone.getTimeZone("Asia/Shanghai"));
118
                 beijingFormat.setTimeZone(TimeZone.getTimeZone("Asia/Shanghai"));
103
                 String beijingTime = beijingFormat.format(describeResponse.getCreatedUtcTimestamp());
119
                 String beijingTime = beijingFormat.format(describeResponse.getCreatedUtcTimestamp());
126
      */
142
      */
127
     public void deleteCollectionName(String collectionName) {
143
     public void deleteCollectionName(String collectionName) {
128
         DropCollectionParam param = DropCollectionParam.newBuilder()
144
         DropCollectionParam param = DropCollectionParam.newBuilder()
129
-                .withCollectionName(collectionName).build();
145
+                .withCollectionName(collectionName)
146
+                .build();
130
         milvusClient.dropCollection(param);
147
         milvusClient.dropCollection(param);
131
     }
148
     }
132
 
149
 
150
+    /**
151
+     * 查询知识库文件
152
+     */
153
+    public List<String> listDocument(String collectionName, String fileType) {
154
+        List<String> documentList = new ArrayList<>();
155
+        loadCollectionName(collectionName);
156
+        QuerySimpleParam queryParam = QuerySimpleParam.newBuilder()
157
+                .withCollectionName(collectionName)
158
+                .withFilter(String.format("file_type == \"%s\"", fileType))
159
+                .withOutputFields(Arrays.asList("file_type", "file_name"))
160
+                .withLimit(16384L)
161
+                .build();
162
+        QueryResponse queryResults = milvusClient.query(queryParam).getData();
163
+        List<QueryResultsWrapper.RowRecord> rowRecordList = queryResults.getRowRecords();
164
+        for (QueryResultsWrapper.RowRecord rowRecord : rowRecordList) {
165
+            documentList.add(rowRecord.get("file_name").toString());
166
+        }
167
+        releaseCollectionName(collectionName);
168
+        return documentList.stream().distinct().collect(Collectors.toList());
169
+    }
170
+
171
+    /**
172
+     * 删除知识库文件
173
+     */
174
+    public void removeDocument(String collectionName, String fileName) {
175
+        DeleteParam deleteParam = DeleteParam.newBuilder()
176
+                .withCollectionName(collectionName)
177
+                .withExpr(String.format("file_name == \"%s\"", fileName))
178
+                .build();
179
+        milvusClient.delete(deleteParam);
180
+    }
181
+
182
+    /**
183
+     * 加载知识库Collection
184
+     */
185
+    public void loadCollectionName(String collectionName) {
186
+        LoadCollectionParam param = LoadCollectionParam.newBuilder()
187
+                .withCollectionName(collectionName)
188
+                .build();
189
+        milvusClient.loadCollection(param);
190
+    }
191
+
133
     /**
192
     /**
134
      * 释放知识库Collection
193
      * 释放知识库Collection
135
      */
194
      */

Loading…
취소
저장