From 1b8a4677f362d234d834120deac4880d7ae89a50 Mon Sep 17 00:00:00 2001
From: Junjie <fallin.jie@qq.com>
Date: 星期四, 12 三月 2026 17:03:59 +0800
Subject: [PATCH] #
---
src/main/java/com/zy/ai/service/WcsDiagnosisService.java | 176 ++++++++++++++++++++++++++--------------------------------
1 files changed, 79 insertions(+), 97 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..3f65ed0 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,18 +24,16 @@
@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) {
List<ChatCompletionRequest.Message> messages = new ArrayList<>();
@@ -61,25 +56,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 +75,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);
}
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) {
@@ -178,6 +114,7 @@
if (tools.isEmpty()) {
throw new IllegalStateException("No MCP tools registered");
}
+ AgentUsageStats usageStats = new AgentUsageStats();
baseMessages.add(systemPrompt);
baseMessages.add(userQuestion);
@@ -195,6 +132,7 @@
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);
@@ -256,45 +194,34 @@
} 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);
+ aiChatStoreService.saveConversation(chatId,
+ buildTitleFromPrompt(userQuestion.getContent()),
+ userQuestion,
+ 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 +238,28 @@
} 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 sendLargeText(SseEmitter emitter, String text) {
@@ -429,6 +378,39 @@
}
}
+ 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