EAI-Backend/src/test/java/com/njuzr/eaibackend/service/AIRequestServiceTest.java

package com.njuzr.eaibackend.service;

import com.njuzr.eaibackend.exception.MyException;
import com.njuzr.eaibackend.po.AIEntry;
import lombok.extern.slf4j.Slf4j;
import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.http.HttpStatus;
import org.springframework.test.context.ActiveProfiles;

import java.util.ArrayList;
import java.util.List;

/**
 * @author: Leonezhurui
 * @Date: 2024/3/2 - 00:30
 * @Package: EAI-Backend
 */

@Slf4j
@ActiveProfiles("dev")
@SpringBootTest
public class AIRequestServiceTest {
    private final AIRequestService aiRequestService;

    @Autowired
    public AIRequestServiceTest(AIRequestService aiRequestService) {
        this.aiRequestService = aiRequestService;
    }

    @Test
    public void testChatGLM() {
        List<AIEntry> ms = new ArrayList<>();
        AIEntry a1 = new AIEntry();
        a1.setRole("user");
        a1.setContent("123");
        ms.add(a1);
        AIRequestService.AIResponse response = aiRequestService.requestChatGLM4(ms);
        System.out.println(response);
    }

    @Test
    public void testTranslate() {
        List<AIEntry> ms = new ArrayList<>();
        AIEntry a1 = new AIEntry();
        a1.setRole("user");
        String content = "任务说明：\n" +
                "\n" +
                "如果用户提供中文句子，请将该句子翻译为英文。\n" +
                "如果用户提供英文句子，请将该句子翻译为中文。\n" +
                "如果用户提供中文或英文单词，请提供该单词的所有可能翻译，并给出相关的例句。\n" +
                "示例：\n" +
                "\n" +
                "中文句子翻译：\n" +
                "用户：我今天去学校了。\n" +
                "模型：I went to school today.\n" +
                "\n" +
                "英文句子翻译：\n" +
                "用户：She is studying at the library.\n" +
                "模型：她正在图书馆学习。\n" +
                "\n" +
                "中文单词翻译与例句：\n" +
                "用户：书\n" +
                "模型：\n" +
                "\n" +
                "翻译： book, paper (某些上下文中可指书面材料)\n" +
                "例句：\n" +
                "I bought a new book yesterday. (我昨天买了一本新书。)\n" +
                "He wrote a letter on a piece of paper. (他在一张纸上写了封信。)\n" +
                "英文单词翻译与例句：\n" +
                "用户：apple\n" +
                "模型：\n" +
                "\n" +
                "翻译： 苹果 (名词), 苹果公司 (名词)\n" +
                "例句：\n" +
                "I ate an apple this morning. (我今天早上吃了一个苹果。)\n" +
                "Apple is a well-known technology company. (苹果公司是一个知名的科技公司。)\n" +
                "请确保根据用户输入的语言（中文或英文）进行适当的翻译或解释。\n" +
                "必须直接给出翻译结果，不要交互，不要任何其他信息。\n" +
                "以下是用户输入：\n\n";
        content += "回忆";
        a1.setContent(content);
        ms.add(a1);
        AIRequestService.AIResponse response = aiRequestService.requestChatGLM4(ms);
        System.out.println(response);
    }

    @Test
    void testRewrite() {
        String content = "hello world";
        String assignmentDescription = "英文读写第1次作业：软院时光机\n" +
                "穿梭时光，回到那美好而又艰苦的2002\n";
        // 3 构造提示词
        String promptTemplate =
                "假设你现在是一位英文专业的教师，正在批阅一位同学的写作作业，作业描述是：%s"
                        + "请做如下工作：\n"
                        + "1. 满分100分，请综合英文用词、句式等方面，给出评分和评分原因\n"
                        + "2. 针对文章，给出改进的内容，格式为“xxx”可以修改为“xxx”\n"
                        + "3. 给出修改后的完整文章内容\n"
                        + "其余的任何描述都不需要，也不需要任何交互！\n"
                        + "以下是学生作业内容：\n\n"
                        + "%s";

        String prompt = String.format(promptTemplate, assignmentDescription, content);
        log.info("内置提示词为：" + prompt);

        // 4 请求AI
        List<AIEntry> entry = new ArrayList<>();
        entry.add(new AIEntry("user", prompt));
        AIRequestService.AIResponse response = aiRequestService.requestChatGLM4(entry);

        // 5 获取返回数据
        String retContent;
        String role;
        Long timestamp;
        try {
            retContent = response.getChoices().get(0).getMessage().getContent();
            role = response.getChoices().get(0).getMessage().getRole();
            timestamp = response.getCreated();
            log.info("AI返回的数据为：" + retContent);
        } catch (Exception e) {
            log.error("AI请求返回的数据内容有误，数据为{}", response);
            throw new MyException(HttpStatus.INTERNAL_SERVER_ERROR.value(), HttpStatus.INTERNAL_SERVER_ERROR.getReasonPhrase() + "：" + "请求失败");
        }
    }
}