xiaoshi-weixinback/src/main/java/com/example/xiaoshiweixinback/service/common/EsDenseVectorService.java

package com.example.xiaoshiweixinback.service.common;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.InlineScript;
import co.elastic.clients.elasticsearch._types.Script;
import co.elastic.clients.elasticsearch._types.query_dsl.Query;
import co.elastic.clients.elasticsearch._types.query_dsl.QueryBuilders;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.elasticsearch.core.search.Hit;
import co.elastic.clients.json.JsonData;
import com.example.xiaoshiweixinback.business.utils.BeanUtil;
import com.example.xiaoshiweixinback.business.utils.parseQueryToTree.expressManager;
import com.example.xiaoshiweixinback.business.utils.parseQueryToTree.operateNode;
import com.example.xiaoshiweixinback.business.utils.parseQueryToTree.treeNode;
import com.example.xiaoshiweixinback.domain.es.PatentVector;
import com.example.xiaoshiweixinback.entity.dto.esPicture.EsPatentSearchDTO;
import com.example.xiaoshiweixinback.entity.dto.esPicture.EsPictureNoDTO;
import com.example.xiaoshiweixinback.entity.vo.esPicture.EsPictureNoVo;
import com.example.xiaoshiweixinback.entity.vo.esPicture.EsPictureVectorVo;
import com.example.xiaoshiweixinback.service.importPatent.FormatQueryService;
import lombok.RequiredArgsConstructor;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Lazy;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

@Service
@RequiredArgsConstructor(onConstructor_ = {@Lazy})
public class EsDenseVectorService {

    private final ElasticsearchClient client;

    @Autowired
    private FormatQueryService formatQueryService;

    @Autowired
    private GetVectorService getVectorService;

    public List<EsPictureVectorVo> getPatentList(EsPatentSearchDTO dto) throws Exception {
        Long pageNum = dto.getPageNum();
        Long pageSize = dto.getPageSize();
//        String key = dto.getKey().replaceAll("[，。、；,./;\\s]"," OR ");
        String key = dto.getKey().replaceAll("[，。、；,./;]"," OR ");
        String s = "TI = " + "(" + key + ")";
        System.out.println(s);

        SearchRequest.Builder builder = new SearchRequest.Builder();
        //设置查询索引
        builder.index("patent_vector");

        //1. 解析检索条件
        treeNode tree = expressManager.getInstance().Parse(s, false);
        //3. 从es中检索数据
        Query query = formatQueryService.EsQueryToQuery((operateNode) tree, "patentVector", null);
        builder.query(query);

        //分页
        if (pageNum != null && pageSize != null && pageNum > 0 && pageSize > 0) {
            builder.from((pageNum.intValue() - 1) * pageSize.intValue()).size(pageSize.intValue());
        }
//        else {
//            builder.from(0).size(9);
//        }

        SearchResponse<PatentVector> response = client.search(builder.build(), PatentVector.class);
        List<EsPictureVectorVo> vectorVos = new ArrayList<>();
        List<Hit<PatentVector>> hits = response.hits().hits();
        for (Hit<PatentVector> hit : hits) {
            PatentVector vector = hit.source();
            EsPictureVectorVo vectorVo = new EsPictureVectorVo();
            BeanUtil.copy(vector,vectorVo);
            vectorVos.add(vectorVo);
        }

        return vectorVos;
    }

    public List<EsPictureVectorVo> getPatentVectorSort(File file, String description) throws IOException {
        List<Float> imageList = new ArrayList<>();
        List<String> stringList = new ArrayList<>();
        if (file != null) {
            stringList = getVectorService.getVectorByFile(file);
        } else if (StringUtils.isNotEmpty(description)) {
            stringList = getVectorService.getVectorByText(description);
        }
        stringList.forEach(item -> {
            Float a = Float.parseFloat(item);
            imageList.add(a);
        });

        List<EsPictureVectorVo> list = new ArrayList<>();
        SearchRequest.Builder builder = new SearchRequest.Builder();
        //设置查询索引
        builder.index("patent_vector");
        String source = "cosineSimilarity(params.queryVector, 'my_vector') + 1.0";
        InlineScript inlineScript = InlineScript.of(i -> i.lang("painless").params("queryVector", JsonData.of(imageList)).source(source));
        Script script = Script.of(i -> i.inline(inlineScript));
        Query query = QueryBuilders.scriptScore(i -> i.script(script).query(org.springframework.data.elasticsearch.client.elc.QueryBuilders.matchAllQueryAsQuery()));
        builder.query(query);
        builder.size(100);
        SearchResponse<PatentVector> response = client.search(builder.build(), PatentVector.class);
        List<Hit<PatentVector>> hits = response.hits().hits();
        for (Hit<PatentVector> hit : hits) {
            PatentVector vector = hit.source();
            EsPictureVectorVo vectorVo = new EsPictureVectorVo();
            BeanUtil.copy(vector,vectorVo);
            list.add(vectorVo);
        }
        return list;
    }

    public List<EsPictureNoVo> getPictureByNo(EsPictureNoDTO noDTO) throws IOException {
        List<EsPictureNoVo> pictureNoVos = new ArrayList<>();
        SearchRequest.Builder builder = new SearchRequest.Builder();
        //设置查询索引
        builder.index("patent_vector");
        Query query = QueryBuilders.term(i -> i.field("app_no.keyword").value(noDTO.getAppNo()));
        builder.query(query);
        builder.size(100);
        SearchResponse<PatentVector> response = client.search(builder.build(), PatentVector.class);
        List<Hit<PatentVector>> hits = response.hits().hits();
        for (Hit<PatentVector> hit : hits) {
            PatentVector vector = hit.source();
            EsPictureNoVo noVo = new EsPictureNoVo();
            noVo.setGuid(vector.getGuid());
            noVo.setImageIndex(vector.getImageIndex());
            pictureNoVos.add(noVo);
        }
        return pictureNoVos;
    }
}