fix

virgo.api/src/main/java/com/bosshand/virgo/api/workark/service/DifyDatasetService.java

@@ -1,5 +1,6 @@
 package com.bosshand.virgo.api.workark.service;
 
+import com.alibaba.fastjson.JSONArray;
 import com.alibaba.fastjson.JSONObject;
 import com.bosshand.virgo.api.workark.dao.DifyDatasetDao;
 import com.bosshand.virgo.api.workark.dao.DifyDatasetDocumentDao;
@@ -309,7 +310,7 @@ public class DifyDatasetService {
 
     }
 
-    public RetrieveResponse retrieveDataset(String datasetId, JSONObject query) throws IOException {
+    public JSONObject retrieveDataset(String datasetId, JSONObject query) throws IOException {
         String js = "{\n" +
                 "    \"retrieval_model\": {\n" +
                 "        \"metadata_filtering_conditions\": {},\n" +
@@ -332,7 +333,15 @@ public class DifyDatasetService {
         String data = JSONObject.toJSONString(json);
         String response = retrieve(apiUrl, "dataset-SWjJp6FOFqT85n7KxxyCFPSS", datasetId, data);
         RetrieveResponse retrieveResponse = JSONObject.parseObject(response, RetrieveResponse.class);
-        return retrieveResponse;
+        JSONObject responseJson = new JSONObject();
+        responseJson.put("query", retrieveResponse.getQuery().getContent());
+        JSONArray array = new JSONArray();
+        List<RetrieveResponse.Record> records = retrieveResponse.getRecords();
+        for (RetrieveResponse.Record r : records) {
+            array.add(r.getSegment().getContent());
+        }
+        responseJson.put("records", array);
+        return responseJson;
     }
 
 }