dcs 20 小时之前
父节点
当前提交
1e3b126b6f

+ 9 - 0
virgo.api/src/main/java/com/bosshand/virgo/api/workark/controller/DifyController.java

@@ -245,6 +245,15 @@ public class DifyController {
         return Response.ok();
     }
 
+    @ApiOperation("检索知识库")
+    @RequestMapping(value = "/datasets/{datasetId}/retrieve", method = RequestMethod.POST)
+    @ApiImplicitParams({
+            @ApiImplicitParam(name = "datasetId", value = "知识库id")
+    })
+    public Response retrieveDataset(@PathVariable String datasetId, @RequestBody RetrieveDatasetDto retrieveDatasetDto) {
+        return Response.ok(difyDatasetService.retrieveDataset(datasetId, retrieveDatasetDto));
+    }
+
     private static final List<String> ALLOWED_TYPES = Arrays.asList("txt", "markdown", "md", "mdx", "pdf", "html", "xlsx", "xls", "docx", "csv", "vtt", "properties", "htm");
 
     @ApiOperation("新增文档")

+ 26 - 0
virgo.api/src/main/java/com/bosshand/virgo/api/workark/model/RetrieveDatasetDto.java

@@ -0,0 +1,26 @@
+package com.bosshand.virgo.api.workark.model;
+
+import io.github.imfangs.dify.client.model.datasets.RetrievalModel;
+
+public class RetrieveDatasetDto {
+
+    private String query;
+
+    private RetrievalModel retrievalModel;
+
+    public String getQuery() {
+        return query;
+    }
+
+    public void setQuery(String query) {
+        this.query = query;
+    }
+
+    public RetrievalModel getRetrievalModel() {
+        return retrievalModel;
+    }
+
+    public void setRetrievalModel(RetrievalModel retrievalModel) {
+        this.retrievalModel = retrievalModel;
+    }
+}

+ 25 - 1
virgo.api/src/main/java/com/bosshand/virgo/api/workark/service/DifyDatasetService.java

@@ -4,6 +4,7 @@ import com.bosshand.virgo.api.workark.dao.DifyDatasetDao;
 import com.bosshand.virgo.api.workark.dao.DifyDatasetDocumentDao;
 import com.bosshand.virgo.api.workark.model.DifyDataset;
 import com.bosshand.virgo.api.workark.model.DifyDatasetDocument;
+import com.bosshand.virgo.api.workark.model.RetrieveDatasetDto;
 import io.github.imfangs.dify.client.DifyClientFactory;
 import io.github.imfangs.dify.client.DifyDatasetsClient;
 import io.github.imfangs.dify.client.exception.DifyApiException;
@@ -25,7 +26,7 @@ public class DifyDatasetService {
     DifyDatasetDocumentDao difyDatasetDocumentDao;
 
     private DifyDatasetsClient getClient() {
-        return DifyClientFactory.createDatasetsClient("http://203.110.233.149:9000/v1", "dataset-SWjJp6FOFqT85n7KxxyCFPSS");
+        return DifyClientFactory.createDatasetsClient("http://203.110.233.149:80/v1", "dataset-SWjJp6FOFqT85n7KxxyCFPSS");
     }
 
     /**
@@ -204,4 +205,27 @@ public class DifyDatasetService {
         return difyDatasetDocumentDao.getDatasetId(datasetId);
     }
 
+    public RetrieveResponse retrieveDataset(String datasetId, RetrieveDatasetDto dto) {
+
+        String query = dto.getQuery();
+
+        RetrievalModel retrievalModel = dto.getRetrievalModel();
+
+        RetrieveRequest request = RetrieveRequest.builder()
+                .query(query)
+                .retrievalModel(retrievalModel)
+                .build();
+
+        // 发送请求
+        try {
+            return getClient().retrieveDataset(datasetId, request);
+        } catch (IOException e) {
+            e.printStackTrace();
+        } catch (DifyApiException e) {
+            e.printStackTrace();
+        }
+        return null;
+    }
+
+
 }