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 |  382 ++++++++++++++++++++---------------------------------
 1 files changed, 145 insertions(+), 237 deletions(-)

diff --git a/src/main/java/com/zy/ai/service/WcsDiagnosisService.java b/src/main/java/com/zy/ai/service/WcsDiagnosisService.java
index 189e57d..6288698 100644
--- a/src/main/java/com/zy/ai/service/WcsDiagnosisService.java
+++ b/src/main/java/com/zy/ai/service/WcsDiagnosisService.java
@@ -2,16 +2,14 @@
 
 import com.alibaba.fastjson.JSON;
 import com.alibaba.fastjson.JSONObject;
+import com.zy.ai.entity.AiPromptTemplate;
 import com.zy.ai.entity.ChatCompletionRequest;
 import com.zy.ai.entity.ChatCompletionResponse;
 import com.zy.ai.entity.WcsDiagnosisRequest;
-import com.zy.ai.mcp.controller.McpController;
-import com.zy.ai.utils.AiPromptUtils;
+import com.zy.ai.enums.AiPromptScene;
+import com.zy.ai.mcp.service.SpringAiMcpToolManager;
 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.beans.factory.annotation.Value;
 import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
 import lombok.RequiredArgsConstructor;
 import lombok.extern.slf4j.Slf4j;
@@ -26,228 +24,82 @@
 @Slf4j
 public class WcsDiagnosisService {
 
-    private static final long CHAT_TTL_SECONDS = 7L * 24 * 3600;
-
-    @Value("${llm.platform}")
-    private String platform;
     @Autowired
     private LlmChatService llmChatService;
     @Autowired
-    private RedisUtil redisUtil;
-    @Autowired
-    private AiPromptUtils aiPromptUtils;
-    @Autowired
     private AiUtils aiUtils;
-    @Autowired(required = false)
-    private McpController mcpController;
     @Autowired
-    private PythonService pythonService;
+    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();
         mcpSystem.setRole("system");
-        mcpSystem.setContent(aiPromptUtils.getAiDiagnosePromptMcp());
+        mcpSystem.setContent(promptTemplate.getContent());
 
         ChatCompletionRequest.Message mcpUser = new ChatCompletionRequest.Message();
         mcpUser.setRole("user");
         mcpUser.setContent(aiUtils.buildDiagnosisUserContentMcp(request));
 
-        if (runMcpStreamingDiagnosis(messages, mcpSystem, mcpUser, 0.3, 2048, emitter, null)) {
-            return;
-        }
+        ChatCompletionRequest.Message storedUser = new ChatCompletionRequest.Message();
+        storedUser.setRole("user");
+        storedUser.setContent(buildDiagnoseDisplayPrompt(request));
 
-        messages = new ArrayList<>();
-        ChatCompletionRequest.Message system = new ChatCompletionRequest.Message();
-        system.setRole("system");
-        system.setContent(aiPromptUtils.getAiDiagnosePrompt());
-        messages.add(system);
-
-        ChatCompletionRequest.Message user = new ChatCompletionRequest.Message();
-        user.setRole("user");
-        user.setContent(aiUtils.buildDiagnosisUserContent(request));
-        messages.add(user);
-
-        llmChatService.chatStream(messages, 0.3, 2048, s -> {
-            try {
-                String safe = s == null ? "" : s.replace("\r", "").replace("\n", "\\n");
-                if (!safe.isEmpty()) {
-                    emitter.send(SseEmitter.event().data(safe));
-                }
-            } catch (Exception ignore) {}
-        }, () -> {
-            try {
-                log.info("AI diagnose stream stopped: normal end");
-                emitter.complete();
-            } catch (Exception ignore) {}
-        }, e -> {
-            try {
-                try { emitter.send(SseEmitter.event().data("銆怉I銆戣繍琛屽凡鍋滄锛堝紓甯革級")); } catch (Exception ignore) {}
-                log.error("AI diagnose stream stopped: error", e);
-                emitter.completeWithError(e);
-            } catch (Exception ignore) {}
-        });
+        runMcpStreamingDiagnosis(messages, mcpSystem, mcpUser, storedUser, promptTemplate, 0.3, 2048, emitter, chatId);
     }
 
-    public void askStream(WcsDiagnosisRequest request,
-                          String prompt,
+    public void askStream(String prompt,
                           String chatId,
                           boolean reset,
                           SseEmitter emitter) {
-        if (platform.equals("python")) {
-            pythonService.runPython(prompt, chatId, emitter);
-            return;
-        }
-
         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);
             }
         }
 
-        StringBuilder assistantBuffer = new StringBuilder();
         final String finalChatId = chatId;
-        final String finalHistoryKey = historyKey;
-        final String finalMetaKey = metaKey;
-        final String finalPrompt = prompt;
+        AiPromptTemplate promptTemplate = aiPromptTemplateService.resolvePublished(AiPromptScene.SENSOR_CHAT.getCode());
 
         ChatCompletionRequest.Message mcpSystem = new ChatCompletionRequest.Message();
         mcpSystem.setRole("system");
-        mcpSystem.setContent(aiPromptUtils.getWcsSensorPromptMcp());
+        mcpSystem.setContent(promptTemplate.getContent());
 
         ChatCompletionRequest.Message mcpUser = new ChatCompletionRequest.Message();
         mcpUser.setRole("user");
-        mcpUser.setContent("銆愮敤鎴锋彁闂�慭n" + (prompt == null ? "" : prompt));
+        mcpUser.setContent(prompt == null ? "" : prompt);
 
-        if (runMcpStreamingDiagnosis(messages, mcpSystem, mcpUser, 0.3, 2048, emitter, finalChatId)) {
-            return;
-        }
-
-        messages = new ArrayList<>();
-        ChatCompletionRequest.Message system = new ChatCompletionRequest.Message();
-        system.setRole("system");
-        system.setContent(aiPromptUtils.getWcsSensorPrompt());
-        messages.add(system);
-
-        ChatCompletionRequest.Message questionMsg = new ChatCompletionRequest.Message();
-        questionMsg.setRole("user");
-        questionMsg.setContent("銆愮敤鎴锋彁闂�慭n" + (prompt == null ? "" : prompt));
-        messages.add(questionMsg);
-
-        llmChatService.chatStream(messages, 0.3, 2048, s -> {
-            try {
-                String safe = s == null ? "" : s.replace("\r", "").replace("\n", "\\n");
-                if (!safe.isEmpty()) {
-                    emitter.send(SseEmitter.event().data(safe));
-                    assistantBuffer.append(s);
-                }
-            } catch (Exception ignore) {}
-        }, () -> {
-            try {
-                if (finalChatId != null && !finalChatId.isEmpty()) {
-                    ChatCompletionRequest.Message q = new ChatCompletionRequest.Message();
-                    q.setRole("user");
-                    q.setContent(finalPrompt == null ? "" : finalPrompt);
-                    ChatCompletionRequest.Message a = new ChatCompletionRequest.Message();
-                    a.setRole("assistant");
-                    a.setContent(assistantBuffer.toString());
-                    redisUtil.lSet(finalHistoryKey, q);
-                    redisUtil.lSet(finalHistoryKey, a);
-                    redisUtil.expire(finalHistoryKey, CHAT_TTL_SECONDS);
-                    Map<Object, Object> old = redisUtil.hmget(finalMetaKey);
-                    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", finalChatId);
-                    meta.put("title", buildTitleFromPrompt(finalPrompt));
-                    meta.put("createdAt", createdAt);
-                    meta.put("updatedAt", System.currentTimeMillis());
-                    redisUtil.hmset(finalMetaKey, meta, CHAT_TTL_SECONDS);
-                }
-                emitter.complete();
-            } catch (Exception ignore) {}
-        }, e -> {
-            try { emitter.completeWithError(e); } catch (Exception ignore) {}
-        });
+        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) {
@@ -256,17 +108,25 @@
         return p.length() > 20 ? p.substring(0, 20) : p;
     }
  
-    private boolean runMcpStreamingDiagnosis(List<ChatCompletionRequest.Message> baseMessages,
-                                             ChatCompletionRequest.Message systemPrompt,
-                                             ChatCompletionRequest.Message userQuestion,
-                                             Double temperature,
-                                             Integer maxTokens,
-                                             SseEmitter emitter,
-                                             String chatId) {
+    private void runMcpStreamingDiagnosis(List<ChatCompletionRequest.Message> baseMessages,
+                                          ChatCompletionRequest.Message systemPrompt,
+                                          ChatCompletionRequest.Message userQuestion,
+                                          ChatCompletionRequest.Message storedUserQuestion,
+                                          AiPromptTemplate promptTemplate,
+                                          Double temperature,
+                                          Integer maxTokens,
+                                          SseEmitter emitter,
+                                          String chatId) {
         try {
-            if (mcpController == null) return false;
-            List<Object> tools = buildOpenAiTools();
-            if (tools.isEmpty()) return false;
+            if (mcpToolManager == null) {
+                throw new IllegalStateException("Spring AI MCP tool manager is unavailable");
+            }
+            List<Object> tools = mcpToolManager.buildOpenAiTools();
+            if (tools.isEmpty()) {
+                throw new IllegalStateException("No MCP tools registered");
+            }
+            AgentUsageStats usageStats = new AgentUsageStats();
+            StringBuilder reasoningBuffer = new StringBuilder();
 
             baseMessages.add(systemPrompt);
             baseMessages.add(userQuestion);
@@ -275,20 +135,23 @@
             messages.addAll(baseMessages);
  
             sse(emitter, "<think>\\n姝e湪鍒濆鍖栬瘖鏂笌宸ュ叿鐜...\\n");
+            appendReasoning(reasoningBuffer, "姝e湪鍒濆鍖栬瘖鏂笌宸ュ叿鐜...\n");
 
             int maxRound = 10;
             int i = 0;
             while(true) {
                 sse(emitter, "\\n姝e湪鍒嗘瀽锛堢" + (i + 1) + "杞級...\\n");
+                appendReasoning(reasoningBuffer, "\n姝e湪鍒嗘瀽锛堢" + (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) {
-                    sse(emitter, "\\n鍒嗘瀽鍑洪敊锛屾鍦ㄥ洖閫�...\\n");
-                    return false;
+                    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()) {
@@ -299,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 {
@@ -310,7 +174,7 @@
                     }
                     Object output;
                     try {
-                        output = mcpController.callTool(toolName, args);
+                        output = mcpToolManager.callTool(toolName, args);
                     } catch (Exception e) {
                         java.util.LinkedHashMap<String, Object> err = new java.util.LinkedHashMap<String, Object>();
                         err.put("tool", toolName);
@@ -318,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());
@@ -329,6 +194,7 @@
             }
 
             sse(emitter, "\\n姝e湪鏍规嵁鏁版嵁杩涜鍒嗘瀽...\\n</think>\\n\\n");
+            appendReasoning(reasoningBuffer, "\n姝e湪鏍规嵁鏁版嵁杩涜鍒嗘瀽...\n");
 
             ChatCompletionRequest.Message diagnosisMessage = new ChatCompletionRequest.Message();
             diagnosisMessage.setRole("system");
@@ -346,43 +212,41 @@
                 } 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()));
-                        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 -> {
-                sse(emitter, "\\n\\n銆怉I銆戝垎鏋愬嚭閿欙紝姝e湪鍥為��...\\n\\n");
-            });
-            return true;
+                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");
                 log.error("AI MCP diagnose stopped: error", e);
-                emitter.completeWithError(e);
+                emitter.complete();
             } catch (Exception ignore) {}
-            return true;
         }
     }
 
@@ -395,29 +259,33 @@
         }
     }
 
-    private List<Object> buildOpenAiTools() {
-        if (mcpController == null) return java.util.Collections.emptyList();
-        List<Map<String, Object>> mcpTools = mcpController.listTools();
-        if (mcpTools == null || mcpTools.isEmpty()) return java.util.Collections.emptyList();
-
-        List<Object> tools = new ArrayList<>();
-        for (Map<String, Object> t : mcpTools) {
-            if (t == null) continue;
-            Object name = t.get("name");
-            if (name == null) continue;
-            Object inputSchema = t.get("inputSchema");
-            java.util.LinkedHashMap<String, Object> function = new java.util.LinkedHashMap<String, Object>();
-            function.put("name", String.valueOf(name));
-            Object desc = t.get("description");
-            if (desc != null) function.put("description", String.valueOf(desc));
-            function.put("parameters", inputSchema == null ? new java.util.LinkedHashMap<String, Object>() : inputSchema);
-
-            java.util.LinkedHashMap<String, Object> tool = new java.util.LinkedHashMap<String, Object>();
-            tool.put("type", "function");
-            tool.put("function", function);
-            tools.add(tool);
+    private void emitTokenUsage(SseEmitter emitter, AgentUsageStats usageStats) {
+        if (emitter == null || usageStats == null || usageStats.getTotalTokens() <= 0) {
+            return;
         }
-        return tools;
+        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) {
@@ -536,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