Procházet zdrojové kódy

Merge remote-tracking branch 'origin/master'

chenkq před 1 měsícem
rodič
revize
00ae6bae0b

+ 82 - 43
java/storlead-knowledge/storlead-knowledge-api/src/main/java/com/storlead/knowledge/api/ChunkController.java

@@ -101,12 +101,12 @@ public class ChunkController {
 
     @GetMapping("getAllChunks")
     @ApiOperation("获取所有的文本块")
-    public Result<Object> getAllChunks(QueryPageDTO page) {
+    public Result<Object> getAllChunks(String dataset_id,QueryPageDTO page) {
         // 1. 调用Dify API获取所有知识库列表(使用较大分页以获取全部数据)
         String url = difyProperties.getBaseUrl() + "datasets";
         Result<Object> datasetsResult = httpService.get(
                 url,
-                Map.of("page", 1, "limit", 1000),
+                Map.of("page", 1, "limit", 100),
                 "Bearer " + difyProperties.getDatasetApiKey(),
                 new TypeReference<>() {}
         );
@@ -122,48 +122,10 @@ public class ChunkController {
                     for (Object item : dataList) {
                         if (item instanceof Map<?, ?> datasetMap) {
                             String datasets_id = datasetMap.get("id").toString();
-                            //获取知识库中文档的列表
-                            String url1 = difyProperties.getBaseUrl() + "datasets/"+datasets_id+"/documents";
-                            Result<Object> documents_list  = httpService.get(
-                                    url1,
-                                    Map.of("page", 1, "limit", 10000),
-                                    "Bearer " + difyProperties.getDatasetApiKey(),
-                                    new TypeReference<>() {}
-                            );
-
-                            if (documents_list.isSuccess()) {
-                                Object responseData1 = documents_list.getResult();
-                                Map<?, ?> responseMap2 = (Map<?, ?>) responseData1;
-                                Object dataObj1 = responseMap2.get("data");
-                                if (dataObj1 instanceof List<?> dataList1) {
-                                    for (Object item1 : dataList1) {
-                                        if (item1 instanceof Map<?, ?> datasetMap1) {
-                                            String document_id = datasetMap1.get("id").toString();
-                                            String url2 = difyProperties.getBaseUrl() + "datasets/" + datasets_id + "/documents/" + document_id+"/segments";
-                                            Result<Object> chunks_result = httpService.get(
-                                                    url2,
-                                                    page.getKeyword()==null?Map.of("page", 1, "limit", 10000):Map.of("page", 1, "limit", 10000,"keyword", page.getKeyword()),
-                                                    "Bearer " + difyProperties.getDatasetApiKey(),
-                                                    new TypeReference<>() {
-                                                    });
-                                            if (chunks_result.isSuccess()) {
-                                                Object responseData2 = chunks_result.getResult();
-                                                Map<?, ?> responseMap3 = (Map<?, ?>) responseData2;
-                                                Object dataObj2 = responseMap3.get("data");
-                                                if (dataObj2 instanceof List<?> dataList2) {
-                                                    //results.addAll(dataList2);
-                                                    for (Object item2 : dataList2) {
-                                                        @SuppressWarnings("unchecked")
-                                                        Map<String, Object> datasetMap2 = (Map<String, Object>) item2;
-                                                        datasetMap2.put("dataset_id", datasets_id);
-                                                        results.add(datasetMap2);
-                                                    }
-                                                }
-                                            }
-                                        }
-                                    }
-                                }
+                            if (dataset_id != null && !datasets_id.equals(dataset_id)) {
+                                continue;
                             }
+                            results.addAll(getDataSetsChunks(datasets_id, page.getKeyword()));
                         }
                     }
                 }
@@ -191,4 +153,81 @@ public class ChunkController {
         resultMap.put("has_more", total > (fromIndex + pageSize));
         return Result.ok(resultMap);
     }
+
+
+    /**
+     * 获取单个知识库的所有的文本块
+     * @param dataset_id
+     * @return
+     */
+    public List<Object> getDataSetsChunks(String dataset_id,String keyword){
+        List<Object> results =  new ArrayList<>();
+        //获取知识库中文档的列表
+        String url1 = difyProperties.getBaseUrl() + "datasets/"+dataset_id+"/documents";
+        int page0 = 1;
+        while(true){
+             Result<Object> documents_list  = httpService.get(
+                url1,
+                Map.of("page", 1, "limit", 100),
+                "Bearer " + difyProperties.getDatasetApiKey(),
+                new TypeReference<>() {});
+
+            if (documents_list.isSuccess()) {
+                Object responseData1 = documents_list.getResult();
+                Map<?, ?> responseMap2 = (Map<?, ?>) responseData1;
+                Object dataObj1 = responseMap2.get("data");
+                //分页上线为100
+                Object has_more = responseMap2.get("has_more");
+
+
+                if (dataObj1 instanceof List<?> dataList1) {
+                    for (Object item1 : dataList1) {
+                        if (item1 instanceof Map<?, ?> datasetMap1) {
+                            String document_id = datasetMap1.get("id").toString();
+                            String url2 = difyProperties.getBaseUrl() + "datasets/" + dataset_id + "/documents/" + document_id+"/segments";
+                            int page = 1 ;
+                            while (true) {
+                                Result<Object> chunks_result = httpService.get(
+                                        url2,
+                                        keyword==null?Map.of("page", page, "limit", 100):
+                                                Map.of("page", page, "limit", 100,"keyword",keyword),
+                                        "Bearer " + difyProperties.getDatasetApiKey(),
+                                        new TypeReference<>() {
+                                        });
+                                if (chunks_result.isSuccess()) {
+                                    Object responseData2 = chunks_result.getResult();
+                                    Map<?, ?> responseMap3 = (Map<?, ?>) responseData2;
+                                    Object dataObj2 = responseMap3.get("data");
+                                    if (dataObj2 instanceof List<?> dataList2) {
+                                        //results.addAll(dataList2);
+                                        for (Object item2 : dataList2) {
+                                            @SuppressWarnings("unchecked")
+                                            Map<String, Object> datasetMap2 = (Map<String, Object>) item2;
+                                            datasetMap2.put("dataset_id", dataset_id);
+                                            results.add(datasetMap2);
+                                        }
+                                    }
+                                    //分页上线为100
+                                    Object has_more1 = responseMap3.get("has_more");
+                                    if (has_more1.toString().equals("false")) {
+                                        break;
+                                    }else{
+                                        page +=1 ;
+                                    }
+                                }
+                            }
+
+                        }
+                    }
+                }
+                if (has_more.toString().equals("false")) {
+                    break;
+                }else{
+                    page0 +=1 ;
+                }
+            }
+        }
+        return  results;
+    }
+
 }