long-article-recommend/long-article-recommend-service/src/test/java/com/tzld/longarticle/recommend/server/ArticleVideoAuditTest.java

package com.tzld.longarticle.recommend.server;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.google.common.collect.Lists;
import com.tzld.longarticle.recommend.server.common.enums.recommend.ArticleCategoryStatusEnum;
import com.tzld.longarticle.recommend.server.mapper.longArticle.ArticleAuditMapper;
import com.tzld.longarticle.recommend.server.model.dto.kimi.KimiResult;
import com.tzld.longarticle.recommend.server.model.entity.aigc.ProducePlan;
import com.tzld.longarticle.recommend.server.model.entity.aigc.ProducePlanExeRecord;
import com.tzld.longarticle.recommend.server.model.entity.longArticle.ArticleCategory;
import com.tzld.longarticle.recommend.server.model.entity.longArticle.LongArticleTitleAudit;
import com.tzld.longarticle.recommend.server.remote.DeepSeekApiService;
import com.tzld.longarticle.recommend.server.repository.aigc.ProducePlanExeRecordRepository;
import com.tzld.longarticle.recommend.server.repository.aigc.ProducePlanRepository;
import com.tzld.longarticle.recommend.server.repository.longArticle.ArticleCategoryRepository;
import com.tzld.longarticle.recommend.server.repository.longArticle.LongArticleTitleAuditRepository;
import com.tzld.longarticle.recommend.server.service.recommend.ArticleCategoryService;
import com.tzld.longarticle.recommend.server.util.TitleSimilarCheckUtil;
import lombok.extern.slf4j.Slf4j;
import org.apache.poi.ss.usermodel.Cell;
import org.apache.poi.ss.usermodel.Row;
import org.apache.poi.ss.usermodel.Sheet;
import org.apache.poi.ss.usermodel.Workbook;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.util.StringUtils;

import javax.annotation.Resource;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.stream.Collectors;

@SpringBootTest(classes = Application.class)
@Slf4j
public class ArticleVideoAuditTest {

    @Resource
    private ArticleAuditMapper articleAuditMapper;
    @Resource
    private LongArticleTitleAuditRepository titleAuditRepository;
    @Resource
    private ProducePlanExeRecordRepository exeRecordRepository;
    @Resource
    private ProducePlanRepository producePlanRepository;
    @Resource
    private ArticleCategoryService articleCategoryService;
    @Autowired
    ArticleCategoryRepository articleCategoryRepository;
    @Autowired
    DeepSeekApiService deepSeekApiService;

    @Value("${kimiCategoryPrompt:}")
    private String kimiCategoryPrompt;

    @Test
    public void updateTitleAuditFlowPoolLevel() {
        List<LongArticleTitleAudit> list = titleAuditRepository.getByFlowPoolLevelIsNull();
        List<String> contentIds = list.stream().map(LongArticleTitleAudit::getContentId).collect(Collectors.toList());
        List<ProducePlanExeRecord> contentList = exeRecordRepository.findByPlanExeIdIn(contentIds);
        Map<String, String> contentPlanMap = contentList.stream().collect(Collectors.toMap(ProducePlanExeRecord::getPlanExeId, ProducePlanExeRecord::getPlanId));
        List<String> planIds = contentList.stream().map(ProducePlanExeRecord::getPlanId).collect(Collectors.toList());
        List<ProducePlan> planList = producePlanRepository.findByIdIn(planIds);
        Map<String, String> planMap = planList.stream().collect(Collectors.toMap(ProducePlan::getId, ProducePlan::getPlanTag));
        for (LongArticleTitleAudit titleAudit : list) {
            String contentId = titleAudit.getContentId();
            String planId = contentPlanMap.get(contentId);
            String tag = planMap.get(planId);
            if (!StringUtils.hasText(tag)) {
                continue;
            }
            titleAudit.setFlowPoolLevel(tag);
            titleAuditRepository.save(titleAudit);
        }
        articleAuditMapper.updateTitleAuditFlowPoolLevel();
    }

    @Test
    public void ArticleCategoryTest() {
        List<ArticleCategory> dealList = articleCategoryRepository.getByStatus(
                ArticleCategoryStatusEnum.SUCCESS.getCode());
        dealList = dealList.subList(0, 200);
        List<List<ArticleCategory>> partitionList = Lists.partition(dealList, 50);
        JSONArray result = new JSONArray();
        for (List<ArticleCategory> partition : partitionList) {
            List<String> partitionTitles = partition.stream().map(ArticleCategory::getTitle).distinct().collect(Collectors.toList());
            String prompt = buildKimiPrompt(partitionTitles);
            KimiResult kimiResult = deepSeekApiService.requestOfficialApi(prompt, null, null, true);
            long now = System.currentTimeMillis();
            JSONObject obj = null;
            if (kimiResult.isSuccess()) {
                try {
                    obj = JSONObject.parseObject(kimiResult.getResponse().getChoices().get(0).getMessage().getContent());
                } catch (Exception e) {
                    log.error(kimiResult.getResponse().getChoices().get(0).getMessage().getContent());
                }
            }
            for (ArticleCategory articleCategory : partition) {
                JSONObject res = new JSONObject();
                res.put("title", articleCategory.getTitle());
                res.put("category", articleCategory.getCategory());
                articleCategory.setKimiResult(kimiResult.getResponseStr());
                articleCategory.setUpdateTimestamp(now);
                if (kimiResult.isSuccess() && Objects.nonNull(obj) && obj.containsKey(articleCategory.getTitle())) {
                    articleCategory.setCategory(obj.getString(articleCategory.getTitle()));
                    articleCategory.setStatus(ArticleCategoryStatusEnum.SUCCESS.getCode());
                } else {
                    if (Objects.nonNull(obj)) {
                        for (Map.Entry<String, Object> entry : obj.entrySet()) {
                            if (TitleSimilarCheckUtil.isSimilar(articleCategory.getTitle(), entry.getKey(),
                                    TitleSimilarCheckUtil.ARTICLE_PROMOTION_THRESHOLD)) {
                                articleCategory.setCategory(obj.getString(entry.getKey()));
                                articleCategory.setStatus(ArticleCategoryStatusEnum.SUCCESS.getCode());
                                break;
                            }
                        }
                    }
                }
                if (articleCategory.getStatus().equals(ArticleCategoryStatusEnum.SUCCESS.getCode())) {
                    res.put("newCategory", articleCategory.getCategory());
                }
                result.add(res);
            }
        }

        Workbook workbook = new XSSFWorkbook();
        Sheet sheet = workbook.createSheet("ExampleSheet");
        int rowNum = 0;
        // 创建标题行
        Row titleRow = sheet.createRow(rowNum);
        Cell titleCell = titleRow.createCell(0);
        titleCell.setCellValue("标题");
        titleCell = titleRow.createCell(1);
        titleCell.setCellValue("kimi品类");
        titleCell = titleRow.createCell(2);
        titleCell.setCellValue("deepseek品类");

        for (Object obj : result) {
            JSONObject jsonObject = (JSONObject) obj;
            String title = jsonObject.getString("title");
            String category = jsonObject.getString("category");
            String newCategory = jsonObject.getString("newCategory");
            rowNum++;
            Row row = sheet.createRow(rowNum);
            Cell cell = row.createCell(0);
            cell.setCellValue(title);
            cell = row.createCell(1);
            cell.setCellValue(category);
            cell = row.createCell(2);
            cell.setCellValue(newCategory);
        }

        try (FileOutputStream outputStream = new FileOutputStream("/Users/wangyunpeng/Downloads/category.xlsx")) {
            workbook.write(outputStream);
        } catch (IOException e) {
            e.printStackTrace();
        } finally {
            try {
                workbook.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    private String buildKimiPrompt(List<String> titleList) {
        StringBuilder prompt = new StringBuilder(kimiCategoryPrompt);
        prompt.append("\n");
        for (String title : titleList) {
            prompt.append(title).append("\n");
        }
        return prompt.toString();
    }
}