| | |
| | | import com.zy.ai.domain.autotune.AutoTuneApplyResult; |
| | | import com.zy.ai.domain.autotune.AutoTuneChangeCommand; |
| | | import com.zy.ai.domain.autotune.AutoTuneSnapshot; |
| | | import com.zy.ai.domain.autotune.AutoTuneJobStatus; |
| | | import com.zy.ai.domain.autotune.AutoTuneTriggerType; |
| | | import com.zy.ai.entity.AiAutoTuneChange; |
| | | import com.zy.ai.entity.AiAutoTuneJob; |
| | | import com.zy.ai.entity.AiPromptTemplate; |
| | | import com.zy.ai.entity.ChatCompletionRequest; |
| | | import com.zy.ai.entity.ChatCompletionResponse; |
| | | import com.zy.ai.enums.AiPromptScene; |
| | | import com.zy.ai.mcp.service.SpringAiMcpToolManager; |
| | | import com.zy.ai.mcp.tool.AutoTuneMcpTools; |
| | | import com.zy.ai.service.impl.AutoTuneAgentServiceImpl; |
| | | import com.zy.ai.service.impl.AutoTuneCoordinatorServiceImpl; |
| | | import com.zy.asrs.service.WrkMastService; |
| | | import com.zy.common.utils.RedisUtil; |
| | | import com.zy.core.enums.RedisKeyType; |
| | | import com.zy.system.entity.OperateLog; |
| | | import com.zy.system.service.ConfigService; |
| | | import com.zy.system.service.OperateLogService; |
| | | import org.junit.jupiter.api.BeforeEach; |
| | | import org.junit.jupiter.api.Test; |
| | | import org.junit.jupiter.api.extension.ExtendWith; |
| | | import org.mockito.ArgumentCaptor; |
| | | import org.mockito.Mock; |
| | | import org.mockito.junit.jupiter.MockitoExtension; |
| | | import org.springframework.test.util.ReflectionTestUtils; |
| | | |
| | | import java.util.ArrayList; |
| | | import java.util.Collections; |
| | | import java.util.Date; |
| | | import java.util.LinkedHashMap; |
| | | import java.util.List; |
| | | import java.util.Map; |
| | | import java.util.concurrent.atomic.AtomicLong; |
| | | import java.util.function.LongSupplier; |
| | | |
| | | import static org.junit.jupiter.api.Assertions.assertEquals; |
| | | import static org.junit.jupiter.api.Assertions.assertFalse; |
| | |
| | | import static org.mockito.ArgumentMatchers.any; |
| | | import static org.mockito.ArgumentMatchers.anyDouble; |
| | | import static org.mockito.ArgumentMatchers.anyInt; |
| | | import static org.mockito.ArgumentMatchers.anyLong; |
| | | import static org.mockito.ArgumentMatchers.anyString; |
| | | import static org.mockito.ArgumentMatchers.eq; |
| | | import static org.mockito.Mockito.doThrow; |
| | | import static org.mockito.Mockito.never; |
| | | import static org.mockito.Mockito.times; |
| | | import static org.mockito.Mockito.verify; |
| | |
| | | private SpringAiMcpToolManager mcpToolManager; |
| | | @Mock |
| | | private AiPromptTemplateService aiPromptTemplateService; |
| | | @Mock |
| | | private ConfigService configService; |
| | | @Mock |
| | | private WrkMastService wrkMastService; |
| | | @Mock |
| | | private AutoTuneAgentService autoTuneAgentService; |
| | | @Mock |
| | | private RedisUtil redisUtil; |
| | | @Mock |
| | | private OperateLogService operateLogService; |
| | | |
| | | @BeforeEach |
| | | void setUp() { |
| | |
| | | } |
| | | |
| | | @Test |
| | | void applyToolRejectsExpiredDryRunToken() { |
| | | AtomicLong currentTimeMillis = new AtomicLong(1_000L); |
| | | ReflectionTestUtils.invokeMethod(tools, "setCurrentTimeMillisSupplier", (LongSupplier) currentTimeMillis::get); |
| | | AutoTuneApplyResult dryRunResult = new AutoTuneApplyResult(); |
| | | dryRunResult.setDryRun(true); |
| | | dryRunResult.setSuccess(true); |
| | | when(autoTuneApplyService.apply(any(AutoTuneApplyRequest.class))).thenReturn(dryRunResult); |
| | | List<AutoTuneChangeCommand> changes = Collections.singletonList( |
| | | change("sys_config", null, "conveyorStationTaskLimit", "12")); |
| | | |
| | | AutoTuneApplyResult preview = tools.applyAutoTuneChanges("preview", 10, "agent", true, null, changes); |
| | | currentTimeMillis.addAndGet(10L * 60L * 1000L + 1L); |
| | | IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, |
| | | () -> tools.applyAutoTuneChanges("apply", 10, "agent", false, preview.getDryRunToken(), changes)); |
| | | |
| | | assertTrue(exception.getMessage().contains("expired")); |
| | | verify(autoTuneApplyService, times(1)).apply(any(AutoTuneApplyRequest.class)); |
| | | } |
| | | |
| | | @Test |
| | | void rollbackToolDelegatesToApplyServiceRollback() { |
| | | AutoTuneApplyResult expected = new AutoTuneApplyResult(); |
| | | when(autoTuneApplyService.rollbackLastSuccessfulJob("bad result")).thenReturn(expected); |
| | |
| | | |
| | | assertSame(expected, result); |
| | | verify(autoTuneApplyService).rollbackLastSuccessfulJob("bad result"); |
| | | } |
| | | |
| | | @Test |
| | | void coordinatorSkipsWhenDisabled() { |
| | | when(configService.getConfigValue("aiAutoTuneEnabled", "N")).thenReturn("N"); |
| | | |
| | | AutoTuneCoordinatorService.AutoTuneCoordinatorResult result = coordinatorService().runAutoTuneIfEligible(); |
| | | |
| | | assertTrue(result.getSkipped()); |
| | | assertEquals("disabled", result.getReason()); |
| | | verify(wrkMastService, never()).count(any(Wrapper.class)); |
| | | verify(autoTuneAgentService, never()).runAutoTune(anyString()); |
| | | } |
| | | |
| | | @Test |
| | | void coordinatorSkipsWhenNoActiveTasks() { |
| | | when(configService.getConfigValue("aiAutoTuneEnabled", "N")).thenReturn("Y"); |
| | | when(configService.getConfigValue("aiAutoTuneIntervalMinutes", "10")).thenReturn("10"); |
| | | when(wrkMastService.count(any(Wrapper.class))).thenReturn(0L); |
| | | |
| | | AutoTuneCoordinatorService.AutoTuneCoordinatorResult result = coordinatorService().runAutoTuneIfEligible(); |
| | | |
| | | assertTrue(result.getSkipped()); |
| | | assertEquals("no_active_tasks", result.getReason()); |
| | | verify(autoTuneAgentService, never()).runAutoTune(anyString()); |
| | | } |
| | | |
| | | @Test |
| | | void coordinatorSkipsWhenIntervalNotReached() { |
| | | AiAutoTuneJob recentJob = new AiAutoTuneJob(); |
| | | recentJob.setId(11L); |
| | | recentJob.setFinishTime(new Date()); |
| | | when(configService.getConfigValue("aiAutoTuneEnabled", "N")).thenReturn("true"); |
| | | when(configService.getConfigValue("aiAutoTuneIntervalMinutes", "10")).thenReturn("10"); |
| | | when(wrkMastService.count(any(Wrapper.class))).thenReturn(1L); |
| | | when(redisUtil.get(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key)).thenReturn(null); |
| | | when(aiAutoTuneJobService.list(any(Wrapper.class))).thenReturn(Collections.singletonList(recentJob)); |
| | | |
| | | AutoTuneCoordinatorService.AutoTuneCoordinatorResult result = coordinatorService().runAutoTuneIfEligible(); |
| | | |
| | | assertTrue(result.getSkipped()); |
| | | assertEquals("interval_not_reached", result.getReason()); |
| | | verify(autoTuneAgentService, never()).runAutoTune(anyString()); |
| | | } |
| | | |
| | | @Test |
| | | void coordinatorTriggersAgentWhenEligible() { |
| | | AutoTuneAgentService.AutoTuneAgentResult agentResult = successfulAgentResult(); |
| | | when(configService.getConfigValue("aiAutoTuneEnabled", "N")).thenReturn("1"); |
| | | when(configService.getConfigValue("aiAutoTuneIntervalMinutes", "10")).thenReturn("10"); |
| | | when(wrkMastService.count(any(Wrapper.class))).thenReturn(1L); |
| | | when(redisUtil.get(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key)).thenReturn(null); |
| | | when(aiAutoTuneJobService.list(any(Wrapper.class))).thenReturn(Collections.emptyList()); |
| | | when(redisUtil.trySetStringIfAbsent(anyString(), anyString(), anyLong())).thenReturn(true); |
| | | when(autoTuneAgentService.runAutoTune(AutoTuneTriggerType.AUTO.getCode())).thenReturn(agentResult); |
| | | |
| | | AutoTuneCoordinatorService.AutoTuneCoordinatorResult result = coordinatorService().runAutoTuneIfEligible(); |
| | | |
| | | assertFalse(result.getSkipped()); |
| | | assertTrue(result.getTriggered()); |
| | | assertSame(agentResult, result.getAgentResult()); |
| | | verify(autoTuneAgentService).runAutoTune(AutoTuneTriggerType.AUTO.getCode()); |
| | | verify(operateLogService).save(any()); |
| | | verify(redisUtil).set(anyString(), any(), anyLong()); |
| | | verify(redisUtil).compareAndDelete(anyString(), anyString()); |
| | | } |
| | | |
| | | @Test |
| | | void coordinatorKeepsAgentResultWhenOperateLogFails() { |
| | | AutoTuneAgentService.AutoTuneAgentResult agentResult = successfulAgentResult(); |
| | | when(configService.getConfigValue("aiAutoTuneEnabled", "N")).thenReturn("Y"); |
| | | when(configService.getConfigValue("aiAutoTuneIntervalMinutes", "10")).thenReturn("10"); |
| | | when(wrkMastService.count(any(Wrapper.class))).thenReturn(1L); |
| | | when(redisUtil.get(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key)).thenReturn(null); |
| | | when(aiAutoTuneJobService.list(any(Wrapper.class))).thenReturn(Collections.emptyList()); |
| | | when(redisUtil.trySetStringIfAbsent(anyString(), anyString(), anyLong())).thenReturn(true); |
| | | when(autoTuneAgentService.runAutoTune(AutoTuneTriggerType.AUTO.getCode())).thenReturn(agentResult); |
| | | doThrow(new RuntimeException("log failed")).when(operateLogService).save(any()); |
| | | |
| | | AutoTuneCoordinatorService.AutoTuneCoordinatorResult result = coordinatorService().runAutoTuneIfEligible(); |
| | | |
| | | assertFalse(result.getSkipped()); |
| | | assertTrue(result.getTriggered()); |
| | | assertSame(agentResult, result.getAgentResult()); |
| | | verify(redisUtil).compareAndDelete(anyString(), anyString()); |
| | | } |
| | | |
| | | @Test |
| | | void coordinatorRunsAgentAndReleasesLockWhenGuardWriteFails() { |
| | | AutoTuneAgentService.AutoTuneAgentResult agentResult = successfulAgentResult(); |
| | | when(configService.getConfigValue("aiAutoTuneEnabled", "N")).thenReturn("Y"); |
| | | when(configService.getConfigValue("aiAutoTuneIntervalMinutes", "10")).thenReturn("10"); |
| | | when(wrkMastService.count(any(Wrapper.class))).thenReturn(1L); |
| | | when(redisUtil.get(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key)).thenReturn(null); |
| | | when(aiAutoTuneJobService.list(any(Wrapper.class))).thenReturn(Collections.emptyList()); |
| | | when(redisUtil.trySetStringIfAbsent(anyString(), anyString(), anyLong())).thenReturn(true); |
| | | doThrow(new RuntimeException("guard failed")) |
| | | .when(redisUtil) |
| | | .set(eq(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key), any(), eq(600L)); |
| | | when(autoTuneAgentService.runAutoTune(AutoTuneTriggerType.AUTO.getCode())).thenReturn(agentResult); |
| | | |
| | | AutoTuneCoordinatorService.AutoTuneCoordinatorResult result = coordinatorService().runAutoTuneIfEligible(); |
| | | |
| | | assertFalse(result.getSkipped()); |
| | | assertTrue(result.getTriggered()); |
| | | assertSame(agentResult, result.getAgentResult()); |
| | | verify(autoTuneAgentService).runAutoTune(AutoTuneTriggerType.AUTO.getCode()); |
| | | verify(redisUtil).compareAndDelete(anyString(), anyString()); |
| | | } |
| | | |
| | | @Test |
| | | void coordinatorSetsGuardWhenAgentReturnsFailure() { |
| | | AutoTuneAgentService.AutoTuneAgentResult agentResult = failedAgentResult(); |
| | | when(configService.getConfigValue("aiAutoTuneEnabled", "N")).thenReturn("Y"); |
| | | when(configService.getConfigValue("aiAutoTuneIntervalMinutes", "10")).thenReturn("10"); |
| | | when(wrkMastService.count(any(Wrapper.class))).thenReturn(1L); |
| | | when(redisUtil.get(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key)).thenReturn(null); |
| | | when(aiAutoTuneJobService.list(any(Wrapper.class))).thenReturn(Collections.emptyList()); |
| | | when(redisUtil.trySetStringIfAbsent(anyString(), anyString(), anyLong())).thenReturn(true); |
| | | when(autoTuneAgentService.runAutoTune(AutoTuneTriggerType.AUTO.getCode())).thenReturn(agentResult); |
| | | |
| | | AutoTuneCoordinatorService.AutoTuneCoordinatorResult result = coordinatorService().runAutoTuneIfEligible(); |
| | | |
| | | assertFalse(result.getSkipped()); |
| | | assertSame(agentResult, result.getAgentResult()); |
| | | verify(redisUtil).set(eq(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key), any(), eq(600L)); |
| | | } |
| | | |
| | | @Test |
| | | void coordinatorSetsGuardWhenAgentThrows() { |
| | | when(configService.getConfigValue("aiAutoTuneEnabled", "N")).thenReturn("Y"); |
| | | when(configService.getConfigValue("aiAutoTuneIntervalMinutes", "10")).thenReturn("10"); |
| | | when(wrkMastService.count(any(Wrapper.class))).thenReturn(1L); |
| | | when(redisUtil.get(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key)).thenReturn(null); |
| | | when(aiAutoTuneJobService.list(any(Wrapper.class))).thenReturn(Collections.emptyList()); |
| | | when(redisUtil.trySetStringIfAbsent(anyString(), anyString(), anyLong())).thenReturn(true); |
| | | when(autoTuneAgentService.runAutoTune(AutoTuneTriggerType.AUTO.getCode())).thenThrow(new RuntimeException("agent failed")); |
| | | |
| | | AutoTuneCoordinatorService.AutoTuneCoordinatorResult result = coordinatorService().runAutoTuneIfEligible(); |
| | | |
| | | assertFalse(result.getSkipped()); |
| | | assertFalse(result.getAgentResult().getSuccess()); |
| | | verify(redisUtil).set(eq(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key), any(), eq(600L)); |
| | | verify(redisUtil).compareAndDelete(anyString(), anyString()); |
| | | } |
| | | |
| | | @Test |
| | | void coordinatorSkipsWhenRunningLockIsNotAcquired() { |
| | | when(configService.getConfigValue("aiAutoTuneEnabled", "N")).thenReturn("Y"); |
| | | when(configService.getConfigValue("aiAutoTuneIntervalMinutes", "10")).thenReturn("10"); |
| | | when(wrkMastService.count(any(Wrapper.class))).thenReturn(1L); |
| | | when(redisUtil.get(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key)).thenReturn(null); |
| | | when(aiAutoTuneJobService.list(any(Wrapper.class))).thenReturn(Collections.emptyList()); |
| | | when(redisUtil.trySetStringIfAbsent(anyString(), anyString(), anyLong())).thenReturn(false); |
| | | |
| | | AutoTuneCoordinatorService.AutoTuneCoordinatorResult result = coordinatorService().runAutoTuneIfEligible(); |
| | | |
| | | assertTrue(result.getSkipped()); |
| | | assertEquals("running_lock_not_acquired", result.getReason()); |
| | | verify(autoTuneAgentService, never()).runAutoTune(anyString()); |
| | | verify(redisUtil, never()).compareAndDelete(anyString(), anyString()); |
| | | } |
| | | |
| | | @Test |
| | | void manualTriggerRunsAgentWithRunningLockAndDoesNotWriteSchedulerGuard() { |
| | | AutoTuneAgentService.AutoTuneAgentResult agentResult = successfulAgentResult(); |
| | | agentResult.setTriggerType(AutoTuneTriggerType.MANUAL.getCode()); |
| | | when(configService.getConfigValue("aiAutoTuneIntervalMinutes", "10")).thenReturn("10"); |
| | | when(wrkMastService.count(any(Wrapper.class))).thenReturn(1L); |
| | | when(redisUtil.trySetStringIfAbsent(anyString(), anyString(), anyLong())).thenReturn(true); |
| | | when(autoTuneAgentService.runAutoTune(AutoTuneTriggerType.MANUAL.getCode())).thenReturn(agentResult); |
| | | |
| | | AutoTuneCoordinatorService.AutoTuneCoordinatorResult result = coordinatorService().runManualAutoTune(); |
| | | |
| | | assertFalse(result.getSkipped()); |
| | | assertTrue(result.getTriggered()); |
| | | assertSame(agentResult, result.getAgentResult()); |
| | | verify(autoTuneAgentService).runAutoTune(AutoTuneTriggerType.MANUAL.getCode()); |
| | | verify(redisUtil, never()).set(eq(RedisKeyType.AI_AUTO_TUNE_LAST_TRIGGER_GUARD.key), any(), anyLong()); |
| | | verify(redisUtil).compareAndDelete(anyString(), anyString()); |
| | | |
| | | ArgumentCaptor<OperateLog> operateLogCaptor = ArgumentCaptor.forClass(OperateLog.class); |
| | | verify(operateLogService).save(operateLogCaptor.capture()); |
| | | assertEquals("ai_auto_tune_manual_trigger", operateLogCaptor.getValue().getAction()); |
| | | |
| | | ArgumentCaptor<AiAutoTuneJob> jobCaptor = ArgumentCaptor.forClass(AiAutoTuneJob.class); |
| | | verify(aiAutoTuneJobService).save(jobCaptor.capture()); |
| | | AiAutoTuneJob auditJob = jobCaptor.getValue(); |
| | | assertEquals(AutoTuneTriggerType.MANUAL.getCode(), auditJob.getTriggerType()); |
| | | assertEquals(AutoTuneJobStatus.SUCCESS.getCode(), auditJob.getStatus()); |
| | | assertNotNull(auditJob.getStartTime()); |
| | | assertNotNull(auditJob.getFinishTime()); |
| | | assertEquals(1, auditJob.getHasActiveTasks()); |
| | | assertEquals(AiPromptScene.AUTO_TUNE_DISPATCH.getCode(), auditJob.getPromptSceneCode()); |
| | | assertEquals("no changes needed", auditJob.getSummary()); |
| | | assertEquals(10, auditJob.getIntervalBefore()); |
| | | assertEquals(10, auditJob.getIntervalAfter()); |
| | | assertEquals(0, auditJob.getSuccessCount()); |
| | | assertEquals(0, auditJob.getRejectCount()); |
| | | assertEquals(1, auditJob.getLlmCallCount()); |
| | | assertEquals(10, auditJob.getPromptTokens()); |
| | | assertEquals(5, auditJob.getCompletionTokens()); |
| | | assertEquals(15, auditJob.getTotalTokens()); |
| | | } |
| | | |
| | | @Test |
| | | void manualTriggerSkipsWhenRunningLockIsNotAcquired() { |
| | | when(redisUtil.trySetStringIfAbsent(anyString(), anyString(), anyLong())).thenReturn(false); |
| | | |
| | | AutoTuneCoordinatorService.AutoTuneCoordinatorResult result = coordinatorService().runManualAutoTune(); |
| | | |
| | | assertTrue(result.getSkipped()); |
| | | assertEquals("running_lock_not_acquired", result.getReason()); |
| | | verify(autoTuneAgentService, never()).runAutoTune(anyString()); |
| | | verify(redisUtil, never()).compareAndDelete(anyString(), anyString()); |
| | | } |
| | | |
| | | @Test |
| | |
| | | promptTemplate.setContent("system prompt"); |
| | | when(aiPromptTemplateService.resolvePublished("wcs_auto_tune_dispatch")).thenReturn(promptTemplate); |
| | | when(mcpToolManager.buildOpenAiTools()).thenReturn(allowedOpenAiTools()); |
| | | when(mcpToolManager.callTool(any(), any(JSONObject.class))).thenReturn(Collections.singletonMap("ok", true)); |
| | | when(mcpToolManager.callTool(any(), any(JSONObject.class))).thenAnswer(invocation -> { |
| | | String toolName = invocation.getArgument(0); |
| | | JSONObject arguments = invocation.getArgument(1); |
| | | if ("wcs_local_dispatch_apply_auto_tune_changes".equals(toolName) |
| | | && Boolean.TRUE.equals(arguments.getBoolean("dryRun"))) { |
| | | LinkedHashMap<String, Object> dryRunOutput = new LinkedHashMap<>(); |
| | | dryRunOutput.put("success", true); |
| | | dryRunOutput.put("dryRun", true); |
| | | dryRunOutput.put("dryRunToken", "token-123"); |
| | | return dryRunOutput; |
| | | } |
| | | return Collections.singletonMap("ok", true); |
| | | }); |
| | | when(llmChatService.chatCompletion(any(), anyDouble(), anyInt(), any())) |
| | | .thenReturn( |
| | | response("read snapshot", toolCall("call_1", "wcs_local_dispatch_get_auto_tune_snapshot", "{}"), 10, 5), |
| | | response("dry-run first", toolCall("call_2", "wcs_local_dispatch_apply_auto_tune_changes", |
| | | "{\"dryRun\":true,\"changes\":[{\"targetType\":\"sys_config\",\"targetKey\":\"conveyorStationTaskLimit\",\"newValue\":\"12\"}]}"), 11, 6), |
| | | response("apply after dry-run", toolCall("call_3", "wcs_local_dispatch_apply_auto_tune_changes", |
| | | "{\"dryRun\":false,\"changes\":[{\"targetType\":\"sys_config\",\"targetKey\":\"conveyorStationTaskLimit\",\"newValue\":\"12\"}]}"), 12, 7), |
| | | "{\"dryRun\":false,\"dryRunToken\":\"token-123\",\"changes\":[{\"targetType\":\"sys_config\",\"targetKey\":\"conveyorStationTaskLimit\",\"newValue\":\"12\"}]}"), 12, 7), |
| | | response("已完成自动调参", null, 13, 8) |
| | | ); |
| | | |
| | |
| | | assertEquals("wcs_local_dispatch_apply_auto_tune_changes", toolNameCaptor.getAllValues().get(1)); |
| | | assertEquals(Boolean.TRUE, argumentCaptor.getAllValues().get(1).getBoolean("dryRun")); |
| | | assertEquals(Boolean.FALSE, argumentCaptor.getAllValues().get(2).getBoolean("dryRun")); |
| | | assertEquals("token-123", argumentCaptor.getAllValues().get(2).getString("dryRunToken")); |
| | | |
| | | ArgumentCaptor<List<Object>> toolsCaptor = ArgumentCaptor.forClass(List.class); |
| | | verify(llmChatService, times(4)).chatCompletion(any(), anyDouble(), anyInt(), toolsCaptor.capture()); |
| | | List<String> visibleToolNames = toolNames(toolsCaptor.getAllValues().get(0)); |
| | | assertEquals(4, visibleToolNames.size()); |
| | | assertTrue(visibleToolNames.contains("wcs_local_dispatch_get_auto_tune_snapshot")); |
| | | assertTrue(visibleToolNames.contains("wcs_local_dispatch_get_recent_auto_tune_jobs")); |
| | | assertTrue(visibleToolNames.contains("wcs_local_dispatch_apply_auto_tune_changes")); |
| | | assertTrue(visibleToolNames.contains("wcs_local_dispatch_revert_last_auto_tune_job")); |
| | | assertFalse(visibleToolNames.contains("wcs_local_device_get_crn_status")); |
| | | } |
| | | |
| | | @Test |
| | | void agentForcesAutoTriggerTypeOnApplyTools() { |
| | | AutoTuneAgentServiceImpl service = agentService(); |
| | | when(mcpToolManager.buildOpenAiTools()).thenReturn(allowedOpenAiTools()); |
| | | when(mcpToolManager.callTool(any(), any(JSONObject.class))).thenReturn(Collections.singletonMap("ok", true)); |
| | | when(llmChatService.chatCompletion(any(), anyDouble(), anyInt(), any())) |
| | | .thenReturn( |
| | | response("snapshot", toolCall("call_1", "wcs_local_dispatch_get_auto_tune_snapshot", |
| | | "{}"), 10, 5), |
| | | response("dry run", toolCall("call_2", "wcs_local_dispatch_apply_auto_tune_changes", |
| | | "{\"dryRun\":true,\"changes\":[]}"), 10, 5), |
| | | response("done", null, 10, 5) |
| | | ); |
| | | |
| | | AutoTuneAgentService.AutoTuneAgentResult result = service.runAutoTune("auto"); |
| | | |
| | | assertTrue(result.getSuccess()); |
| | | ArgumentCaptor<JSONObject> argumentCaptor = ArgumentCaptor.forClass(JSONObject.class); |
| | | verify(mcpToolManager).callTool(eq("wcs_local_dispatch_apply_auto_tune_changes"), argumentCaptor.capture()); |
| | | assertEquals("auto", argumentCaptor.getValue().getString("triggerType")); |
| | | } |
| | | |
| | | @Test |
| | |
| | | return new AutoTuneAgentServiceImpl(llmChatService, mcpToolManager, aiPromptTemplateService); |
| | | } |
| | | |
| | | private AutoTuneCoordinatorServiceImpl coordinatorService() { |
| | | return new AutoTuneCoordinatorServiceImpl( |
| | | configService, |
| | | wrkMastService, |
| | | aiAutoTuneJobService, |
| | | autoTuneAgentService, |
| | | redisUtil, |
| | | operateLogService); |
| | | } |
| | | |
| | | private AutoTuneAgentService.AutoTuneAgentResult successfulAgentResult() { |
| | | AutoTuneAgentService.AutoTuneAgentResult result = new AutoTuneAgentService.AutoTuneAgentResult(); |
| | | result.setSuccess(true); |
| | | result.setTriggerType(AutoTuneTriggerType.AUTO.getCode()); |
| | | result.setSummary("no changes needed"); |
| | | result.setToolCallCount(1); |
| | | result.setLlmCallCount(1); |
| | | result.setPromptTokens(10L); |
| | | result.setCompletionTokens(5L); |
| | | result.setTotalTokens(15L); |
| | | result.setMaxRoundsReached(false); |
| | | return result; |
| | | } |
| | | |
| | | private AutoTuneAgentService.AutoTuneAgentResult failedAgentResult() { |
| | | AutoTuneAgentService.AutoTuneAgentResult result = successfulAgentResult(); |
| | | result.setSuccess(false); |
| | | result.setSummary("failed"); |
| | | return result; |
| | | } |
| | | |
| | | private AutoTuneChangeCommand change(String targetType, String targetId, String targetKey, String newValue) { |
| | | AutoTuneChangeCommand command = new AutoTuneChangeCommand(); |
| | | command.setTargetType(targetType); |
| | |
| | | return tool; |
| | | } |
| | | |
| | | private List<String> toolNames(List<Object> tools) { |
| | | List<String> names = new ArrayList<>(); |
| | | for (Object tool : tools) { |
| | | Map<?, ?> toolMap = (Map<?, ?>) tool; |
| | | Map<?, ?> function = (Map<?, ?>) toolMap.get("function"); |
| | | names.add(String.valueOf(function.get("name"))); |
| | | } |
| | | return names; |
| | | } |
| | | |
| | | private ChatCompletionResponse response(String content, |
| | | ChatCompletionRequest.ToolCall toolCall, |
| | | int promptTokens, |