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("退休");
    }

    public String itemExtractBusiness(String question) {
        List<MessageDTO> messageList = new ArrayList<>();
        messageList.add(new MessageDTO("user", "我现在是要进行实体抽取任务，并且精通社保业务，现在需要抽取[业务]这个实体类型的内容。\n" +
                "\n" +
                "下面是一些示例：\n" +
                "\n" +
                "输入：女性工人一般多少岁可以退休？\n" +
                "输出：{\"business\":\"退休\"}\n" +
                "输入：办理退休一般需要什么条件？\n" +
                "输出：{\"business\":\"退休\"}\n" +
                "输入：身份证挂失怎么办？\n" +
                "输出：{\"business\":\"身份证挂失\"}\n" +
                "输入：退休金没发，是什么原因？\n" +
                "输出：{\"business\":\"退休金发放\"}\n" +
                "\n" +
                "现在有一句话，请进行抽取。\n" +
                "输入：" + question + "\n" +
                "输出："));
        log.info("itemExtractBusiness查询语句为:{}", JSONUtil.toJsonStr(messageList));
        String item = AiUtil.chatByMessage(messageList);
        log.info("itemExtractBusiness查询结果为:{}", item);
        boolean typeJSON = JSONUtil.isTypeJSON(item);
        if (typeJSON) {
            String business = JSONUtil.parseObj(item).getStr("business");
            if (StrUtil.isNotBlank(business)) {
                return business;
            }
        }
        throw new ItemExtractException("未从问题中找到业务事项");
    }


}