zhuyongxin
a3abe3f7a2
- AgentLoggingHook 改为从 config.metadata 读取 sessionId（线程安全） - 移除 AgentLoggingHook 调试用的 metadata 日志 - TokenTrackingChatModel 日志降为 debug - SessionContextHolder 移除未使用的 setAgentName/getAgentName - ChatService 清理无用 import - 修复 stream 路径下 ThreadLocal NPE
51 lines
1.5 KiB
Java
51 lines
1.5 KiB
Java
package com.superbiz.agent.hook;
|
|
|
|
import org.slf4j.Logger;
|
|
import org.slf4j.LoggerFactory;
|
|
import org.springframework.ai.chat.model.ChatModel;
|
|
import org.springframework.ai.chat.model.ChatResponse;
|
|
import org.springframework.ai.chat.prompt.Prompt;
|
|
import reactor.core.publisher.Flux;
|
|
|
|
/**
|
|
* ChatModel 包装器 — 捕获每次模型调用的实际 token 用量
|
|
* 通过 TokenUsageHolder 传递给 AgentLoggingHook
|
|
*/
|
|
public class TokenTrackingChatModel implements ChatModel {
|
|
|
|
private static final Logger log = LoggerFactory.getLogger(TokenTrackingChatModel.class);
|
|
|
|
private final ChatModel delegate;
|
|
|
|
public TokenTrackingChatModel(ChatModel delegate) {
|
|
this.delegate = delegate;
|
|
}
|
|
|
|
@Override
|
|
public ChatResponse call(Prompt prompt) {
|
|
ChatResponse response = delegate.call(prompt);
|
|
captureTokenUsage(response);
|
|
return response;
|
|
}
|
|
|
|
@Override
|
|
public Flux<ChatResponse> stream(Prompt prompt) {
|
|
return delegate.stream(prompt);
|
|
}
|
|
|
|
private void captureTokenUsage(ChatResponse response) {
|
|
try {
|
|
if (response.getMetadata() == null || response.getMetadata().getUsage() == null) {
|
|
return;
|
|
}
|
|
var usage = response.getMetadata().getUsage();
|
|
Integer total = usage.getTotalTokens();
|
|
if (total != null && total > 0) {
|
|
TokenUsageHolder.set(total);
|
|
}
|
|
} catch (Exception e) {
|
|
log.debug("捕获 token 用量失败", e);
|
|
}
|
|
}
|
|
}
|