From caf3bdd9bbb629c8bc6f1a19b3ccdf441bf7650c Mon Sep 17 00:00:00 2001
From: Junjie <fallin.jie@qq.com>
Date: 星期日, 15 三月 2026 17:46:47 +0800
Subject: [PATCH] #
---
src/main/java/com/zy/ai/service/WcsDiagnosisService.java | 215 ++++++++++++++++++++++++++++-------------------------
1 files changed, 115 insertions(+), 100 deletions(-)
diff --git a/src/main/java/com/zy/ai/service/WcsDiagnosisService.java b/src/main/java/com/zy/ai/service/WcsDiagnosisService.java
index 757e877..6288698 100644
--- a/src/main/java/com/zy/ai/service/WcsDiagnosisService.java
+++ b/src/main/java/com/zy/ai/service/WcsDiagnosisService.java
@@ -8,10 +8,7 @@
import com.zy.ai.entity.WcsDiagnosisRequest;
import com.zy.ai.enums.AiPromptScene;
import com.zy.ai.mcp.service.SpringAiMcpToolManager;
-import com.zy.ai.service.AiPromptTemplateService;
import com.zy.ai.utils.AiUtils;
-import com.zy.common.utils.RedisUtil;
-import com.zy.core.enums.RedisKeyType;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import lombok.RequiredArgsConstructor;
@@ -27,21 +24,25 @@
@Slf4j
public class WcsDiagnosisService {
- private static final long CHAT_TTL_SECONDS = 7L * 24 * 3600;
-
@Autowired
private LlmChatService llmChatService;
- @Autowired
- private RedisUtil redisUtil;
@Autowired
private AiUtils aiUtils;
@Autowired
private SpringAiMcpToolManager mcpToolManager;
@Autowired
private AiPromptTemplateService aiPromptTemplateService;
+ @Autowired
+ private AiChatStoreService aiChatStoreService;
- public void diagnoseStream(WcsDiagnosisRequest request, SseEmitter emitter) {
+ public void diagnoseStream(WcsDiagnosisRequest request,
+ String chatId,
+ boolean reset,
+ SseEmitter emitter) {
List<ChatCompletionRequest.Message> messages = new ArrayList<>();
+ if (chatId != null && !chatId.isEmpty() && reset) {
+ aiChatStoreService.deleteChat(chatId);
+ }
AiPromptTemplate promptTemplate = aiPromptTemplateService.resolvePublished(AiPromptScene.DIAGNOSE_STREAM.getCode());
ChatCompletionRequest.Message mcpSystem = new ChatCompletionRequest.Message();
@@ -52,7 +53,11 @@
mcpUser.setRole("user");
mcpUser.setContent(aiUtils.buildDiagnosisUserContentMcp(request));
- runMcpStreamingDiagnosis(messages, mcpSystem, mcpUser, promptTemplate, 0.3, 2048, emitter, null);
+ ChatCompletionRequest.Message storedUser = new ChatCompletionRequest.Message();
+ storedUser.setRole("user");
+ storedUser.setContent(buildDiagnoseDisplayPrompt(request));
+
+ runMcpStreamingDiagnosis(messages, mcpSystem, mcpUser, storedUser, promptTemplate, 0.3, 2048, emitter, chatId);
}
public void askStream(String prompt,
@@ -61,25 +66,13 @@
SseEmitter emitter) {
List<ChatCompletionRequest.Message> messages = new ArrayList<>();
- List<ChatCompletionRequest.Message> history = null;
- String historyKey = null;
- String metaKey = null;
if (chatId != null && !chatId.isEmpty()) {
- historyKey = RedisKeyType.AI_CHAT_HISTORY.key + chatId;
- metaKey = RedisKeyType.AI_CHAT_META.key + chatId;
if (reset) {
- redisUtil.del(historyKey, metaKey);
+ aiChatStoreService.deleteChat(chatId);
}
- List<Object> stored = redisUtil.lGet(historyKey, 0, -1);
- if (stored != null && !stored.isEmpty()) {
- history = new ArrayList<>(stored.size());
- for (Object o : stored) {
- ChatCompletionRequest.Message m = convertToMessage(o);
- if (m != null) history.add(m);
- }
- if (!history.isEmpty()) messages.addAll(history);
- } else {
- history = new ArrayList<>();
+ List<ChatCompletionRequest.Message> history = aiChatStoreService.getChatHistory(chatId);
+ if (history != null && !history.isEmpty()) {
+ messages.addAll(history);
}
}
@@ -92,68 +85,21 @@
ChatCompletionRequest.Message mcpUser = new ChatCompletionRequest.Message();
mcpUser.setRole("user");
- mcpUser.setContent("銆愮敤鎴锋彁闂�慭n" + (prompt == null ? "" : prompt));
+ mcpUser.setContent(prompt == null ? "" : prompt);
- runMcpStreamingDiagnosis(messages, mcpSystem, mcpUser, promptTemplate, 0.3, 2048, emitter, finalChatId);
+ runMcpStreamingDiagnosis(messages, mcpSystem, mcpUser, mcpUser, promptTemplate, 0.3, 2048, emitter, finalChatId);
}
public List<Map<String, Object>> listChats() {
- java.util.Set<String> keys = redisUtil.scanKeys(RedisKeyType.AI_CHAT_META.key, 1000);
- List<Map<String, Object>> resp = new ArrayList<>();
- if (keys != null) {
- for (String key : keys) {
- Map<Object, Object> m = redisUtil.hmget(key);
- if (m != null && !m.isEmpty()) {
- java.util.HashMap<String, Object> item = new java.util.HashMap<>();
- for (Map.Entry<Object, Object> e : m.entrySet()) {
- item.put(String.valueOf(e.getKey()), e.getValue());
- }
- String chatId = String.valueOf(item.get("chatId"));
- String historyKey = RedisKeyType.AI_CHAT_HISTORY.key + chatId;
- item.put("size", redisUtil.lGetListSize(historyKey));
- resp.add(item);
- }
- }
- }
- return resp;
+ return aiChatStoreService.listChats();
}
public boolean deleteChat(String chatId) {
- if (chatId == null || chatId.isEmpty()) return false;
- String historyKey = RedisKeyType.AI_CHAT_HISTORY.key + chatId;
- String metaKey = RedisKeyType.AI_CHAT_META.key + chatId;
- redisUtil.del(historyKey, metaKey);
- return true;
+ return aiChatStoreService.deleteChat(chatId);
}
public List<ChatCompletionRequest.Message> getChatHistory(String chatId) {
- if (chatId == null || chatId.isEmpty()) return java.util.Collections.emptyList();
- String historyKey = RedisKeyType.AI_CHAT_HISTORY.key + chatId;
- List<Object> stored = redisUtil.lGet(historyKey, 0, -1);
- List<ChatCompletionRequest.Message> result = new ArrayList<>();
- if (stored != null) {
- for (Object o : stored) {
- ChatCompletionRequest.Message m = convertToMessage(o);
- if (m != null) result.add(m);
- }
- }
- return result;
- }
-
- private ChatCompletionRequest.Message convertToMessage(Object o) {
- if (o instanceof ChatCompletionRequest.Message) {
- return (ChatCompletionRequest.Message) o;
- }
- if (o instanceof Map) {
- Map<?, ?> map = (Map<?, ?>) o;
- ChatCompletionRequest.Message m = new ChatCompletionRequest.Message();
- Object role = map.get("role");
- Object content = map.get("content");
- m.setRole(role == null ? null : String.valueOf(role));
- m.setContent(content == null ? null : String.valueOf(content));
- return m;
- }
- return null;
+ return aiChatStoreService.getChatHistory(chatId);
}
private String buildTitleFromPrompt(String prompt) {
@@ -165,6 +111,7 @@
private void runMcpStreamingDiagnosis(List<ChatCompletionRequest.Message> baseMessages,
ChatCompletionRequest.Message systemPrompt,
ChatCompletionRequest.Message userQuestion,
+ ChatCompletionRequest.Message storedUserQuestion,
AiPromptTemplate promptTemplate,
Double temperature,
Integer maxTokens,
@@ -178,6 +125,8 @@
if (tools.isEmpty()) {
throw new IllegalStateException("No MCP tools registered");
}
+ AgentUsageStats usageStats = new AgentUsageStats();
+ StringBuilder reasoningBuffer = new StringBuilder();
baseMessages.add(systemPrompt);
baseMessages.add(userQuestion);
@@ -186,19 +135,23 @@
messages.addAll(baseMessages);
sse(emitter, "<think>\\n姝ｅ湪鍒濆鍖栬瘖鏂笌宸ュ叿鐜...\\n");
+ appendReasoning(reasoningBuffer, "姝ｅ湪鍒濆鍖栬瘖鏂笌宸ュ叿鐜...\n");
int maxRound = 10;
int i = 0;
while(true) {
sse(emitter, "\\n姝ｅ湪鍒嗘瀽锛堢" + (i + 1) + "杞級...\\n");
+ appendReasoning(reasoningBuffer, "\n姝ｅ湪鍒嗘瀽锛堢" + (i + 1) + "杞級...\n");
ChatCompletionResponse resp = llmChatService.chatCompletion(messages, temperature, maxTokens, tools);
if (resp == null || resp.getChoices() == null || resp.getChoices().isEmpty() || resp.getChoices().get(0).getMessage() == null) {
throw new IllegalStateException("LLM returned empty response");
}
+ usageStats.add(resp.getUsage());
ChatCompletionRequest.Message assistant = resp.getChoices().get(0).getMessage();
messages.add(assistant);
sse(emitter, assistant.getContent());
+ appendReasoning(reasoningBuffer, assistant == null ? null : assistant.getContent());
List<ChatCompletionRequest.ToolCall> toolCalls = assistant.getTool_calls();
if (toolCalls == null || toolCalls.isEmpty()) {
@@ -209,6 +162,7 @@
String toolName = tc != null && tc.getFunction() != null ? tc.getFunction().getName() : null;
if (toolName == null || toolName.trim().isEmpty()) continue;
sse(emitter, "\\n鍑嗗璋冪敤宸ュ叿锛�" + toolName + "\\n");
+ appendReasoning(reasoningBuffer, "\n鍑嗗璋冪敤宸ュ叿锛�" + toolName + "\n");
JSONObject args = new JSONObject();
if (tc.getFunction() != null && tc.getFunction().getArguments() != null && !tc.getFunction().getArguments().trim().isEmpty()) {
try {
@@ -228,6 +182,7 @@
output = err;
}
sse(emitter, "\\n宸ュ叿杩斿洖锛屾鍦ㄧ户缁帹鐞�...\\n");
+ appendReasoning(reasoningBuffer, "\n宸ュ叿杩斿洖锛屾鍦ㄧ户缁帹鐞�...\n");
ChatCompletionRequest.Message toolMsg = new ChatCompletionRequest.Message();
toolMsg.setRole("tool");
toolMsg.setTool_call_id(tc == null ? null : tc.getId());
@@ -239,6 +194,7 @@
}
sse(emitter, "\\n姝ｅ湪鏍规嵁鏁版嵁杩涜鍒嗘瀽...\\n</think>\\n\\n");
+ appendReasoning(reasoningBuffer, "\n姝ｅ湪鏍规嵁鏁版嵁杩涜鍒嗘瀽...\n");
ChatCompletionRequest.Message diagnosisMessage = new ChatCompletionRequest.Message();
diagnosisMessage.setRole("system");
@@ -256,45 +212,35 @@
} catch (Exception ignore) {}
}, () -> {
try {
+ emitTokenUsage(emitter, usageStats);
sse(emitter, "\\n\\n銆怉I銆戣繍琛屽凡鍋滄锛堟甯哥粨鏉燂級\\n\\n");
log.info("AI MCP diagnose stopped: final end");
emitter.complete();
if (chatId != null) {
- String historyKey = RedisKeyType.AI_CHAT_HISTORY.key + chatId;
- String metaKey = RedisKeyType.AI_CHAT_META.key + chatId;
-
ChatCompletionRequest.Message a = new ChatCompletionRequest.Message();
a.setRole("assistant");
a.setContent(assistantBuffer.toString());
- redisUtil.lSet(historyKey, userQuestion);
- redisUtil.lSet(historyKey, a);
- redisUtil.expire(historyKey, CHAT_TTL_SECONDS);
- Map<Object, Object> old = redisUtil.hmget(metaKey);
- Long createdAt = old != null && old.get("createdAt") != null ?
- (old.get("createdAt") instanceof Number ? ((Number) old.get("createdAt")).longValue() : Long.valueOf(String.valueOf(old.get("createdAt"))))
- : System.currentTimeMillis();
- Map<String, Object> meta = new java.util.HashMap<>();
- meta.put("chatId", chatId);
- meta.put("title", buildTitleFromPrompt(userQuestion.getContent()));
- if (promptTemplate != null) {
- meta.put("promptTemplateId", promptTemplate.getId());
- meta.put("promptSceneCode", promptTemplate.getSceneCode());
- meta.put("promptVersion", promptTemplate.getVersion());
- meta.put("promptName", promptTemplate.getName());
- }
- meta.put("createdAt", createdAt);
- meta.put("updatedAt", System.currentTimeMillis());
- redisUtil.hmset(metaKey, meta, CHAT_TTL_SECONDS);
+ a.setReasoningContent(reasoningBuffer.toString());
+ aiChatStoreService.saveConversation(chatId,
+ buildTitleFromPrompt(storedUserQuestion == null ? null : storedUserQuestion.getContent()),
+ storedUserQuestion == null ? userQuestion : storedUserQuestion,
+ a,
+ promptTemplate,
+ usageStats.getPromptTokens(),
+ usageStats.getCompletionTokens(),
+ usageStats.getTotalTokens(),
+ usageStats.getLlmCallCount());
}
} catch (Exception ignore) {}
}, e -> {
try {
+ emitTokenUsage(emitter, usageStats);
sse(emitter, "\\n\\n銆怉I銆戝垎鏋愬嚭閿欙紝杩愯宸插仠姝紙寮傚父锛塡\n\\n");
log.error("AI MCP diagnose stopped: stream error", e);
emitter.complete();
} catch (Exception ignore) {}
- });
+ }, usageStats::add);
} catch (Exception e) {
try {
sse(emitter, "\\n\\n銆怉I銆戣繍琛屽凡鍋滄锛堝紓甯革級\\n\\n");
@@ -311,6 +257,35 @@
} catch (Exception e) {
log.warn("SSE send failed", e);
}
+ }
+
+ private void emitTokenUsage(SseEmitter emitter, AgentUsageStats usageStats) {
+ if (emitter == null || usageStats == null || usageStats.getTotalTokens() <= 0) {
+ return;
+ }
+ try {
+ emitter.send(SseEmitter.event()
+ .name("token_usage")
+ .data(JSON.toJSONString(buildTokenUsagePayload(usageStats))));
+ } catch (Exception e) {
+ log.warn("SSE token usage send failed", e);
+ }
+ }
+
+ private Map<String, Object> buildTokenUsagePayload(AgentUsageStats usageStats) {
+ java.util.LinkedHashMap<String, Object> payload = new java.util.LinkedHashMap<>();
+ payload.put("promptTokens", usageStats.getPromptTokens());
+ payload.put("completionTokens", usageStats.getCompletionTokens());
+ payload.put("totalTokens", usageStats.getTotalTokens());
+ payload.put("llmCallCount", usageStats.getLlmCallCount());
+ return payload;
+ }
+
+ private void appendReasoning(StringBuilder reasoningBuffer, String text) {
+ if (reasoningBuffer == null || text == null || text.isEmpty()) {
+ return;
+ }
+ reasoningBuffer.append(text);
}
private void sendLargeText(SseEmitter emitter, String text) {
@@ -429,6 +404,46 @@
}
}
+ private String buildDiagnoseDisplayPrompt(WcsDiagnosisRequest request) {
+ if (request == null || request.getAlarmMessage() == null || request.getAlarmMessage().trim().isEmpty()) {
+ return "瀵瑰綋鍓嶇郴缁熻繘琛屽贰妫�";
+ }
+ return request.getAlarmMessage().trim();
+ }
+
+ private static class AgentUsageStats {
+ private long promptTokens;
+ private long completionTokens;
+ private long totalTokens;
+ private int llmCallCount;
+
+ void add(ChatCompletionResponse.Usage usage) {
+ 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();
+ llmCallCount++;
+ }
+
+ long getPromptTokens() {
+ return promptTokens;
+ }
+
+ long getCompletionTokens() {
+ return completionTokens;
+ }
+
+ long getTotalTokens() {
+ return totalTokens;
+ }
+
+ int getLlmCallCount() {
+ return llmCallCount;
+ }
+ }
+
private boolean isConclusionText(String content) {
if (content == null) return false;
String c = content;
--
Gitblit v1.9.1