zhou zhou
昨天 6e5ff559023efd2d24fdca2adcb7268d06420e46
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
package com.vincent.rsf.server.ai.service.impl.chat;
 
import com.vincent.rsf.framework.exception.CoolException;
import com.vincent.rsf.server.ai.service.AiCallLogService;
import com.vincent.rsf.server.ai.service.MountedToolCallback;
import com.vincent.rsf.server.ai.store.AiCachedToolResult;
import com.vincent.rsf.server.ai.store.AiToolResultStore;
import lombok.RequiredArgsConstructor;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
 
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;
 
@Component
@RequiredArgsConstructor
public class AiToolObservationService {
 
    private final AiToolResultStore aiToolResultStore;
    private final AiCallLogService aiCallLogService;
 
    public ToolCallback[] wrapToolCallbacks(ToolCallback[] toolCallbacks, String requestId,
                                            Long sessionId, AtomicLong toolCallSequence,
                                            AtomicLong toolSuccessCount, AtomicLong toolFailureCount,
                                            Long callLogId, Long userId, Long tenantId,
                                            AiChatTraceEmitter traceEmitter) {
        if (toolCallbacks == null || toolCallbacks.length == 0) {
            return toolCallbacks;
        }
        List<ToolCallback> wrappedCallbacks = new ArrayList<>();
        for (ToolCallback callback : toolCallbacks) {
            if (callback == null) {
                continue;
            }
            wrappedCallbacks.add(new ObservableToolCallback(callback, requestId, sessionId, toolCallSequence,
                    toolSuccessCount, toolFailureCount, callLogId, userId, tenantId, traceEmitter));
        }
        return wrappedCallbacks.toArray(new ToolCallback[0]);
    }
 
    private String summarizeToolPayload(String content, int maxLength) {
        if (!StringUtils.hasText(content)) {
            return null;
        }
        String normalized = content.trim()
                .replace("\r", " ")
                .replace("\n", " ")
                .replaceAll("\\s+", " ");
        return normalized.length() > maxLength ? normalized.substring(0, maxLength) : normalized;
    }
 
    private class ObservableToolCallback implements ToolCallback {
 
        private final ToolCallback delegate;
        private final String requestId;
        private final Long sessionId;
        private final AtomicLong toolCallSequence;
        private final AtomicLong toolSuccessCount;
        private final AtomicLong toolFailureCount;
        private final Long callLogId;
        private final Long userId;
        private final Long tenantId;
        private final AiChatTraceEmitter traceEmitter;
 
        private ObservableToolCallback(ToolCallback delegate, String requestId,
                                       Long sessionId, AtomicLong toolCallSequence,
                                       AtomicLong toolSuccessCount, AtomicLong toolFailureCount,
                                       Long callLogId, Long userId, Long tenantId,
                                       AiChatTraceEmitter traceEmitter) {
            this.delegate = delegate;
            this.requestId = requestId;
            this.sessionId = sessionId;
            this.toolCallSequence = toolCallSequence;
            this.toolSuccessCount = toolSuccessCount;
            this.toolFailureCount = toolFailureCount;
            this.callLogId = callLogId;
            this.userId = userId;
            this.tenantId = tenantId;
            this.traceEmitter = traceEmitter;
        }
 
        @Override
        public org.springframework.ai.tool.definition.ToolDefinition getToolDefinition() {
            return delegate.getToolDefinition();
        }
 
        @Override
        public org.springframework.ai.tool.metadata.ToolMetadata getToolMetadata() {
            return delegate.getToolMetadata();
        }
 
        @Override
        public String call(String toolInput) {
            return call(toolInput, null);
        }
 
        @Override
        public String call(String toolInput, ToolContext toolContext) {
            String toolName = delegate.getToolDefinition() == null ? "unknown" : delegate.getToolDefinition().name();
            String mountName = delegate instanceof MountedToolCallback ? ((MountedToolCallback) delegate).getMountName() : null;
            String toolCallId = requestId + "-tool-" + toolCallSequence.incrementAndGet();
            long startedAt = System.currentTimeMillis();
            String inputSummary = summarizeToolPayload(toolInput, 400);
            AiCachedToolResult cachedToolResult = aiToolResultStore.getToolResult(tenantId, requestId, toolName, toolInput);
            if (cachedToolResult != null) {
                String outputSummary = summarizeToolPayload(cachedToolResult.getOutput(), 600);
                String errorMessage = cachedToolResult.getErrorMessage();
                if (traceEmitter != null) {
                    traceEmitter.onToolResult(toolName, mountName, toolCallId, inputSummary, outputSummary,
                            errorMessage, 0L, System.currentTimeMillis(), !cachedToolResult.isSuccess());
                }
                if (cachedToolResult.isSuccess()) {
                    toolSuccessCount.incrementAndGet();
                    aiCallLogService.saveMcpCallLog(callLogId, requestId, sessionId, toolCallId, mountName, toolName,
                            "COMPLETED", inputSummary, outputSummary,
                            null, 0L, userId, tenantId);
                    return cachedToolResult.getOutput();
                }
                toolFailureCount.incrementAndGet();
                aiCallLogService.saveMcpCallLog(callLogId, requestId, sessionId, toolCallId, mountName, toolName,
                        "FAILED", inputSummary, null, errorMessage,
                        0L, userId, tenantId);
                throw new CoolException(errorMessage);
            }
            if (traceEmitter != null) {
                traceEmitter.onToolStart(toolName, mountName, toolCallId, inputSummary, startedAt);
            }
            try {
                String output = toolContext == null ? delegate.call(toolInput) : delegate.call(toolInput, toolContext);
                long durationMs = System.currentTimeMillis() - startedAt;
                String outputSummary = summarizeToolPayload(output, 600);
                if (traceEmitter != null) {
                    traceEmitter.onToolResult(toolName, mountName, toolCallId, inputSummary, outputSummary,
                            null, durationMs, System.currentTimeMillis(), false);
                }
                aiToolResultStore.cacheToolResult(tenantId, requestId, toolName, toolInput, true, output, null);
                toolSuccessCount.incrementAndGet();
                aiCallLogService.saveMcpCallLog(callLogId, requestId, sessionId, toolCallId, mountName, toolName,
                        "COMPLETED", inputSummary, outputSummary,
                        null, durationMs, userId, tenantId);
                return output;
            } catch (RuntimeException e) {
                long durationMs = System.currentTimeMillis() - startedAt;
                String errorMessage = e.getMessage();
                if (traceEmitter != null) {
                    traceEmitter.onToolResult(toolName, mountName, toolCallId, inputSummary, null,
                            errorMessage, durationMs, System.currentTimeMillis(), true);
                }
                aiToolResultStore.cacheToolResult(tenantId, requestId, toolName, toolInput, false, null, errorMessage);
                toolFailureCount.incrementAndGet();
                aiCallLogService.saveMcpCallLog(callLogId, requestId, sessionId, toolCallId, mountName, toolName,
                        "FAILED", inputSummary, null, errorMessage,
                        durationMs, userId, tenantId);
                throw e;
            }
        }
    }
}