dcs 1 day ago
parent
commit
1d6fedff36

+ 21 - 6
virgo.api/src/main/java/com/bosshand/virgo/api/workark/service/DifyDatasetService.java

@@ -309,15 +309,30 @@ public class DifyDatasetService {
 
     }
 
-    public RetrieveResponse retrieveDataset(String datasetId, JSONObject jsonobject) throws IOException {
+    public RetrieveResponse retrieveDataset(String datasetId, JSONObject query) throws IOException {
+        String js = "{\n" +
+                "    \"retrieval_model\": {\n" +
+                "        \"metadata_filtering_conditions\": {},\n" +
+                "        \"reranking_enable\": false,\n" +
+                "        \"reranking_mode\": \"None\",\n" +
+                "        \"reranking_model\": {\n" +
+                "            \"reranking_model_name\": \"\",\n" +
+                "            \"reranking_provider_name\": \"\"\n" +
+                "        },\n" +
+                "        \"score_threshold\": \"None\",\n" +
+                "        \"score_threshold_enabled\": false,\n" +
+                "        \"search_method\": \"keyword_search\",\n" +
+                "        \"top_k\": 1,\n" +
+                "        \"weights\": \"None\"\n" +
+                "    }\n" +
+                "}";
+        JSONObject json = JSONObject.parseObject(js);
+        json.put("query", query.get("query"));
         String apiUrl = "http://203.110.233.149:80/v1/datasets/{dataset_id}/retrieve";
-        String json = JSONObject.toJSONString(jsonobject);
-        String response = retrieve(apiUrl, "dataset-SWjJp6FOFqT85n7KxxyCFPSS", datasetId, json);
+        String data = JSONObject.toJSONString(json);
+        String response = retrieve(apiUrl, "dataset-SWjJp6FOFqT85n7KxxyCFPSS", datasetId, data);
         RetrieveResponse retrieveResponse = JSONObject.parseObject(response, RetrieveResponse.class);
         return retrieveResponse;
     }
 
-
-
-
 }