BJ_ZTEITS-Developers / parking-cloud-yudao

Commit 74e32f502c0273831b9a7c7ea95b5548f3ec6d5b

Authored by YunaiV
1 parent f838a4b9

test(ai): 补充聊天工具调用的同步与流式回归测试

Inline Side-by-side
Showing 1 changed file with 387 additions and 0 deletions
yudao-module-ai/src/test/java/cn/iocoder/yudao/module/ai/service/chat/AiChatMessageServiceImplTest.java 0 → 100644
  View file @74e32f5
  1 +package cn.iocoder.yudao.module.ai.service.chat;
  2 +
  3 +import cn.hutool.core.collection.CollUtil;
  4 +import cn.iocoder.yudao.framework.common.pojo.CommonResult;
  5 +import cn.iocoder.yudao.framework.security.core.LoginUser;
  6 +import cn.iocoder.yudao.framework.security.core.util.SecurityFrameworkUtils;
  7 +import cn.iocoder.yudao.framework.tenant.core.context.TenantContextHolder;
  8 +import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest;
  9 +import cn.iocoder.yudao.module.ai.controller.admin.chat.vo.message.AiChatMessageSendReqVO;
  10 +import cn.iocoder.yudao.module.ai.controller.admin.chat.vo.message.AiChatMessageSendRespVO;
  11 +import cn.iocoder.yudao.module.ai.dal.dataobject.chat.AiChatConversationDO;
  12 +import cn.iocoder.yudao.module.ai.dal.dataobject.chat.AiChatMessageDO;
  13 +import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiChatRoleDO;
  14 +import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiModelDO;
  15 +import cn.iocoder.yudao.module.ai.dal.dataobject.model.AiToolDO;
  16 +import cn.iocoder.yudao.module.ai.dal.mysql.chat.AiChatMessageMapper;
  17 +import cn.iocoder.yudao.module.ai.enums.model.AiPlatformEnum;
  18 +import cn.iocoder.yudao.module.ai.service.knowledge.AiKnowledgeDocumentService;
  19 +import cn.iocoder.yudao.module.ai.service.model.AiChatRoleService;
  20 +import cn.iocoder.yudao.module.ai.service.model.AiModelService;
  21 +import cn.iocoder.yudao.module.ai.service.model.AiToolService;
  22 +import lombok.Data;
  23 +import org.junit.jupiter.api.AfterEach;
  24 +import org.junit.jupiter.api.Test;
  25 +import org.mockito.ArgumentCaptor;
  26 +import org.mockito.InjectMocks;
  27 +import org.mockito.Mock;
  28 +import org.springframework.ai.chat.model.ToolContext;
  29 +import org.springframework.ai.deepseek.DeepSeekChatModel;
  30 +import org.springframework.ai.deepseek.api.DeepSeekApi;
  31 +import org.springframework.ai.deepseek.api.DeepSeekApi.ChatCompletion;
  32 +import org.springframework.ai.deepseek.api.DeepSeekApi.ChatCompletionChunk;
  33 +import org.springframework.ai.deepseek.api.DeepSeekApi.ChatCompletionFinishReason;
  34 +import org.springframework.ai.deepseek.api.DeepSeekApi.ChatCompletionMessage;
  35 +import org.springframework.ai.deepseek.api.DeepSeekApi.ChatCompletionMessage.ChatCompletionFunction;
  36 +import org.springframework.ai.deepseek.api.DeepSeekApi.ChatCompletionMessage.Role;
  37 +import org.springframework.ai.deepseek.api.DeepSeekApi.ChatCompletionMessage.ToolCall;
  38 +import org.springframework.ai.deepseek.api.DeepSeekApi.ChatCompletionRequest;
  39 +import org.springframework.ai.tool.ToolCallback;
  40 +import org.springframework.ai.tool.function.FunctionToolCallback;
  41 +import org.springframework.ai.tool.resolution.ToolCallbackResolver;
  42 +import org.springframework.http.ResponseEntity;
  43 +import org.springframework.mock.web.MockHttpServletRequest;
  44 +import org.springframework.security.core.context.SecurityContextHolder;
  45 +import reactor.core.publisher.Flux;
  46 +
  47 +import java.util.List;
  48 +import java.util.Map;
  49 +import java.util.concurrent.CopyOnWriteArrayList;
  50 +import java.util.concurrent.atomic.AtomicLong;
  51 +import java.util.stream.Collectors;
  52 +
  53 +import static cn.iocoder.yudao.framework.test.core.util.RandomUtils.randomLongId;
  54 +import static cn.iocoder.yudao.module.ai.enums.ErrorCodeConstants.CHAT_STREAM_ERROR;
  55 +import static cn.iocoder.yudao.module.ai.util.AiUtils.TOOL_CONTEXT_LOGIN_USER;
  56 +import static cn.iocoder.yudao.module.ai.util.AiUtils.TOOL_CONTEXT_TENANT_ID;
  57 +import static java.util.Arrays.asList;
  58 +import static java.util.Collections.singletonList;
  59 +import static org.junit.jupiter.api.Assertions.*;
  60 +import static org.mockito.ArgumentMatchers.any;
  61 +import static org.mockito.ArgumentMatchers.eq;
  62 +import static org.mockito.Mockito.*;
  63 +
  64 +/**
  65 + * {@link AiChatMessageServiceImpl} 的单元测试类
  66 + *
  67 + * 使用真实的 {@link DeepSeekChatModel} 和工具调用，仅 mock 底层的 {@link DeepSeekApi}
  68 + *
  69 + * @author 芋道源码
  70 + */
  71 +public class AiChatMessageServiceImplTest extends BaseMockitoUnitTest {
  72 +
  73 + private static final String MODEL_NAME = "deepseek-test";
  74 + private static final Double TEMPERATURE = 0.5;
  75 + private static final Integer MAX_TOKENS = 1024;
  76 +
  77 + private static final String TOOL_NAME = "query_weather";
  78 + private static final String TOOL_RESULT = "杭州：晴，25℃";
  79 + private static final ChatCompletionMessage TOOL_CALL_MESSAGE = new ChatCompletionMessage("我来查询一下杭州的天气",
  80 + Role.ASSISTANT, null, null, singletonList(new ToolCall("call_1", "function",
  81 + new ChatCompletionFunction(TOOL_NAME, "{\"city\":\"杭州\"}"))));
  82 +
  83 + private static final List<String> TOOL_ANSWER_CHUNKS = asList("杭州今天", "晴，25℃");
  84 + private static final String TOOL_ANSWER = String.join("", TOOL_ANSWER_CHUNKS);
  85 + private static final List<String> PLAIN_ANSWER_CHUNKS = asList("你好，", "我是 AI 助手");
  86 + private static final String PLAIN_ANSWER = String.join("", PLAIN_ANSWER_CHUNKS);
  87 +
  88 + /**
  89 + * assistant 消息编号：先插入 user 消息（编号 1），再插入 assistant 消息（编号 2）
  90 + */
  91 + private static final Long ASSISTANT_MESSAGE_ID = 2L;
  92 +
  93 + @InjectMocks
  94 + private AiChatMessageServiceImpl chatMessageService;
  95 +
  96 + @Mock
  97 + private AiChatMessageMapper chatMessageMapper;
  98 + @Mock
  99 + private AiChatConversationService chatConversationService;
  100 + @Mock
  101 + private AiChatRoleService chatRoleService;
  102 + @Mock
  103 + private AiModelService modalService;
  104 + @Mock
  105 + private AiKnowledgeDocumentService knowledgeDocumentService;
  106 + @Mock
  107 + private AiToolService toolService;
  108 + @Mock
  109 + private ToolCallbackResolver toolCallbackResolver;
  110 + @Mock
  111 + private DeepSeekApi deepSeekApi;
  112 +
  113 + @AfterEach
  114 + public void after() {
  115 + SecurityContextHolder.clearContext();
  116 + TenantContextHolder.clear();
  117 + }
  118 +
  119 + @Test
  120 + public void testSendMessage_toolCall() {
  121 + // 准备参数
  122 + Long userId = randomLongId();
  123 + Long roleId = randomLongId();
  124 + AiChatMessageSendReqVO reqVO = new AiChatMessageSendReqVO().setConversationId(randomLongId())
  125 + .setContent("杭州天气怎么样？");
  126 + // mock 对话、模型
  127 + mockConversation(reqVO.getConversationId(), userId, roleId);
  128 + // mock 工具
  129 + List<ToolContext> toolContexts = new CopyOnWriteArrayList<>();
  130 + mockWeatherTool(roleId, toolContexts);
  131 + // mock 登录用户、租户
  132 + Long tenantId = randomLongId();
  133 + mockLoginUser(userId, tenantId);
  134 + // mock deepSeekApi 的方法：先请求工具，再返回回答
  135 + when(deepSeekApi.chatCompletionEntity(any())).thenReturn(buildChatCompletion(TOOL_CALL_MESSAGE),
  136 + buildChatCompletion(new ChatCompletionMessage(TOOL_ANSWER, Role.ASSISTANT)));
  137 +
  138 + // 调用
  139 + AiChatMessageSendRespVO respVO = chatMessageService.sendMessage(reqVO, userId);
  140 + // 断言
  141 + assertEquals(TOOL_ANSWER, respVO.getReceive().getContent());
  142 + ArgumentCaptor<ChatCompletionRequest> requestCaptor = ArgumentCaptor.forClass(ChatCompletionRequest.class);
  143 + verify(deepSeekApi, times(2)).chatCompletionEntity(requestCaptor.capture());
  144 + assertToolExecuted(requestCaptor.getAllValues(), toolContexts, userId, tenantId);
  145 + // 断言模型参数
  146 + ChatCompletionRequest request = requestCaptor.getAllValues().get(0);
  147 + assertEquals(MODEL_NAME, request.model());
  148 + assertEquals(TEMPERATURE, request.temperature());
  149 + assertEquals(MAX_TOKENS, request.maxTokens());
  150 + // 断言消息保存
  151 + assertAssistantMessageUpdated(TOOL_ANSWER);
  152 + }
  153 +
  154 + @Test
  155 + public void testSendMessage_withoutTool() {
  156 + // 准备参数
  157 + Long userId = randomLongId();
  158 + AiChatMessageSendReqVO reqVO = new AiChatMessageSendReqVO().setConversationId(randomLongId())
  159 + .setContent("你好");
  160 + // mock 对话、模型
  161 + mockConversation(reqVO.getConversationId(), userId, null);
  162 + // mock deepSeekApi 的方法
  163 + when(deepSeekApi.chatCompletionEntity(any()))
  164 + .thenReturn(buildChatCompletion(new ChatCompletionMessage(PLAIN_ANSWER, Role.ASSISTANT)));
  165 +
  166 + // 调用
  167 + AiChatMessageSendRespVO respVO = chatMessageService.sendMessage(reqVO, userId);
  168 + // 断言
  169 + assertEquals(PLAIN_ANSWER, respVO.getReceive().getContent());
  170 + verify(deepSeekApi).chatCompletionEntity(any());
  171 + // 断言消息保存
  172 + assertAssistantMessageUpdated(PLAIN_ANSWER);
  173 + }
  174 +
  175 + @Test
  176 + public void testSendMessage_toolCallThenModelError() {
  177 + // 准备参数
  178 + Long userId = randomLongId();
  179 + Long roleId = randomLongId();
  180 + AiChatMessageSendReqVO reqVO = new AiChatMessageSendReqVO().setConversationId(randomLongId())
  181 + .setContent("杭州天气怎么样？");
  182 + // mock 对话、模型
  183 + mockConversation(reqVO.getConversationId(), userId, roleId);
  184 + // mock 工具
  185 + List<ToolContext> toolContexts = new CopyOnWriteArrayList<>();
  186 + mockWeatherTool(roleId, toolContexts);
  187 + // mock 登录用户、租户
  188 + mockLoginUser(userId, randomLongId());
  189 + // mock deepSeekApi 的方法：先请求工具，再抛出异常
  190 + RuntimeException modelException = new IllegalStateException("模型服务异常");
  191 + when(deepSeekApi.chatCompletionEntity(any())).thenReturn(buildChatCompletion(TOOL_CALL_MESSAGE))
  192 + .thenThrow(modelException);
  193 +
  194 + // 调用，并断言异常
  195 + RuntimeException exception = assertThrows(RuntimeException.class,
  196 + () -> chatMessageService.sendMessage(reqVO, userId));
  197 + assertSame(modelException, exception);
  198 + // 断言工具已执行、消息未保存
  199 + verify(deepSeekApi, times(2)).chatCompletionEntity(any());
  200 + assertEquals(1, toolContexts.size());
  201 + verify(chatMessageMapper, never()).updateById(any(AiChatMessageDO.class));
  202 + }
  203 +
  204 + @Test
  205 + public void testSendChatMessageStream_toolCall() {
  206 + // 准备参数
  207 + Long userId = randomLongId();
  208 + Long roleId = randomLongId();
  209 + AiChatMessageSendReqVO reqVO = new AiChatMessageSendReqVO().setConversationId(randomLongId())
  210 + .setContent("杭州天气怎么样？");
  211 + // mock 对话、模型
  212 + mockConversation(reqVO.getConversationId(), userId, roleId);
  213 + // mock 工具
  214 + List<ToolContext> toolContexts = new CopyOnWriteArrayList<>();
  215 + mockWeatherTool(roleId, toolContexts);
  216 + // mock 登录用户、租户
  217 + Long tenantId = randomLongId();
  218 + mockLoginUser(userId, tenantId);
  219 + // mock deepSeekApi 的方法：先请求工具，再流式返回回答
  220 + when(deepSeekApi.chatCompletionStream(any())).thenReturn(Flux.just(buildChatCompletionChunk(TOOL_CALL_MESSAGE)),
  221 + buildAnswerChunks(TOOL_ANSWER_CHUNKS));
  222 +
  223 + // 调用
  224 + List<CommonResult<AiChatMessageSendRespVO>> results = chatMessageService
  225 + .sendChatMessageStream(reqVO, userId).collectList().block();
  226 + // 断言（工具调用轮次不下发）
  227 + assertNotNull(results);
  228 + assertTrue(results.stream().allMatch(CommonResult::isSuccess));
  229 + assertEquals(TOOL_ANSWER, joinContent(results));
  230 + ArgumentCaptor<ChatCompletionRequest> requestCaptor = ArgumentCaptor.forClass(ChatCompletionRequest.class);
  231 + verify(deepSeekApi, times(2)).chatCompletionStream(requestCaptor.capture());
  232 + assertToolExecuted(requestCaptor.getAllValues(), toolContexts, userId, tenantId);
  233 + // 断言消息保存
  234 + assertAssistantMessageUpdated(TOOL_ANSWER);
  235 + }
  236 +
  237 + @Test
  238 + public void testSendChatMessageStream_withoutTool() {
  239 + // 准备参数
  240 + Long userId = randomLongId();
  241 + AiChatMessageSendReqVO reqVO = new AiChatMessageSendReqVO().setConversationId(randomLongId())
  242 + .setContent("你好");
  243 + // mock 对话、模型
  244 + mockConversation(reqVO.getConversationId(), userId, null);
  245 + // mock deepSeekApi 的方法
  246 + when(deepSeekApi.chatCompletionStream(any())).thenReturn(buildAnswerChunks(PLAIN_ANSWER_CHUNKS));
  247 +
  248 + // 调用
  249 + List<CommonResult<AiChatMessageSendRespVO>> results = chatMessageService
  250 + .sendChatMessageStream(reqVO, userId).collectList().block();
  251 + // 断言
  252 + assertNotNull(results);
  253 + assertTrue(results.stream().allMatch(CommonResult::isSuccess));
  254 + assertEquals(PLAIN_ANSWER, joinContent(results));
  255 + verify(deepSeekApi).chatCompletionStream(any());
  256 + // 断言消息保存
  257 + assertAssistantMessageUpdated(PLAIN_ANSWER);
  258 + }
  259 +
  260 + @Test
  261 + public void testSendChatMessageStream_toolCallThenModelError() {
  262 + // 准备参数
  263 + Long userId = randomLongId();
  264 + Long roleId = randomLongId();
  265 + AiChatMessageSendReqVO reqVO = new AiChatMessageSendReqVO().setConversationId(randomLongId())
  266 + .setContent("杭州天气怎么样？");
  267 + // mock 对话、模型
  268 + mockConversation(reqVO.getConversationId(), userId, roleId);
  269 + // mock 工具
  270 + List<ToolContext> toolContexts = new CopyOnWriteArrayList<>();
  271 + mockWeatherTool(roleId, toolContexts);
  272 + // mock 登录用户、租户
  273 + mockLoginUser(userId, randomLongId());
  274 + // mock deepSeekApi 的方法：先请求工具，再返回异常
  275 + when(deepSeekApi.chatCompletionStream(any())).thenReturn(Flux.just(buildChatCompletionChunk(TOOL_CALL_MESSAGE)),
  276 + Flux.error(new IllegalStateException("模型服务异常")));
  277 +
  278 + // 调用
  279 + List<CommonResult<AiChatMessageSendRespVO>> results = chatMessageService
  280 + .sendChatMessageStream(reqVO, userId).collectList().block();
  281 + // 断言
  282 + assertNotNull(results);
  283 + assertEquals(CHAT_STREAM_ERROR.getCode(), CollUtil.getLast(results).getCode());
  284 + verify(deepSeekApi, times(2)).chatCompletionStream(any());
  285 + assertEquals(1, toolContexts.size());
  286 + // 断言消息删除（无内容）
  287 + verify(chatMessageMapper).deleteById(ASSISTANT_MESSAGE_ID);
  288 + verify(chatMessageMapper, never()).updateById(any(AiChatMessageDO.class));
  289 + }
  290 +
  291 + private void mockConversation(Long conversationId, Long userId, Long roleId) {
  292 + // mock 对话
  293 + AiChatConversationDO conversation = new AiChatConversationDO().setId(conversationId).setUserId(userId)
  294 + .setRoleId(roleId).setModelId(randomLongId()).setTemperature(TEMPERATURE).setMaxTokens(MAX_TOKENS);
  295 + when(chatConversationService.validateChatConversationExists(eq(conversationId))).thenReturn(conversation);
  296 + // mock 模型：真实的 DeepSeekChatModel
  297 + AiModelDO model = new AiModelDO().setId(conversation.getModelId()).setModel(MODEL_NAME)
  298 + .setPlatform(AiPlatformEnum.DEEP_SEEK.getPlatform());
  299 + when(modalService.validateModel(eq(model.getId()))).thenReturn(model);
  300 + when(modalService.getChatModel(eq(model.getId())))
  301 + .thenReturn(DeepSeekChatModel.builder().deepSeekApi(deepSeekApi).build());
  302 + // mock 消息插入：依次分配编号
  303 + AtomicLong messageId = new AtomicLong();
  304 + doAnswer(invocation -> {
  305 + invocation.<AiChatMessageDO>getArgument(0).setId(messageId.incrementAndGet());
  306 + return 1;
  307 + }).when(chatMessageMapper).insert(any(AiChatMessageDO.class));
  308 + }
  309 +
  310 + private void mockWeatherTool(Long roleId, List<ToolContext> toolContexts) {
  311 + // mock 工具：真实的 FunctionToolCallback
  312 + Long toolId = randomLongId();
  313 + ToolCallback toolCallback = FunctionToolCallback.<WeatherRequest, String>builder(TOOL_NAME,
  314 + (request, toolContext) -> {
  315 + toolContexts.add(toolContext);
  316 + return request.getCity() + "：晴，25℃";
  317 + }).description("查询城市天气").inputType(WeatherRequest.class).build();
  318 + when(chatRoleService.getChatRole(eq(roleId))).thenReturn(new AiChatRoleDO().setId(roleId)
  319 + .setToolIds(singletonList(toolId)));
  320 + when(toolService.getToolList(eq(singletonList(toolId))))
  321 + .thenReturn(singletonList(new AiToolDO().setId(toolId).setName(TOOL_NAME)));
  322 + when(toolCallbackResolver.resolve(eq(TOOL_NAME))).thenReturn(toolCallback);
  323 + }
  324 +
  325 + private static void mockLoginUser(Long userId, Long tenantId) {
  326 + SecurityFrameworkUtils.setLoginUser(new LoginUser().setId(userId), new MockHttpServletRequest());
  327 + TenantContextHolder.setTenantId(tenantId);
  328 + }
  329 +
  330 + private static void assertToolExecuted(List<ChatCompletionRequest> requests, List<ToolContext> toolContexts,
  331 + Long userId, Long tenantId) {
  332 + // 断言首轮请求带上工具
  333 + assertEquals(1, requests.get(0).tools().size());
  334 + // 断言工具执行，并拿到登录用户、租户
  335 + assertEquals(1, toolContexts.size());
  336 + Map<String, Object> context = toolContexts.get(0).getContext();
  337 + assertEquals(userId, ((LoginUser) context.get(TOOL_CONTEXT_LOGIN_USER)).getId());
  338 + assertEquals(tenantId, context.get(TOOL_CONTEXT_TENANT_ID));
  339 + // 断言次轮请求回传工具结果
  340 + ChatCompletionMessage toolMessage = CollUtil.getLast(requests.get(1).messages());
  341 + assertEquals(Role.TOOL, toolMessage.role());
  342 + assertTrue(toolMessage.content().contains(TOOL_RESULT));
  343 + }
  344 +
  345 + private void assertAssistantMessageUpdated(String content) {
  346 + ArgumentCaptor<AiChatMessageDO> messageCaptor = ArgumentCaptor.forClass(AiChatMessageDO.class);
  347 + verify(chatMessageMapper).updateById(messageCaptor.capture());
  348 + assertEquals(ASSISTANT_MESSAGE_ID, messageCaptor.getValue().getId());
  349 + assertEquals(content, messageCaptor.getValue().getContent());
  350 + }
  351 +
  352 + private static String joinContent(List<CommonResult<AiChatMessageSendRespVO>> results) {
  353 + return results.stream().map(result -> result.getData().getReceive().getContent())
  354 + .collect(Collectors.joining());
  355 + }
  356 +
  357 + private static ResponseEntity<ChatCompletion> buildChatCompletion(ChatCompletionMessage message) {
  358 + ChatCompletion.Choice choice = new ChatCompletion.Choice(buildFinishReason(message), 0, message, null);
  359 + return ResponseEntity.ok(new ChatCompletion("chatcmpl-1", singletonList(choice), null, MODEL_NAME,
  360 + null, null, null));
  361 + }
  362 +
  363 + private static Flux<ChatCompletionChunk> buildAnswerChunks(List<String> contents) {
  364 + return Flux.fromIterable(contents)
  365 + .map(content -> buildChatCompletionChunk(new ChatCompletionMessage(content, Role.ASSISTANT)));
  366 + }
  367 +
  368 + private static ChatCompletionChunk buildChatCompletionChunk(ChatCompletionMessage delta) {
  369 + ChatCompletionChunk.ChunkChoice choice = new ChatCompletionChunk.ChunkChoice(buildFinishReason(delta),
  370 + 0, delta, null);
  371 + return new ChatCompletionChunk("chatcmpl-1", singletonList(choice), null, MODEL_NAME,
  372 + null, null, null, null);
  373 + }
  374 +
  375 + private static ChatCompletionFinishReason buildFinishReason(ChatCompletionMessage message) {
  376 + return CollUtil.isNotEmpty(message.toolCalls()) ? ChatCompletionFinishReason.TOOL_CALLS
  377 + : ChatCompletionFinishReason.STOP;
  378 + }
  379 +
  380 + @Data
  381 + public static class WeatherRequest {
  382 +
  383 + private String city;
  384 +
  385 + }
  386 +
  387 +}