KBQA/kbqa-graph/src/main/java/com/supervision/handler/gpt/ItemExtractHandler.java

package com.supervision.handler.gpt;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.supervision.ai.AiUtil;
import com.supervision.ai.dto.MessageDTO;
import com.supervision.exception.ItemExtractException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.util.*;

/**
 * 意图提取Handler
 */
@Slf4j
@Component
public class ItemExtractHandler {

    /**
     * 直接从问题中提取事项
     *
     * @param question 问题
     * @return 提取结果
     */
    public String itemExtract(String question) {
        List<MessageDTO> messageList = new ArrayList<>();
        messageList.add(new MessageDTO("system", "假设你是一个命名实体识别模型且精通社保实体识别." +
                "现在我会给你一个句子，请根据我的要求识别出句子中的实体数据在社保业务中属于哪一类业务事项，并给我事项名称,如果未找到事项,就回复:未找到;其他什么都不要说!"));
        messageList.add(new MessageDTO("assistant", "好的"));
        messageList.add(new MessageDTO("user", "句子是:" + question + ";请识别这句话中的社保业务事项.除了识别到的业务事项,其他什么都不要说"));
        String item = AiUtil.chatByMessage(messageList);
        if (StrUtil.equals("未找到", item)) {
            throw new ItemExtractException("未从问题中找到事项");
        }
        return item;
    }

    /**
     * 提供可能的意图之后,从问题中提取事项
     *
     * @param question 问题
     * @return 提取结果
     */
    public List<String> itemExtractByPossibleItem(String question, Collection<String> possible) {
        // TODO 这里暂时没办法,暂时先当做是退休
//        List<MessageDTO> messageList = new ArrayList<>();
//        messageList.add(new MessageDTO("system", "假设你是一个命名实体识别模型且精通社保实体识别." +
//                "现在我会给你一个句子，请根据我的要求识别出句子中的实体数据在社保业务中属于哪一类业务事项，并给我事项名称,同时,我会给你一些事项可供参考,尽量从这些事项中进行匹配.如果未找到事项,就回复:未找到;其他什么都不要说!"));
//        messageList.add(new MessageDTO("assistant", "好的"));
//        messageList.add(new MessageDTO("user", StrUtil.format("可能的事项有:[{}]", CollUtil.join(possible, ";"))));
//        messageList.add(new MessageDTO("user", "句子是:" + question + ";请识别这句话中的社保业务事项,如果识别出多个,请用;分隔.除了识别到的业务事项,其他什么都不要说"));
//        log.info("itemExtractByPossibleItem查询语句为:{}", JSONUtil.toJsonStr(messageList));
//        String item = AiUtil.chatByMessage(messageList);
//        if (StrUtil.equals("未找到", item)) {
//            throw new ItemExtractException("未从问题中找到事项");
//        }
//        if (StrUtil.contains(item, ";")){
//            return Arrays.asList(item.split(";"));
//        }
//        return Collections.singletonList(item);
        return Collections.singletonList("退休");
    }


}