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.CreateCollectionParam;
  11 import org.onap.usecaseui.llmadaptation.bean.fastgpt.CreateDataSetParam;
  12 import org.onap.usecaseui.llmadaptation.bean.fastgpt.CreateDataSetResponse;
  13 import org.onap.usecaseui.llmadaptation.constant.FastGptConstant;
  14 import org.onap.usecaseui.llmadaptation.mapper.FastGptDatasetMapper;
  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.util.CollectionUtils;
  25 import org.springframework.web.reactive.function.BodyInserters;
  26 import org.springframework.web.reactive.function.client.WebClient;
  27 import reactor.core.publisher.Flux;
  28 import reactor.core.publisher.Mono;
  29
  30 import java.util.List;
  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 FastGptDatasetMapper fastGptDatasetMapper;
  39
  40     @Autowired
  41     private WebClient webClient;
  42
  43     public Mono<ServiceResult> createDataset(Flux<FilePart> fileParts, String metaData) {
  44         KnowledgeBase knowledgeBase = JSONObject.parseObject(metaData, KnowledgeBase.class);
  45         knowledgeBase.setUpdateTime(TimeUtil.getNowTime());
  46         CreateDataSetParam dataSetParam = new CreateDataSetParam();
  47         dataSetParam.setAgentModel("qwen2:7b");
  48         dataSetParam.setType("dataset");
  49         dataSetParam.setAvatar("core/dataset/commonDatasetColor");
  50         dataSetParam.setVectorModel("m3e");
  51         dataSetParam.setIntro(knowledgeBase.getKnowledgeBaseDescription());
  52         dataSetParam.setName(knowledgeBase.getKnowledgeBaseName());
  53         return webClient.post()
  54                 .uri(FastGptConstant.CREATE_DATASET_URL)
  55                 .contentType(APPLICATION_JSON)
  56                 .header(FastGptConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
  57                 .bodyValue(dataSetParam)
  58                 .retrieve()
  59                 .bodyToMono(CreateDataSetResponse.class)
  60                 .flatMap(response -> {
  61                     if (response.getCode() == 200) {
  62                         String knowledgeBaseId = String.valueOf(response.getData());
  63                         return fileParts.flatMap(filePart -> uploadFile(filePart, knowledgeBaseId))
  64                                 .then(Mono.defer(() -> {
  65                                     knowledgeBase.setKnowledgeBaseId(knowledgeBaseId);
  66                                     fastGptDatasetMapper.insertKnowledgeBaseRecord(knowledgeBase);
  67                                     return Mono.just(new ServiceResult(new ResultHeader(200, "create success")));
  68                                 }))
  69                                 .onErrorResume(e -> {
  70                                     log.error("Error occurred during file upload: {}", e.getMessage());
  71                                     return Mono.just(new ServiceResult(new ResultHeader(500, "file upload failed")));
  72                                 });
  73                     } else {
  74                         return Mono.just(new ServiceResult(new ResultHeader(500, response.getMessage())));
  75                     }
  76                 })
  77                 .onErrorResume(e -> {
  78                     log.error("Error occurred while creating dataset: {}", e.getMessage());
  79                     return Mono.just(new ServiceResult(new ResultHeader(500, "create failed")));
  80                 });
  81     }
  82
  83     private Mono<Void> uploadFile(FilePart filePart, String knowledgeBaseId) {
  84         String filename = filePart.filename();
  85         Flux<DataBuffer> content = filePart.content();
  86
  87         MultipartBodyBuilder builder = new MultipartBodyBuilder();
  88         builder.part("metadata", "", APPLICATION_JSON);
  89         builder.part("bucketName", "dataset");
  90         builder.asyncPart("file", content, DataBuffer.class)
  91                 .headers(headers -> {
  92                     ContentDisposition contentDisposition = ContentDisposition
  93                             .builder("form-data")
  94                             .name("file")
  95                             .filename(filename)
  96                             .build();
  97                     headers.setContentDisposition(contentDisposition);
  98                     headers.setContentType(MediaType.TEXT_PLAIN);
  99                 });
 100
 101         return webClient.post()
 102                 .uri(FastGptConstant.UPLOAD_FILE_URL)
 103                 .contentType(MediaType.MULTIPART_FORM_DATA)
 104                 .header(FastGptConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
 105                 .body(BodyInserters.fromMultipartData(builder.build()))
 106                 .retrieve()
 107                 .bodyToMono(CreateDataSetResponse.class)
 108                 .flatMap(response -> {
 109                     if (response.getCode() != 200) {
 110                         log.error("response is {}", response);
 111                         return Mono.empty();
 112                     }
 113                     Object data = response.getData();
 114                     JSONObject jsonObject = JSON.parseObject(JSON.toJSONString(data));
 115                     String fileId = jsonObject.getString("fileId");
 116                     CreateCollectionParam createCollectionParam = getCreateCollectionParam(knowledgeBaseId, fileId);
 117
 118                     return webClient.post()
 119                             .uri(FastGptConstant.CRATE_COLLECTION_URL)
 120                             .contentType(APPLICATION_JSON)
 121                             .header(FastGptConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
 122                             .bodyValue(createCollectionParam)
 123                             .retrieve()
 124                             .bodyToMono(CreateDataSetResponse.class)
 125                             .flatMap(responseData -> {
 126                                 if (responseData.getCode() == 200) {
 127                                     fastGptDatasetMapper.insertFileName(fileId, filename, knowledgeBaseId);
 128                                 }
 129                                 return Mono.empty();
 130                             });
 131                 });
 132     }
 133
 134     @NotNull
 135     private static CreateCollectionParam getCreateCollectionParam(String knowledgeBaseId, String fileId) {
 136         CreateCollectionParam createCollectionParam = new CreateCollectionParam();
 137         createCollectionParam.setTrainingType("chunk");
 138         createCollectionParam.setDatasetId(knowledgeBaseId);
 139         createCollectionParam.setChunkSize(700);
 140         createCollectionParam.setChunkSplitter("");
 141         createCollectionParam.setFileId(fileId);
 142         createCollectionParam.setName("");
 143         createCollectionParam.setQaPrompt("");
 144         return createCollectionParam;
 145     }
 146
 147     public Mono<ServiceResult> removeDataset(String knowledgeBaseId) {
 148         String url = FastGptConstant.DELETE_DATASET_URL + knowledgeBaseId;
 149         return webClient.delete()
 150                 .uri(url)
 151                 .header(FastGptConstant.COOKIE, FastGptConstant.COOKIE_VALUE)
 152                 .retrieve()
 153                 .bodyToMono(CreateDataSetResponse.class)
 154                 .flatMap(response -> {
 155                     if (response.getCode() == 200) {
 156                         return Mono.fromRunnable(() -> {
 157                             try {
 158                                 fastGptDatasetMapper.deleteKnowledgeBaseByUuid(knowledgeBaseId);
 159                                 fastGptDatasetMapper.deleteFileById(knowledgeBaseId);
 160                             } catch (Exception dbException) {
 161                                 throw new RuntimeException("Database operation failed", dbException); // 抛出新异常
 162                             }
 163                         }).then(Mono.just(new ServiceResult(new ResultHeader(200, "delete success"))));
 164                     } else {
 165                         return Mono.just(new ServiceResult(new ResultHeader(500, response.getStatusText())));
 166                     }
 167                 })
 168                 .onErrorResume(e -> {
 169                     log.error("Error occurred while delete dataset: {}", e.getMessage());
 170                     return Mono.just(new ServiceResult(new ResultHeader(500, "delete failed")));
 171                 });
 172     }
 173
 174     public ServiceResult getDataSetRecord() {
 175         List<KnowledgeBase> knowledgeBaseRecords = fastGptDatasetMapper.getKnowledgeBaseRecords();
 176         if (CollectionUtils.isEmpty(knowledgeBaseRecords)) {
 177             return new ServiceResult(new ResultHeader(200, "get datasets failed"), knowledgeBaseRecords);
 178         }
 179
 180         knowledgeBaseRecords.forEach(knowledgeBase -> {
 181             List<String> fileNamesByKnowledgeBaseId = fastGptDatasetMapper.getFileNamesByKnowledgeBaseId(knowledgeBase.getKnowledgeBaseId());
 182             knowledgeBase.setFilesName(fileNamesByKnowledgeBaseId);
 183         });
 184         return new ServiceResult(new ResultHeader(200, "success"), knowledgeBaseRecords);
 185     }
 186
 187     public ServiceResult geDatasetById(String knowledgeBaseId) {
 188         KnowledgeBase knowledgeBase = fastGptDatasetMapper.getKnowledgeBaseRecordById(knowledgeBaseId);
 189         if (knowledgeBase == null) {
 190             return new ServiceResult(new ResultHeader(200, "get dataset failed"));
 191         }
 192         List<String> fileNamesByKnowledgeBaseId = fastGptDatasetMapper.getFileNamesByKnowledgeBaseId(knowledgeBase.getKnowledgeBaseId());
 193         knowledgeBase.setFilesName(fileNamesByKnowledgeBaseId);
 194
 195         return new ServiceResult(new ResultHeader(200, "success"), knowledgeBase);
 196     }
 197
 198 }