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, List<String> possible) {
        String template = "请根据以下事项列表，从句子中识别并标注出对应的社保业务事项名称。若句子中不能直接对应到任何事项，就找可能性最高的事项。请确保只识别并标注一种可能性最高的社保业务事项。" +
                "事项名称列表：\n" +
                "{}" +
                "句子是：{}";
        List<MessageDTO> messageList = new ArrayList<>();
        messageList.add(new MessageDTO("user", StrUtil.format(template, CollUtil.join(possible, ";"), question)));

        log.info("itemExtractByPossibleItem查询语句为:{}", JSONUtil.toJsonStr(messageList));
        String item = AiUtil.chatByMessage(messageList);
        log.info("itemExtractByPossibleItem结果为:{}", item);
        if (StrUtil.equals("未找到", item)) {
            throw new ItemExtractException("未从问题中找到事项");
        }
        return Collections.singletonList(item);
    }

    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("未从问题中找到业务事项");
    }


}