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;
|
}
|
}
|
}
|
}
|
