Bladeren bron

添加查询所有文档的接口

183207892172 1 maand geleden
bovenliggende
commit
219adeaafa

+ 106 - 0
java/storlead-knowledge/storlead-knowledge-api/src/main/java/com/storlead/knowledge/api/DocumentController.java

@@ -16,6 +16,9 @@ import org.springframework.web.bind.annotation.*;
 import org.springframework.web.multipart.MultipartFile;
 
 import javax.annotation.Resource;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 
 
@@ -142,4 +145,107 @@ public class DocumentController {
         );
     }
 
+
+
+    @GetMapping("getAllDocumnets")
+    @ApiOperation("获取所有的文档")
+    public Result<Object> getAllDocumnets(String dataset_id,QueryPageDTO page) {
+        // 1. 调用Dify API获取所有知识库列表(使用较大分页以获取全部数据)
+        String url = difyProperties.getBaseUrl() + "datasets";
+        Result<Object> datasetsResult = httpService.get(
+                url,
+                Map.of("page", 1, "limit", 100),
+                "Bearer " + difyProperties.getDatasetApiKey(),
+                new TypeReference<>() {}
+        );
+
+        List<Object> results =  new ArrayList<>();
+        if (datasetsResult.isSuccess()) {
+            Object responseData = datasetsResult.getResult();
+            if (responseData instanceof Map) {
+                Map<?, ?> responseMap = (Map<?, ?>) responseData;
+                // 遍历数据集,累加文档数和文本块数
+                Object dataObj = responseMap.get("data");
+                if (dataObj instanceof List) {
+                    List<?> dataList = (List<?>) dataObj;
+                    for (Object item : dataList) {
+                        if (item instanceof Map ) {
+                            Map<?, ?> datasetMap = (Map<?, ?>) item;
+                            String datasets_id = datasetMap.get("id").toString();
+                            if (dataset_id != null && !datasets_id.equals(dataset_id)) {
+                                continue;
+                            }
+                            results.addAll(getDataSetsDocuments(datasets_id, page.getKeyword()));
+                        }
+                    }
+                }
+            }
+        }
+
+        // 对结果进行分页
+        int total = results.size();
+        int pageIndex = page.getPageIndex() != null ? page.getPageIndex() : 1;
+        int pageSize = page.getPageSize() != null ? page.getPageSize() : 10;
+        int fromIndex = (pageIndex - 1) * pageSize;
+        int toIndex = Math.min(fromIndex + pageSize, total);
+
+        List<Object> pagedResults;
+        if (fromIndex >= total) {
+            pagedResults = new ArrayList<>();
+        } else {
+            pagedResults = results.subList(fromIndex, toIndex);
+        }
+        Map<String,Object> resultMap = new HashMap<>() ;
+        resultMap.put("data", pagedResults);
+        resultMap.put("total", total);
+        resultMap.put("limit", pageSize);
+        resultMap.put("page", pageIndex);
+        resultMap.put("has_more", total > (fromIndex + pageSize));
+        return Result.ok(resultMap);
+    }
+
+
+    /**
+     * 获取单个知识库的所有的文本块
+     * @param dataset_id
+     * @return
+     */
+    public List<Object> getDataSetsDocuments(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,
+                    keyword==null?Map.of("page", 1, "limit", 100):Map.of("page", 1, "limit", 100,"keyword",keyword),
+                    "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) {
+                    List<?> dataList1 = (List<?>) dataObj1;
+                    for (Object item2 : dataList1) {
+                        @SuppressWarnings("unchecked")
+                        Map<String, Object> datasetMap2 = (Map<String, Object>) item2;
+                        datasetMap2.put("dataset_id", dataset_id);
+                        results.add(datasetMap2);
+                    }
+                }
+                if (has_more.toString().equals("false")) {
+                    break;
+                }else{
+                    page0 +=1 ;
+                }
+            }
+        }
+        return  results;
+    }
+
 }