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.JSONObject;
   5 import lombok.extern.slf4j.Slf4j;
   6 import org.jetbrains.annotations.NotNull;
   7 import org.onap.usecaseui.llmadaptation.bean.KnowledgeBase;
   8 import org.onap.usecaseui.llmadaptation.bean.ResultHeader;
   9 import org.onap.usecaseui.llmadaptation.bean.ServiceResult;
  10 import org.onap.usecaseui.llmadaptation.bean.fastgpt.dataset.CreateCollectionParam;
  11 import org.onap.usecaseui.llmadaptation.bean.fastgpt.dataset.CreateDataSetParam;
  12 import org.onap.usecaseui.llmadaptation.bean.fastgpt.dataset.CreateDataSetResponse;
  13 import org.onap.usecaseui.llmadaptation.constant.CommonConstant;
  14 import org.onap.usecaseui.llmadaptation.constant.FastGptConstant;
  15 import org.onap.usecaseui.llmadaptation.constant.ServerConstant;
  16 import org.onap.usecaseui.llmadaptation.mapper.DatasetMapper;
  17 import org.onap.usecaseui.llmadaptation.service.FastGptDatasetService;
  18 import org.onap.usecaseui.llmadaptation.util.TimeUtil;
  19 import org.springframework.beans.factory.annotation.Autowired;
  20 import org.springframework.core.io.buffer.DataBuffer;
  21 import org.springframework.http.ContentDisposition;
  22 import org.springframework.http.MediaType;
  23 import org.springframework.http.client.MultipartBodyBuilder;
  24 import org.springframework.http.codec.multipart.FilePart;
  25 import org.springframework.stereotype.Service;
  26 import org.springframework.web.reactive.function.BodyInserters;
  27 import org.springframework.web.reactive.function.client.WebClient;
  28 import reactor.core.publisher.Flux;
  29 import reactor.core.publisher.Mono;
  30
  31
  32 import static org.springframework.http.MediaType.APPLICATION_JSON;
  33
  34 @Slf4j
  35 @Service
  36 public class FastGptDatasetServiceImpl implements FastGptDatasetService {
  37     @Autowired
  38     private DatasetMapper datasetMapper;
  39
  40     @Autowired
  41     private WebClient webClient;
  42
  43     @Autowired
  44     private ServerConstant serverConstant;
  45
  46     @Override
  47     public Mono<ServiceResult> createDataset(Flux<FilePart> fileParts, String metaData) {
  48         KnowledgeBase knowledgeBase = JSONObject.parseObject(metaData, KnowledgeBase.class);
  49         knowledgeBase.setUpdateTime(TimeUtil.getNowTime());
  50         CreateDataSetParam dataSetParam = new CreateDataSetParam();
  51         dataSetParam.setAgentModel(serverConstant.getFastGptModel());
  52         dataSetParam.setType("dataset");
  53         dataSetParam.setAvatar("core/dataset/commonDatasetColor");
  54         dataSetParam.setVectorModel("m3e");
  55         dataSetParam.setIntro(knowledgeBase.getKnowledgeBaseDescription());
  56         dataSetParam.setName(knowledgeBase.getKnowledgeBaseName());
  57         return webClient.post()
  58                 .uri(serverConstant.getFastGptServer() + FastGptConstant.CREATE_DATASET_URL)
  59                 .contentType(APPLICATION_JSON)
  60                 .header(CommonConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
  61                 .bodyValue(dataSetParam)
  62                 .retrieve()
  63                 .bodyToMono(CreateDataSetResponse.class)
  64                 .flatMap(response -> {
  65                     if (response.getCode() == 200) {
  66                         String knowledgeBaseId = String.valueOf(response.getData());
  67                         return fileParts.flatMap(filePart -> uploadFile(filePart, knowledgeBaseId))
  68                                 .then(Mono.defer(() -> {
  69                                     knowledgeBase.setKnowledgeBaseId(knowledgeBaseId);
  70                                     datasetMapper.insertKnowledgeBaseRecord(knowledgeBase);
  71                                     return Mono.just(new ServiceResult(new ResultHeader(200, "create success")));
  72                                 }))
  73                                 .onErrorResume(e -> {
  74                                     log.error("Error occurred during file upload: {}", e.getMessage());
  75                                     return Mono.just(new ServiceResult(new ResultHeader(500, "file upload failed")));
  76                                 });
  77                     } else {
  78                         return Mono.just(new ServiceResult(new ResultHeader(500, response.getMessage())));
  79                     }
  80                 })
  81                 .onErrorResume(e -> {
  82                     log.error("Error occurred while creating dataset: {}", e.getMessage());
  83                     return Mono.just(new ServiceResult(new ResultHeader(500, "create failed")));
  84                 });
  85     }
  86
  87     private Mono<Void> uploadFile(FilePart filePart, String knowledgeBaseId) {
  88         String filename = filePart.filename();
  89         Flux<DataBuffer> content = filePart.content();
  90
  91         MultipartBodyBuilder builder = new MultipartBodyBuilder();
  92         builder.part("metadata", "", APPLICATION_JSON);
  93         builder.part("bucketName", "dataset");
  94         builder.asyncPart("file", content, DataBuffer.class)
  95                 .headers(headers -> {
  96                     ContentDisposition contentDisposition = ContentDisposition
  97                             .builder("form-data")
  98                             .name("file")
  99                             .filename(filename)
 100                             .build();
 101                     headers.setContentDisposition(contentDisposition);
 102                     headers.setContentType(MediaType.TEXT_PLAIN);
 103                 });
 104
 105         return webClient.post()
 106                 .uri(serverConstant.getFastGptServer() + FastGptConstant.UPLOAD_FILE_URL)
 107                 .contentType(MediaType.MULTIPART_FORM_DATA)
 108                 .header(CommonConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
 109                 .body(BodyInserters.fromMultipartData(builder.build()))
 110                 .retrieve()
 111                 .bodyToMono(CreateDataSetResponse.class)
 112                 .flatMap(response -> {
 113                     if (response.getCode() != 200) {
 114                         log.error("response is {}", response);
 115                         return Mono.empty();
 116                     }
 117                     Object data = response.getData();
 118                     JSONObject jsonObject = JSON.parseObject(JSON.toJSONString(data));
 119                     String fileId = jsonObject.getString("fileId");
 120                     CreateCollectionParam createCollectionParam = getCreateCollectionParam(knowledgeBaseId, fileId);
 121
 122                     return webClient.post()
 123                             .uri(serverConstant.getFastGptServer() + FastGptConstant.CRATE_COLLECTION_URL)
 124                             .contentType(APPLICATION_JSON)
 125                             .header(CommonConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
 126                             .bodyValue(createCollectionParam)
 127                             .retrieve()
 128                             .bodyToMono(CreateDataSetResponse.class)
 129                             .flatMap(responseData -> {
 130                                 if (responseData.getCode() == 200) {
 131                                     datasetMapper.insertFileName(fileId, filename, knowledgeBaseId);
 132                                 }
 133                                 return Mono.empty();
 134                             });
 135                 });
 136     }
 137
 138     @NotNull
 139     private static CreateCollectionParam getCreateCollectionParam(String knowledgeBaseId, String fileId) {
 140         CreateCollectionParam createCollectionParam = new CreateCollectionParam();
 141         createCollectionParam.setTrainingType("chunk");
 142         createCollectionParam.setDatasetId(knowledgeBaseId);
 143         createCollectionParam.setChunkSize(700);
 144         createCollectionParam.setChunkSplitter("");
 145         createCollectionParam.setFileId(fileId);
 146         createCollectionParam.setName("");
 147         createCollectionParam.setQaPrompt("");
 148         return createCollectionParam;
 149     }
 150
 151     @Override
 152     public Mono<ServiceResult> removeDataset(String knowledgeBaseId) {
 153         String url = serverConstant.getFastGptServer() + FastGptConstant.DELETE_DATASET_URL + knowledgeBaseId;
 154         return webClient.delete()
 155                 .uri(url)
 156                 .header(CommonConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
 157                 .retrieve()
 158                 .bodyToMono(CreateDataSetResponse.class)
 159                 .flatMap(response -> {
 160                     if (response.getCode() == 200) {
 161                         return Mono.fromRunnable(() -> {
 162                             try {
 163                                 datasetMapper.deleteKnowledgeBaseByUuid(knowledgeBaseId);
 164                                 datasetMapper.deleteFileById(knowledgeBaseId);
 165                             } catch (Exception dbException) {
 166                                 throw new RuntimeException("Database operation failed", dbException);
 167                             }
 168                         }).then(Mono.just(new ServiceResult(new ResultHeader(200, "delete success"))));
 169                     } else {
 170                         return Mono.just(new ServiceResult(new ResultHeader(500, response.getStatusText())));
 171                     }
 172                 })
 173                 .onErrorResume(e -> {
 174                     log.error("Error occurred while delete dataset: {}", e.getMessage());
 175                     return Mono.just(new ServiceResult(new ResultHeader(500, "delete failed")));
 176                 });
 177     }
 178
 179     @Override
 180     public Mono<ServiceResult> editDataset(KnowledgeBase knowledgeBase) {
 181         KnowledgeBase knowledgeBaseRecordById = datasetMapper.getKnowledgeBaseRecordById(knowledgeBase.getKnowledgeBaseId());
 182         if (knowledgeBaseRecordById == null) {
 183             return Mono.just(new ServiceResult(new ResultHeader(500, "dataset is not exist")));
 184         }
 185         JSONObject updateParam = new JSONObject();
 186         updateParam.put("id", knowledgeBase.getKnowledgeBaseId());
 187         updateParam.put("name", knowledgeBase.getKnowledgeBaseName());
 188         updateParam.put("intro", knowledgeBase.getKnowledgeBaseDescription());
 189         updateParam.put("avatar", "core/dataset/commonDatasetColor");
 190
 191         return webClient.put()
 192                 .uri(serverConstant.getFastGptServer() + FastGptConstant.UPDATE_DATASET_URL)
 193                 .contentType(APPLICATION_JSON)
 194                 .header(CommonConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
 195                 .bodyValue(updateParam)
 196                 .retrieve()
 197                 .bodyToMono(CreateDataSetResponse.class)
 198                 .flatMap(response -> {
 199                     if (response.getCode() == 200) {
 200                         return Mono.fromRunnable(() -> {
 201                             knowledgeBase.setUpdateTime(TimeUtil.getNowTime());
 202                             datasetMapper.updateKnowledgeBase(knowledgeBase);
 203                         }).then(Mono.just(new ServiceResult(new ResultHeader(200, "update success"))));
 204                     } else {
 205                         return Mono.just(new ServiceResult(new ResultHeader(500, response.getStatusText())));
 206                     }
 207                 })
 208                 .onErrorResume(e -> {
 209                     log.error("Error occurred while delete dataset: {}", e.getMessage());
 210                     return Mono.just(new ServiceResult(new ResultHeader(500, "update failed")));
 211                 });
 212     }
 213 }