package com.zy.ai.service.impl;
|
|
import com.alibaba.fastjson.JSON;
|
import com.alibaba.fastjson.JSONObject;
|
import com.zy.ai.domain.autotune.AutoTuneTriggerType;
|
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.service.AiPromptTemplateService;
|
import com.zy.ai.service.AutoTuneAgentService;
|
import com.zy.ai.service.LlmChatService;
|
import lombok.RequiredArgsConstructor;
|
import lombok.extern.slf4j.Slf4j;
|
import org.springframework.stereotype.Service;
|
|
import java.util.ArrayList;
|
import java.util.List;
|
import java.util.Map;
|
import java.util.Set;
|
|
@Slf4j
|
@Service
|
@RequiredArgsConstructor
|
public class AutoTuneAgentServiceImpl implements AutoTuneAgentService {
|
|
private static final int MAX_TOOL_ROUNDS = 10;
|
private static final double TEMPERATURE = 0.2D;
|
private static final int MAX_TOKENS = 2048;
|
private static final String TOOL_GET_SNAPSHOT = "wcs_local_dispatch_get_auto_tune_snapshot";
|
private static final String TOOL_GET_RECENT_JOBS = "wcs_local_dispatch_get_recent_auto_tune_jobs";
|
private static final String TOOL_APPLY_CHANGES = "wcs_local_dispatch_apply_auto_tune_changes";
|
private static final String TOOL_REVERT_LAST_JOB = "wcs_local_dispatch_revert_last_auto_tune_job";
|
private static final Set<String> ALLOWED_TOOL_NAMES = Set.of(
|
TOOL_GET_SNAPSHOT,
|
TOOL_GET_RECENT_JOBS,
|
TOOL_APPLY_CHANGES,
|
TOOL_REVERT_LAST_JOB
|
);
|
|
private final LlmChatService llmChatService;
|
private final SpringAiMcpToolManager mcpToolManager;
|
private final AiPromptTemplateService aiPromptTemplateService;
|
|
@Override
|
public AutoTuneAgentResult runAutoTune(String triggerType) {
|
String normalizedTriggerType = normalizeTriggerType(triggerType);
|
UsageCounter usageCounter = new UsageCounter();
|
RunState runState = new RunState();
|
boolean maxRoundsReached = false;
|
StringBuilder summaryBuffer = new StringBuilder();
|
|
try {
|
List<Object> tools = filterAllowedTools(mcpToolManager.buildOpenAiTools());
|
if (tools == null || tools.isEmpty()) {
|
throw new IllegalStateException("No auto-tune MCP tools registered");
|
}
|
|
AiPromptTemplate promptTemplate = aiPromptTemplateService.resolvePublished(AiPromptScene.AUTO_TUNE_DISPATCH.getCode());
|
List<ChatCompletionRequest.Message> messages = buildMessages(promptTemplate, normalizedTriggerType);
|
|
for (int round = 0; round < MAX_TOOL_ROUNDS; round++) {
|
ChatCompletionResponse response = llmChatService.chatCompletion(messages, TEMPERATURE, MAX_TOKENS, tools);
|
ChatCompletionRequest.Message assistantMessage = extractAssistantMessage(response);
|
usageCounter.add(response.getUsage());
|
messages.add(assistantMessage);
|
appendSummary(summaryBuffer, assistantMessage.getContent());
|
|
List<ChatCompletionRequest.ToolCall> toolCalls = assistantMessage.getTool_calls();
|
if (toolCalls == null || toolCalls.isEmpty()) {
|
return buildResult(runState.isSuccessful(), normalizedTriggerType, summaryBuffer, runState,
|
usageCounter, false);
|
}
|
|
for (ChatCompletionRequest.ToolCall toolCall : toolCalls) {
|
Object toolOutput = callMountedTool(toolCall, runState, normalizedTriggerType);
|
messages.add(buildToolMessage(toolCall, toolOutput));
|
}
|
}
|
maxRoundsReached = true;
|
return buildResult(false, normalizedTriggerType, summaryBuffer, runState, usageCounter, maxRoundsReached);
|
} catch (Exception exception) {
|
log.error("Auto tune agent stopped with error", exception);
|
appendSummary(summaryBuffer, "自动调参 Agent 执行异常: " + exception.getMessage());
|
runState.markToolError();
|
return buildResult(false, normalizedTriggerType, summaryBuffer, runState, usageCounter, maxRoundsReached);
|
}
|
}
|
|
private List<ChatCompletionRequest.Message> buildMessages(AiPromptTemplate promptTemplate, String triggerType) {
|
List<ChatCompletionRequest.Message> messages = new ArrayList<>();
|
|
ChatCompletionRequest.Message systemMessage = new ChatCompletionRequest.Message();
|
systemMessage.setRole("system");
|
systemMessage.setContent(promptTemplate == null ? "" : promptTemplate.getContent());
|
messages.add(systemMessage);
|
|
ChatCompletionRequest.Message userMessage = new ChatCompletionRequest.Message();
|
userMessage.setRole("user");
|
userMessage.setContent("请执行一次后台 WCS 自动调参。triggerType=" + triggerType
|
+ "。必须先调用 wcs_local_dispatch_get_auto_tune_snapshot 获取事实；如需提交变更，"
|
+ "必须先 dry-run，再根据 dry-run 结果决定是否实际应用；实际应用时必须带上 dry-run 返回的 dryRunToken。"
|
+ "不要输出自由格式 JSON 供外层解析。");
|
messages.add(userMessage);
|
return messages;
|
}
|
|
private ChatCompletionRequest.Message extractAssistantMessage(ChatCompletionResponse response) {
|
if (response == null || response.getChoices() == null || response.getChoices().isEmpty()) {
|
throw new IllegalStateException("LLM returned empty response");
|
}
|
ChatCompletionRequest.Message message = response.getChoices().get(0).getMessage();
|
if (message == null) {
|
throw new IllegalStateException("LLM returned empty message");
|
}
|
return message;
|
}
|
|
private Object callMountedTool(ChatCompletionRequest.ToolCall toolCall, RunState runState, String triggerType) {
|
String toolName = resolveToolName(toolCall);
|
if (!ALLOWED_TOOL_NAMES.contains(toolName)) {
|
throw new IllegalArgumentException("Disallowed auto-tune MCP tool: " + toolName);
|
}
|
JSONObject arguments = parseArguments(toolCall);
|
applySchedulerTriggerType(toolName, triggerType, arguments);
|
try {
|
Object output = mcpToolManager.callTool(toolName, arguments);
|
runState.markToolSuccess(toolName);
|
return output;
|
} catch (Exception exception) {
|
throw new IllegalStateException("Auto-tune MCP tool failed: " + toolName + ", " + exception.getMessage(),
|
exception);
|
}
|
}
|
|
private void applySchedulerTriggerType(String toolName, String triggerType, JSONObject arguments) {
|
if (!TOOL_APPLY_CHANGES.equals(toolName)) {
|
return;
|
}
|
if (!AutoTuneTriggerType.AUTO.getCode().equals(triggerType)) {
|
return;
|
}
|
arguments.put("triggerType", AutoTuneTriggerType.AUTO.getCode());
|
}
|
|
private ChatCompletionRequest.Message buildToolMessage(ChatCompletionRequest.ToolCall toolCall, Object toolOutput) {
|
ChatCompletionRequest.Message toolMessage = new ChatCompletionRequest.Message();
|
toolMessage.setRole("tool");
|
toolMessage.setTool_call_id(toolCall == null ? null : toolCall.getId());
|
toolMessage.setContent(JSON.toJSONString(toolOutput));
|
return toolMessage;
|
}
|
|
private String resolveToolName(ChatCompletionRequest.ToolCall toolCall) {
|
if (toolCall == null || toolCall.getFunction() == null || isBlank(toolCall.getFunction().getName())) {
|
throw new IllegalArgumentException("missing tool name");
|
}
|
return toolCall.getFunction().getName();
|
}
|
|
private JSONObject parseArguments(ChatCompletionRequest.ToolCall toolCall) {
|
String rawArguments = toolCall == null || toolCall.getFunction() == null
|
? null
|
: toolCall.getFunction().getArguments();
|
if (isBlank(rawArguments)) {
|
return new JSONObject();
|
}
|
try {
|
return JSON.parseObject(rawArguments);
|
} catch (Exception exception) {
|
JSONObject arguments = new JSONObject();
|
arguments.put("_raw", rawArguments);
|
return arguments;
|
}
|
}
|
|
private AutoTuneAgentResult buildResult(boolean success,
|
String triggerType,
|
StringBuilder summaryBuffer,
|
RunState runState,
|
UsageCounter usageCounter,
|
boolean maxRoundsReached) {
|
AutoTuneAgentResult result = new AutoTuneAgentResult();
|
result.setSuccess(success);
|
result.setTriggerType(triggerType);
|
result.setToolCallCount(runState.getToolCallCount());
|
result.setLlmCallCount(usageCounter.getLlmCallCount());
|
result.setPromptTokens(usageCounter.getPromptTokens());
|
result.setCompletionTokens(usageCounter.getCompletionTokens());
|
result.setTotalTokens(usageCounter.getTotalTokens());
|
result.setMaxRoundsReached(maxRoundsReached);
|
|
String summary = summaryBuffer == null ? "" : summaryBuffer.toString().trim();
|
if (runState.getToolCallCount() <= 0) {
|
summary = "自动调参 Agent 未调用任何允许的 MCP 工具，未执行调参。" + (summary.isEmpty() ? "" : "\n" + summary);
|
} else if (!runState.isSnapshotCalled()) {
|
summary = summary + "\n自动调参 Agent 未调用快照工具，结果不完整。";
|
}
|
if (runState.hasToolError()) {
|
summary = summary + "\n自动调参 Agent 存在工具调用错误，已标记为失败。";
|
}
|
if (maxRoundsReached) {
|
summary = summary + "\n自动调参 Agent 达到最大工具调用轮次，已停止。";
|
}
|
result.setSummary(summary);
|
return result;
|
}
|
|
private List<Object> filterAllowedTools(List<Object> tools) {
|
List<Object> allowedTools = new ArrayList<>();
|
if (tools == null || tools.isEmpty()) {
|
return allowedTools;
|
}
|
for (Object tool : tools) {
|
String toolName = resolveOpenAiToolName(tool);
|
if (ALLOWED_TOOL_NAMES.contains(toolName)) {
|
allowedTools.add(tool);
|
}
|
}
|
return allowedTools;
|
}
|
|
private String resolveOpenAiToolName(Object tool) {
|
if (!(tool instanceof Map<?, ?> toolMap)) {
|
return null;
|
}
|
Object function = toolMap.get("function");
|
if (!(function instanceof Map<?, ?> functionMap)) {
|
return null;
|
}
|
Object name = functionMap.get("name");
|
return name == null ? null : String.valueOf(name);
|
}
|
|
private void appendSummary(StringBuilder summaryBuffer, String content) {
|
if (summaryBuffer == null || isBlank(content)) {
|
return;
|
}
|
if (summaryBuffer.length() > 0) {
|
summaryBuffer.append('\n');
|
}
|
summaryBuffer.append(content.trim());
|
}
|
|
private String normalizeTriggerType(String triggerType) {
|
return isBlank(triggerType) ? "agent" : triggerType.trim();
|
}
|
|
private boolean isBlank(String value) {
|
return value == null || value.trim().isEmpty();
|
}
|
|
private static class UsageCounter {
|
private long promptTokens;
|
private long completionTokens;
|
private long totalTokens;
|
private int llmCallCount;
|
|
void add(ChatCompletionResponse.Usage usage) {
|
llmCallCount++;
|
if (usage == null) {
|
return;
|
}
|
promptTokens += usage.getPromptTokens() == null ? 0L : usage.getPromptTokens();
|
completionTokens += usage.getCompletionTokens() == null ? 0L : usage.getCompletionTokens();
|
totalTokens += usage.getTotalTokens() == null ? 0L : usage.getTotalTokens();
|
}
|
|
long getPromptTokens() {
|
return promptTokens;
|
}
|
|
long getCompletionTokens() {
|
return completionTokens;
|
}
|
|
long getTotalTokens() {
|
return totalTokens;
|
}
|
|
int getLlmCallCount() {
|
return llmCallCount;
|
}
|
}
|
|
private static class RunState {
|
private int toolCallCount;
|
private boolean snapshotCalled;
|
private boolean toolError;
|
|
void markToolSuccess(String toolName) {
|
toolCallCount++;
|
if (TOOL_GET_SNAPSHOT.equals(toolName)) {
|
snapshotCalled = true;
|
}
|
}
|
|
void markToolError() {
|
toolError = true;
|
}
|
|
boolean isSuccessful() {
|
return toolCallCount > 0 && snapshotCalled && !toolError;
|
}
|
|
int getToolCallCount() {
|
return toolCallCount;
|
}
|
|
boolean isSnapshotCalled() {
|
return snapshotCalled;
|
}
|
|
boolean hasToolError() {
|
return toolError;
|
}
|
}
|
}
|
