gerrit.onap Code Review - usecase-ui/llm-adaptation.git/blob

   1 package org.onap.usecaseui.llmadaptation.service.impl;
   2
   3 import com.alibaba.fastjson2.JSON;
   4 import com.alibaba.fastjson2.JSONArray;
   5 import com.alibaba.fastjson2.JSONObject;
   6 import lombok.extern.slf4j.Slf4j;
   7 import org.jetbrains.annotations.NotNull;
   8 import org.onap.usecaseui.llmadaptation.bean.*;
   9 import org.onap.usecaseui.llmadaptation.bean.fastgpt.dataset.CreateCollectionParam;
  10 import org.onap.usecaseui.llmadaptation.bean.fastgpt.dataset.CreateDataSetParam;
  11 import org.onap.usecaseui.llmadaptation.bean.fastgpt.dataset.CreateDataSetResponse;
  12 import org.onap.usecaseui.llmadaptation.constant.CommonConstant;
  13 import org.onap.usecaseui.llmadaptation.constant.FastGptConstant;
  14 import org.onap.usecaseui.llmadaptation.mapper.DatasetMapper;
  15 import org.onap.usecaseui.llmadaptation.service.FastGptDatasetService;
  16 import org.onap.usecaseui.llmadaptation.util.TimeUtil;
  17 import org.springframework.beans.factory.annotation.Autowired;
  18 import org.springframework.core.io.buffer.DataBuffer;
  19 import org.springframework.http.ContentDisposition;
  20 import org.springframework.http.MediaType;
  21 import org.springframework.http.client.MultipartBodyBuilder;
  22 import org.springframework.http.codec.multipart.FilePart;
  23 import org.springframework.stereotype.Service;
  24 import org.springframework.web.reactive.function.BodyInserters;
  25 import org.springframework.web.reactive.function.client.WebClient;
  26 import reactor.core.publisher.Flux;
  27 import reactor.core.publisher.Mono;
  28
  29
  30 import java.util.List;
  31 import java.util.Map;
  32 import java.util.stream.Collectors;
  33 import java.util.stream.IntStream;
  34
  35 import static org.springframework.http.MediaType.APPLICATION_JSON;
  36
  37 @Slf4j
  38 @Service
  39 public class FastGptDatasetServiceImpl implements FastGptDatasetService {
  40     @Autowired
  41     private DatasetMapper datasetMapper;
  42
  43     @Autowired
  44     private WebClient webClient;
  45
  46     @Override
  47     public Mono<ServiceResult> createDataset(Flux<FilePart> fileParts, String metaData, MaaSPlatform maaSPlatform) {
  48         KnowledgeBase knowledgeBase = JSONObject.parseObject(metaData, KnowledgeBase.class);
  49         knowledgeBase.setUpdateTime(TimeUtil.getNowTime());
  50         CreateDataSetParam dataSetParam = new CreateDataSetParam();
  51         dataSetParam.setAgentModel(maaSPlatform.getVectorModel());
  52         dataSetParam.setType("dataset");
  53         dataSetParam.setIntro(knowledgeBase.getKnowledgeBaseDescription());
  54         dataSetParam.setName(knowledgeBase.getKnowledgeBaseName());
  55         return webClient.post()
  56                 .uri(maaSPlatform.getServerIp() + FastGptConstant.CREATE_DATASET_URL)
  57                 .contentType(APPLICATION_JSON)
  58                 .header(CommonConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
  59                 .bodyValue(dataSetParam)
  60                 .retrieve()
  61                 .bodyToMono(CreateDataSetResponse.class)
  62                 .flatMap(response -> {
  63                     if (response.getCode() == 200) {
  64                         String knowledgeBaseId = String.valueOf(response.getData());
  65                         return fileParts
  66                                 .flatMap(filePart -> uploadFile(filePart, knowledgeBaseId, maaSPlatform.getServerIp()))
  67                                 .then(Mono.defer(() -> {
  68                                     knowledgeBase.setKnowledgeBaseId(knowledgeBaseId);
  69                                     datasetMapper.insertKnowledgeBaseRecord(knowledgeBase);
  70                                     return handleFileId(knowledgeBaseId, maaSPlatform.getServerIp());
  71                                 }))
  72                                 .onErrorResume(e -> {
  73                                     log.error("Error occurred during file upload: {}", e.getMessage());
  74                                     return Mono.just(new ServiceResult(new ResultHeader(500, "file upload failed")));
  75                                 });
  76                     } else {
  77                         return Mono.just(new ServiceResult(new ResultHeader(500, response.getMessage())));
  78                     }
  79                 })
  80                 .onErrorResume(e -> {
  81                     log.error("Error occurred while creating dataset: {}", e.getMessage());
  82                     return Mono.just(new ServiceResult(new ResultHeader(500, "create failed")));
  83                 });
  84     }
  85
  86     private Mono<ServiceResult> handleFileId(String knowledgeBaseId, String serverIp) {
  87         JSONObject jsonObject = new JSONObject();
  88         jsonObject.put("datasetId", knowledgeBaseId);
  89         return webClient.post()
  90                 .uri(serverIp + FastGptConstant.GET_COLLECTION_LIST_URL)
  91                 .header(CommonConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
  92                 .bodyValue(jsonObject)
  93                 .retrieve()
  94                 .bodyToMono(CreateDataSetResponse.class)
  95                 .flatMap(response -> {
  96                     Object data = response.getData();
  97                     JSONArray jsonArray = JSONObject.parseObject(JSONObject.toJSONString(data)).getJSONArray("data");
  98                     Map<String, String> resultMap = IntStream.range(0, jsonArray.size())
  99                             .mapToObj(jsonArray::getJSONObject)
 100                             .collect(Collectors.toMap(
 101                                     obj -> obj.getString("fileId"),
 102                                     obj -> obj.getString("_id")
 103                             ));
 104                     List<File> fileList = datasetMapper.getFileNamesByKnowledgeBaseId(knowledgeBaseId);
 105                     List<File> updatedFileList = fileList.stream()
 106                             .map(file -> new File(
 107                                     resultMap.getOrDefault(file.getFileId(), file.getFileId()),
 108                                     file.getFileName()))
 109                             .toList();
 110                     datasetMapper.deleteFileById(knowledgeBaseId);
 111                     datasetMapper.insertFileName(updatedFileList, knowledgeBaseId);
 112                     return Mono.just(new ServiceResult(new ResultHeader(200, "create success")));
 113                 });
 114     }
 115
 116     private Mono<Void> uploadFile(FilePart filePart, String knowledgeBaseId, String serverIp) {
 117         String filename = filePart.filename();
 118         Flux<DataBuffer> content = filePart.content();
 119
 120         MultipartBodyBuilder builder = new MultipartBodyBuilder();
 121         builder.part("metadata", "", APPLICATION_JSON);
 122         builder.part("bucketName", "dataset");
 123         builder.asyncPart("file", content, DataBuffer.class)
 124                 .headers(headers -> {
 125                     ContentDisposition contentDisposition = ContentDisposition
 126                             .builder("form-data")
 127                             .name("file")
 128                             .filename(filename)
 129                             .build();
 130                     headers.setContentDisposition(contentDisposition);
 131                     headers.setContentType(MediaType.TEXT_PLAIN);
 132                 });
 133
 134         return webClient.post()
 135                 .uri(serverIp + FastGptConstant.UPLOAD_FILE_URL)
 136                 .contentType(MediaType.MULTIPART_FORM_DATA)
 137                 .header(CommonConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
 138                 .body(BodyInserters.fromMultipartData(builder.build()))
 139                 .retrieve()
 140                 .bodyToMono(CreateDataSetResponse.class)
 141                 .flatMap(response -> {
 142                     if (response.getCode() != 200) {
 143                         log.error("response is {}", response);
 144                         return Mono.empty();
 145                     }
 146                     Object data = response.getData();
 147                     JSONObject jsonObject = JSON.parseObject(JSON.toJSONString(data));
 148                     String fileId = jsonObject.getString("fileId");
 149                     CreateCollectionParam createCollectionParam = getCreateCollectionParam(knowledgeBaseId, fileId, filename);
 150
 151                     return webClient.post()
 152                             .uri(serverIp + FastGptConstant.CRATE_COLLECTION_URL)
 153                             .contentType(APPLICATION_JSON)
 154                             .header(CommonConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
 155                             .bodyValue(createCollectionParam)
 156                             .retrieve()
 157                             .bodyToMono(CreateDataSetResponse.class)
 158                             .flatMap(responseData -> {
 159                                 if (responseData.getCode() == 200) {
 160                                     File file = new File(String.valueOf(fileId), filename);
 161                                     datasetMapper.insertFileName(List.of(file), String.valueOf(knowledgeBaseId));
 162                                 }
 163                                 return Mono.empty();
 164                             });
 165                 });
 166     }
 167
 168     @NotNull
 169     private static CreateCollectionParam getCreateCollectionParam(String knowledgeBaseId, String fileId, String fileName) {
 170         CreateCollectionParam createCollectionParam = new CreateCollectionParam();
 171         createCollectionParam.setTrainingType("chunk");
 172         createCollectionParam.setDatasetId(knowledgeBaseId);
 173         createCollectionParam.setChunkSize(700);
 174         createCollectionParam.setChunkSplitter("");
 175         createCollectionParam.setFileId(fileId);
 176         createCollectionParam.setName(fileName);
 177         createCollectionParam.setQaPrompt("");
 178         return createCollectionParam;
 179     }
 180
 181     @Override
 182     public Mono<ServiceResult> removeDataset(String knowledgeBaseId, String serverIp) {
 183         String url = serverIp + FastGptConstant.DELETE_DATASET_URL + knowledgeBaseId;
 184         return webClient.delete()
 185                 .uri(url)
 186                 .header(CommonConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
 187                 .retrieve()
 188                 .bodyToMono(CreateDataSetResponse.class)
 189                 .flatMap(response -> {
 190                     if (response.getCode() == 200) {
 191                         return Mono.fromRunnable(() -> {
 192                             try {
 193                                 datasetMapper.deleteKnowledgeBaseByUuid(knowledgeBaseId);
 194                                 datasetMapper.deleteFileById(knowledgeBaseId);
 195                             } catch (Exception dbException) {
 196                                 throw new RuntimeException("Database operation failed", dbException);
 197                             }
 198                         }).then(Mono.just(new ServiceResult(new ResultHeader(200, "delete success"))));
 199                     } else {
 200                         return Mono.just(new ServiceResult(new ResultHeader(500, response.getStatusText())));
 201                     }
 202                 })
 203                 .onErrorResume(e -> {
 204                     log.error("Error occurred while delete dataset: {}", e.getMessage());
 205                     return Mono.just(new ServiceResult(new ResultHeader(500, "delete failed")));
 206                 });
 207     }
 208
 209     @Override
 210     public Mono<ServiceResult> editDataset(KnowledgeBase knowledgeBase, MaaSPlatform maaSPlatform) {
 211         KnowledgeBase knowledgeBaseRecordById = datasetMapper.getKnowledgeBaseRecordById(knowledgeBase.getKnowledgeBaseId());
 212         if (knowledgeBaseRecordById == null) {
 213             return Mono.just(new ServiceResult(new ResultHeader(500, "dataset is not exist")));
 214         }
 215         JSONObject updateParam = new JSONObject();
 216         updateParam.put("id", knowledgeBase.getKnowledgeBaseId());
 217         updateParam.put("name", knowledgeBase.getKnowledgeBaseName());
 218         updateParam.put("intro", knowledgeBase.getKnowledgeBaseDescription());
 219         updateParam.put("avatar", "core/dataset/commonDatasetColor");
 220
 221         return webClient.put()
 222                 .uri(maaSPlatform.getServerIp() + FastGptConstant.UPDATE_DATASET_URL)
 223                 .contentType(APPLICATION_JSON)
 224                 .header(CommonConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
 225                 .bodyValue(updateParam)
 226                 .retrieve()
 227                 .bodyToMono(CreateDataSetResponse.class)
 228                 .flatMap(response -> {
 229                     if (response.getCode() == 200) {
 230                         return Mono.fromRunnable(() -> {
 231                             knowledgeBase.setUpdateTime(TimeUtil.getNowTime());
 232                             datasetMapper.updateKnowledgeBase(knowledgeBase);
 233                         }).then(Mono.just(new ServiceResult(new ResultHeader(200, "update success"))));
 234                     } else {
 235                         return Mono.just(new ServiceResult(new ResultHeader(500, response.getStatusText())));
 236                     }
 237                 })
 238                 .onErrorResume(e -> {
 239                     log.error("Error occurred while delete dataset: {}", e.getMessage());
 240                     return Mono.just(new ServiceResult(new ResultHeader(500, "update failed")));
 241                 });
 242     }
 243
 244     @Override
 245     public Mono<ServiceResult> uploadFiles(Flux<FilePart> fileParts, String knowledgeBaseId, String serverIp) {
 246         return fileParts.flatMap(filePart -> uploadFile(filePart, knowledgeBaseId, serverIp))
 247                 .then(Mono.defer(() -> handleFileId(knowledgeBaseId, serverIp)))
 248                 .onErrorResume(e -> {
 249                     log.error("Error occurred during file upload: {}", e.getMessage());
 250                     return Mono.just(new ServiceResult(new ResultHeader(500, "file upload failed")));
 251                 });
 252     }
 253
 254     @Override
 255     public Mono<ServiceResult> deleteFile(String fileId, String serverIp) {
 256         return webClient.delete().uri(serverIp + FastGptConstant.DELETE_FILE_URL + fileId)
 257                 .header(CommonConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
 258                 .retrieve()
 259                 .bodyToMono(CreateDataSetResponse.class)
 260                 .flatMap(response -> {
 261                     if (response.getCode() == 200) {
 262                         return Mono.fromRunnable(() -> datasetMapper.deleteFileByFileId(fileId)).then(Mono.just(new ServiceResult(new ResultHeader(200, "delete file success"))));
 263                     } else {
 264                         return Mono.just(new ServiceResult(new ResultHeader(response.getCode(), response.getStatusText())));
 265                     }
 266                 })
 267                 .onErrorResume(e -> {
 268                     log.error("Error occurred while delete dataset: {}", e.getMessage());
 269                     return Mono.just(new ServiceResult(new ResultHeader(500, "delete file failed")));
 270                 });
 271     }
 272 }