public/liteFlow
1
0
Fork 0
mirror of https://gitee.com/dromara/liteFlow.git synced 2026-06-13 03:11:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(agent): track per-invocation chat usage across ReAct reasoning steps

Browse Source 
Add ChatUsageTrackingHook that accumulates ChatUsage from every
PostReasoningEvent within a single process() call, expose it via
ReActAgentContext#getChatUsage(), and emit a per-step usage line in
ReActLoggingHook. The hook is cached on AgentSession and reset at the
start of each process() so the snapshot reflects the full invocation
(not just the last reasoning step).

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
everywhere.z
2026-05-18 20:50:23 +08:00
parent 8ea01d2e88
commit 083bdfa0f2
6 changed files with 150 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
liteflow-react-agent/liteflow-react-agent-core/src/main/java/com/yomahub/liteflow/agent/component/ReActAgentComponent.java
View File

@@ -1,6 +1,7 @@
package com.yomahub.liteflow.agent.component;
import com.yomahub.liteflow.agent.exception.AgentConfigException;
import com.yomahub.liteflow.agent.hook.ChatUsageTrackingHook;
import com.yomahub.liteflow.agent.hook.ReActLoggingHook;
import com.yomahub.liteflow.agent.skill.SkillBoxFactory;
import com.yomahub.liteflow.agent.skill.SkillLoadResult;
@@ -300,6 +301,7 @@ public abstract class ReActAgentComponent extends NodeComponent {
BuiltAgent built = buildAgent();
agent = built.agent();
session.setSkillTrackingHook(built.skillTrackingHook());
session.setChatUsageTrackingHook(built.chatUsageTrackingHook());
mgr.loadIfExists(session, agent);
session.setAgent(agent);
}
@@ -308,6 +310,11 @@ public abstract class ReActAgentComponent extends NodeComponent {
skillHook.clear();
slot.setAttachment(skillHookKey(), skillHook);
}
ChatUsageTrackingHook usageHook = session.getChatUsageTrackingHook();
if (usageHook != null) {
usageHook.reset();
ctx.setChatUsageTrackingHook(usageHook);
}
Throwable processError = null;
try {
Msg userMsg = Msg.builder().textContent(userPrompt()).build();
@@ -408,7 +415,8 @@ public abstract class ReActAgentComponent extends NodeComponent {
return null;
}
private record BuiltAgent(ReActAgent agent, SkillTrackingHook skillTrackingHook) {
private record BuiltAgent(ReActAgent agent, SkillTrackingHook skillTrackingHook,
ChatUsageTrackingHook chatUsageTrackingHook) {
}
private BuiltAgent buildAgent() {
@@ -430,6 +438,9 @@ public abstract class ReActAgentComponent extends NodeComponent {
allHooks.add(new ReActLoggingHook(ctx.getConversationId() + ":" + ctx.getAgentKey()));
}
ChatUsageTrackingHook chatUsageTrackingHook = new ChatUsageTrackingHook();
allHooks.add(chatUsageTrackingHook);
SkillTrackingHook skillTrackingHook = null;
SkillBox skillBox = null;
if (enableSkills()) {
@@ -452,7 +463,7 @@ public abstract class ReActAgentComponent extends NodeComponent {
builder.skillBox(skillBox);
}
return new BuiltAgent(builder.build(), skillTrackingHook);
return new BuiltAgent(builder.build(), skillTrackingHook, chatUsageTrackingHook);
}
/** 持有单例 AgentSessionManager；首次 process() 时懒创建。 */

24
liteflow-react-agent/liteflow-react-agent-core/src/main/java/com/yomahub/liteflow/agent/component/ReActAgentContext.java
View File

@@ -1,6 +1,8 @@
package com.yomahub.liteflow.agent.component;
import com.yomahub.liteflow.agent.hook.ChatUsageTrackingHook;
import com.yomahub.liteflow.slot.Slot;
import io.agentscope.core.model.ChatUsage;
import java.nio.file.Path;
import java.util.Objects;
@@ -28,6 +30,7 @@ public class ReActAgentContext {
private final String conversationId;
private final String agentKey;
private final Path workspaceDir;
private volatile ChatUsageTrackingHook chatUsageTrackingHook;
public ReActAgentContext(Slot slot, String conversationId, String agentKey, Path workspaceDir) {
this.slot = Objects.requireNonNull(slot, "slot");
@@ -43,4 +46,25 @@ public class ReActAgentContext {
public String getAgentKey() { return agentKey; }
public Path getWorkspaceDir() { return workspaceDir; }
/**
* 由框架注入：本次 {@code process()} 调用使用的 token 累加 hook。
*/
public void setChatUsageTrackingHook(ChatUsageTrackingHook hook) {
this.chatUsageTrackingHook = hook;
}
/**
* 返回本次 {@code process()} 截至当前已累计的 token 用量。
*
* <p>{@link ChatUsage#getInputTokens()} / {@link ChatUsage#getOutputTokens()} /
* {@link ChatUsage#getTotalTokens()} 给出累计 token，{@link ChatUsage#getTime()}
* 给出累计推理耗时（秒）。在 {@code handleReply()} 中调用拿到的就是整次调用的累计值。
*
* @return 累计 ChatUsage；若未观察到任何 usage（模型未上报或 reply 为 null）则返回 {@code null}
*/
public ChatUsage getChatUsage() {
ChatUsageTrackingHook hook = this.chatUsageTrackingHook;
return hook == null ? null : hook.snapshot();
}
}

73
liteflow-react-agent/liteflow-react-agent-core/src/main/java/com/yomahub/liteflow/agent/hook/ChatUsageTrackingHook.java Normal file
View File

@@ -0,0 +1,73 @@
package com.yomahub.liteflow.agent.hook;
import io.agentscope.core.hook.Hook;
import io.agentscope.core.hook.HookEvent;
import io.agentscope.core.hook.PostReasoningEvent;
import io.agentscope.core.message.Msg;
import io.agentscope.core.model.ChatUsage;
import reactor.core.publisher.Mono;
/**
* 累加单次 {@code process()} 调用内所有 reasoning step 的 token 用量。
*
* <p>底层 agentscope 每次 {@code reasoning(iter)} 都新建一个 {@link io.agentscope.core.agent.accumulator.ReasoningContext
* ReasoningContext}，因此 {@link PostReasoningEvent#getReasoningMessage()} 的 metadata 中携带的
* {@link ChatUsage} 只是本步 LLM call 的累计（流式聚合），不是跨多步 ReAct 循环的累计。
* 本 hook 在每次 PostReasoningEvent 触发时把当步 usage 累加到内部计数器，
* 暴露整次调用累计后的 {@link #snapshot()}。
*
* <p>实例与缓存 ReActAgent 同生命周期；每次 {@code process()} 开始前必须调用 {@link #reset()}
* 清零，避免上次调用的余量被带入。
*/
public class ChatUsageTrackingHook implements Hook {
private int inputTokens;
private int outputTokens;
private double time;
private int steps;
@Override
public synchronized <T extends HookEvent> Mono<T> onEvent(T event) {
if (event instanceof PostReasoningEvent e) {
Msg msg = e.getReasoningMessage();
if (msg != null) {
ChatUsage usage = msg.getChatUsage();
if (usage != null) {
inputTokens += usage.getInputTokens();
outputTokens += usage.getOutputTokens();
time += usage.getTime();
steps++;
}
}
}
return Mono.just(event);
}
public synchronized void reset() {
this.inputTokens = 0;
this.outputTokens = 0;
this.time = 0;
this.steps = 0;
}
/**
* 返回到目前为止累计的 token 用量；若尚未观察到任何 usage，返回 {@code null}。
*/
public synchronized ChatUsage snapshot() {
if (steps == 0) {
return null;
}
return ChatUsage.builder()
.inputTokens(inputTokens)
.outputTokens(outputTokens)
.time(time)
.build();
}
/**
* 已经累加过 usage 的 reasoning step 次数。
*/
public synchronized int getSteps() {
return steps;
}
}

10
liteflow-react-agent/liteflow-react-agent-core/src/main/java/com/yomahub/liteflow/agent/hook/ReActLoggingHook.java
View File

@@ -11,6 +11,7 @@ import io.agentscope.core.message.Msg;
import io.agentscope.core.message.ToolResultBlock;
import io.agentscope.core.message.ToolUseBlock;
import io.agentscope.core.message.ThinkingBlock;
import io.agentscope.core.model.ChatUsage;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;
@@ -69,6 +70,15 @@ public class ReActLoggingHook implements Hook {
sessionId, text, summarizeToolUses(tools));
}
}
ChatUsage usage = reply == null ? null : reply.getChatUsage();
if (usage != null) {
LOG.info("[agent:reason][{}] <<< usage input={} output={} total={} time={}s",
sessionId,
usage.getInputTokens(),
usage.getOutputTokens(),
usage.getTotalTokens(),
usage.getTime());
}
} else if (event instanceof PreActingEvent e) {
ToolUseBlock t = e.getToolUse();
LOG.info("[agent:act][{}] >>> tool={} input={}",

10
liteflow-react-agent/liteflow-react-agent-core/src/main/java/com/yomahub/liteflow/agent/session/AgentSession.java
View File

@@ -1,5 +1,6 @@
package com.yomahub.liteflow.agent.session;
import com.yomahub.liteflow.agent.hook.ChatUsageTrackingHook;
import com.yomahub.liteflow.agent.skill.SkillTrackingHook;
import java.nio.file.Path;
@@ -30,6 +31,7 @@ public class AgentSession {
private final ReentrantLock lock = new ReentrantLock();
private volatile Object agent;
private volatile SkillTrackingHook skillTrackingHook;
private volatile ChatUsageTrackingHook chatUsageTrackingHook;
private volatile Instant lastActive = Instant.now();
public AgentSession(String conversationId, String agentKey, String cacheKey, Path workspaceDir) {
@@ -78,6 +80,14 @@ public class AgentSession {
this.skillTrackingHook = skillTrackingHook;
}
public ChatUsageTrackingHook getChatUsageTrackingHook() {
return chatUsageTrackingHook;
}
public void setChatUsageTrackingHook(ChatUsageTrackingHook chatUsageTrackingHook) {
this.chatUsageTrackingHook = chatUsageTrackingHook;
}
public Instant getLastActive() {
return lastActive;
}