添加向量库

qianbao
1 parent d72de0c4
Showing 1 changed file with 226 additions and 41 deletions
src/main/java/com/xly/milvus/service/impl/MilvusServiceImpl.java
@@ -337,30 +337,122 @@ public class MilvusServiceImpl implements MilvusService {
     }
  
     @Override
-    public List<Map<String, Object>> getDataToCollection(String collectionName, String milvusFilter,String searchText,Integer size,List<String> fields,String vectorValue,String sceneName){
-        log.info("开始相似度查询: collection={}, searchText={}", collectionName, searchText);
-        // 2. 设置范围搜索参数
-        Map<String, Object> searchParams = new HashMap<>();
-        searchParams.put("nprobe", 10);
-        // 对于 IP 度量，相似度范围在 [minScore, maxScore]
-        searchParams.put("radius", 0.9);      // 最小相似度
-        searchParams.put("range_filter", 1); // 最大相似度
-        if(ObjectUtil.isEmpty(fields)){
+    public List<Map<String, Object>> getDataToCollection(String collectionName,
+                                                         String milvusFilter,
+                                                         String searchText,
+                                                         Integer size,
+                                                         List<String> fields,
+                                                         String vectorValue,
+                                                         String sceneName) {
+        log.info("开始查询: collection={}, searchText={}, vectorValue={}", collectionName, searchText, vectorValue);
+
+        // 设置输出字段
+        if (ObjectUtil.isEmpty(fields)) {
             fields = new ArrayList<>();
         }
         fields.add("sSlaveId");
         fields.add("metadata");
-        // 1. 构建查询（通用）
+
+        // 根据 vectorValue 判断走哪个查询
+        if (ObjectUtil.isNotEmpty(vectorValue)) {
+            // vectorValue 不为空：走向量相似度搜索（使用 SearchReq）
+            log.info("使用向量相似度搜索，vectorValue 存在");
+            return similaritySearchWithVector(collectionName, milvusFilter, size, fields, vectorValue);
+//        } else if (ObjectUtil.isNotEmpty(searchText)) {
+//            // vectorValue 为空但 searchText 不为空：走文本向量化搜索（使用 SearchReq）
+//            log.info("使用文本向量化搜索，searchText={}", searchText);
+//            return similaritySearchWithText(collectionName, milvusFilter, searchText, size, fields);
+        } else {
+            // vectorValue 和 searchText 都为空：走全量数据查询（使用 QueryReq）
+            log.info("使用全量数据查询，无搜索条件");
+            return queryAllData(collectionName, milvusFilter, fields,size);
+        }
+    }
+
+    /**
+     * 全量数据查询（使用 Query API）
+     */
+    private List<Map<String, Object>> queryAllData(String collectionName,
+                                                   String milvusFilter,
+                                                   List<String> fields,
+                                                   Integer size) {
+        // 使用 QueryReqBuilder（不需要向量数据）
+        QueryReq.QueryReqBuilder builder = QueryReq.builder()
+                .collectionName(collectionName)
+                .outputFields(fields);
+        if (ObjectUtil.isNotEmpty(milvusFilter)) {
+            builder.filter(milvusFilter);
+        }
+        // 可选：如果 size 有值，可以设置 limit
+        if (ObjectUtil.isNotEmpty(size)) {
+            builder.limit(size);
+        }
+        QueryReq queryReq = builder.build();
+        QueryResp queryResp = milvusClient.query(queryReq);
+        return processMilvusResults(queryResp);
+    }
+
+    /**
+     * 使用文本向量化进行相似度搜索
+     */
+    private List<Map<String, Object>> similaritySearchWithText(String collectionName,
+                                                               String milvusFilter,
+                                                               String searchText,
+                                                               Integer size,
+                                                               List<String> fields) {
+        // 向量化搜索文本
+        List<Float> vectorList = vectorizationService.textToVector(searchText);
+        if (vectorList == null || vectorList.isEmpty()) {
+            throw new RuntimeException("向量化失败");
+        }
+
+        // 转换为 float[]
+        float[] floatArray = new float[vectorList.size()];
+        for (int i = 0; i < vectorList.size(); i++) {
+            floatArray[i] = vectorList.get(i);
+        }
+
+        // 设置搜索参数
+        Map<String, Object> searchParams = new HashMap<>();
+        searchParams.put("nprobe", 10);
+        searchParams.put("radius", 0.9);
+        searchParams.put("range_filter", 1);
+
+        // 构建搜索请求
+        FloatVec floatVec = new FloatVec(floatArray);
         SearchReq.SearchReqBuilder builder = SearchReq.builder()
                 .collectionName(collectionName)
                 .topK(size)
                 .metricType(IndexParam.MetricType.IP)
                 .outputFields(fields)
-//                .filterType(FilterType.POST_FILTER)
-                .searchParams(searchParams);
-        if(ObjectUtil.isEmpty(vectorValue)){
-            vectorValue = sceneName;
+                .searchParams(searchParams)
+                .data(Collections.singletonList(floatVec))
+                .annsField("vector");
+
+        if (ObjectUtil.isNotEmpty(milvusFilter)) {
+            builder.filter(milvusFilter);
         }
+
+        SearchReq searchReq = builder.build();
+        SearchResp searchResp = milvusClient.search(searchReq);
+
+        return processMilvusResults(searchResp);
+    }
+
+    /**
+     * 使用已有向量进行相似度搜索
+     */
+    private List<Map<String, Object>> similaritySearchWithVector(String collectionName,
+                                                                 String milvusFilter,
+                                                                 Integer size,
+                                                                 List<String> fields,
+                                                                 String vectorValue) {
+        // 设置搜索参数
+        Map<String, Object> searchParams = new HashMap<>();
+        searchParams.put("nprobe", 10);
+        searchParams.put("radius", 0.9);
+        searchParams.put("range_filter", 1);
+        // 解析向量字符串为 float 数组
         // 1. 向量化搜索文本
         List<Float> vectorList = vectorizationService.textToVector(vectorValue);
         if (vectorList == null || vectorList.isEmpty()) {
@@ -373,31 +465,72 @@ public class MilvusServiceImpl implements MilvusService {
         }
         // 3. 创建 Milvus FloatVec 对象
         FloatVec floatVec = new FloatVec(floatArray);
-        builder.data(Collections.singletonList(floatVec))
-                .annsField("vector");  // 向量字段名
+        SearchReq.SearchReqBuilder builder = SearchReq.builder()
+                .collectionName(collectionName)
+                .topK(size)
+                .metricType(IndexParam.MetricType.IP)
+                .outputFields(fields)
+                .searchParams(searchParams)
+                .data(Collections.singletonList(floatVec))
+                .annsField("vector");
  
-        if(ObjectUtil.isNotEmpty(milvusFilter)){
+        if (ObjectUtil.isNotEmpty(milvusFilter)) {
             builder.filter(milvusFilter);
         }
-        // 4. 构建搜索请求
-//        SearchReq searchReq = SearchReq.builder()
-//                .collectionName(collectionName)
-//                .data(Collections.singletonList(floatVec))
-//                .annsField("vector")  // 向量字段名
-//                .topK(size)  // 返回最相似的10条
-//                .metricType(IndexParam.MetricType.IP)  // 内积相似度
-//                .outputFields(fields)
-//                .searchParams(searchParams)
-//                .filter(milvusFilter)
-//                .build();
-        // 5. 执行搜索
+
         SearchReq searchReq = builder.build();
         SearchResp searchResp = milvusClient.search(searchReq);
-
-        // 6. 处理结果
         return processMilvusResults(searchResp);
     }
  
+//    public List<Map<String, Object>> getDataToCollection(String collectionName, String milvusFilter,String searchText,Integer size,List<String> fields,String vectorValue,String sceneName){
+//        log.info("开始相似度查询: collection={}, searchText={}", collectionName, searchText);
+//        // 2. 设置范围搜索参数
+//        Map<String, Object> searchParams = new HashMap<>();
+//        searchParams.put("nprobe", 10);
+//        // 对于 IP 度量，相似度范围在 [minScore, maxScore]
+//        searchParams.put("radius", 0.9);      // 最小相似度
+//        searchParams.put("range_filter", 1); // 最大相似度
+//        if(ObjectUtil.isEmpty(fields)){
+//            fields = new ArrayList<>();
+//        }
+//        fields.add("sSlaveId");
+//        fields.add("metadata");
+//        // 1. 构建查询（通用）
+//        SearchReq.SearchReqBuilder builder = SearchReq.builder()
+//                .collectionName(collectionName)
+//                .topK(size)
+//                .metricType(IndexParam.MetricType.IP)
+//                .outputFields(fields)
+////                .filterType(FilterType.POST_FILTER)
+//                .searchParams(searchParams);
+//        if(ObjectUtil.isNotEmpty(vectorValue)){
+//            // 1. 向量化搜索文本
+//            List<Float> vectorList = vectorizationService.textToVector(vectorValue);
+//            if (vectorList == null || vectorList.isEmpty()) {
+//                throw new RuntimeException("向量化失败");
+//            }
+//            // 2. 转换为 float[]
+//            float[] floatArray = new float[vectorList.size()];
+//            for (int i = 0; i < vectorList.size(); i++) {
+//                floatArray[i] = vectorList.get(i);
+//            }
+//            // 3. 创建 Milvus FloatVec 对象
+//            FloatVec floatVec = new FloatVec(floatArray);
+//            builder.data(Collections.singletonList(floatVec))
+//                    .annsField("vector");  // 向量字段名
+//        }
+//        if(ObjectUtil.isNotEmpty(milvusFilter)){
+//            builder.filter(milvusFilter);
+//        }
+//        // 5. 执行搜索
+//        SearchReq searchReq = builder.build();
+//        SearchResp searchResp = milvusClient.search(searchReq);
+//
+//        // 6. 处理结果
+//        return processMilvusResults(searchResp);
+//    }
+
  
     /**
      * 判断 Milvus 过滤条件是否有效（支持 TEXT_MATCH 全文检索）
@@ -559,7 +692,7 @@ public class MilvusServiceImpl implements MilvusService {
     }
  
     /**
-     * 处理 Milvus 查询结果
+     * 处理 SearchResp 查询结果（向量相似度搜索）
      */
     private List<Map<String, Object>> processMilvusResults(SearchResp response) {
         List<Map<String, Object>> results = new ArrayList<>();
@@ -572,32 +705,84 @@ public class MilvusServiceImpl implements MilvusService {
             log.warn("Milvus 搜索结果为空");
             return results;
         }
+
         // 遍历每个查询的结果集（通常只有一个查询）
         for (List<SearchResp.SearchResult> resultList : searchResults) {
             // 遍历每个搜索结果
             for (SearchResp.SearchResult result : resultList) {
+                Map<String, Object> resultMap = new HashMap<>();
+
                 // 获取实体字段数据
                 Map<String, Object> entity = result.getEntity();
-                Map<String,Object> metadata = new HashMap<>();
-                if(ObjectUtil.isNotEmpty(entity.get("metadata"))){
-                    JsonObject obj = (JsonObject) entity.get("metadata");
-                    metadata.putAll( jsonObjectToMap(obj));
+                if (entity != null && !entity.isEmpty()) {
+                    resultMap.putAll(entity);
+                }
+
+                // 处理 metadata 字段（如果是 JSON 对象）
+                if (resultMap.containsKey("metadata") && resultMap.get("metadata") instanceof JsonObject) {
+                    JsonObject jsonObject = (JsonObject) resultMap.get("metadata");
+                    Map<String, Object> metadataMap = jsonObjectToMap(jsonObject);
+                    // 可以选择合并或替换
+                    resultMap.putAll(metadataMap);
+                    resultMap.remove("metadata"); // 移除原始的 metadata 对象
                 }
+
                 // 获取相似度分数
                 Float score = result.getScore();
                 if (score != null) {
-                    metadata.put("score", score);
+                    resultMap.put("score", score);
                 }
-                // 将所有字段添加到结果中
-//                item.putAll(entity);
-                results.add(metadata);
+
+                results.add(resultMap);
             }
         }
+
         log.info("处理完成，共 {} 条搜索结果", results.size());
         return results;
     }
  
     /**
+     * 处理 QueryResp 查询结果（全量查询）
+     */
+    private List<Map<String, Object>> processMilvusResults(QueryResp response) {
+        List<Map<String, Object>> results = new ArrayList<>();
+        if (response == null) {
+            log.warn("Milvus 响应为空");
+            return results;
+        }
+
+        List<QueryResp.QueryResult> queryResults = response.getQueryResults();
+        if (queryResults == null || queryResults.isEmpty()) {
+            log.warn("Milvus 查询结果为空");
+            return results;
+        }
+
+        // 遍历每个查询结果
+        for (QueryResp.QueryResult queryResult : queryResults) {
+            Map<String, Object> resultMap = new HashMap<>();
+
+            // 获取实体字段数据
+            Map<String, Object> entity = queryResult.getEntity();
+            if (entity != null && !entity.isEmpty()) {
+                resultMap.putAll(entity);
+            }
+
+            // 处理 metadata 字段（如果是 JSON 对象）
+            if (resultMap.containsKey("metadata") && resultMap.get("metadata") instanceof JsonObject) {
+                JsonObject jsonObject = (JsonObject) resultMap.get("metadata");
+                Map<String, Object> metadataMap = jsonObjectToMap(jsonObject);
+                resultMap.putAll(metadataMap);
+                resultMap.remove("metadata");
+            }
+
+            results.add(resultMap);
+        }
+
+        log.info("处理完成，共 {} 条查询结果", results.size());
+        return results;
+    }
+
+    /**
      * JsonObject 转 Map<String, Object>
      */
     public static Map<String, Object> jsonObjectToMap(JsonObject jsonObject) {
@@ -693,7 +878,7 @@ public class MilvusServiceImpl implements MilvusService {
             }
             vectorText.append(" | ").append(fieldArr[0]).append(":").append(sText);
         }
-        vectorText.append(" ").append(sceneName);
+//        vectorText.append(" ").append(sceneName).append("数据");
     }