mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-05-24 00:45:57 +00:00
feat: refine token usage display modes (#2329)
* feat: refine token usage display modes * docs: clarify token usage accounting semantics * fix: avoid duplicate subtask debug keys * style: format token usage tests * chore: address token attribution review feedback * Update test_token_usage_middleware.py * Update test_token_usage_middleware.py * chore: simplify token attribution fallback * fix token usage metadata follow-up handling --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
YuJitang
@@ -0,0 +1,440 @@
|
||||
import type { Message } from "@langchain/langgraph-sdk";
|
||||
|
||||
import type { Translations } from "@/core/i18n/locales/types";
|
||||
|
||||
import { getUsageMetadata, type TokenUsage } from "./usage";
|
||||
import { hasContent } from "./utils";
|
||||
|
||||
export type TokenUsageInlineMode = "off" | "per_turn" | "step_debug";
|
||||
|
||||
export interface TokenUsagePreferences {
|
||||
headerTotal: boolean;
|
||||
inlineMode: TokenUsageInlineMode;
|
||||
}
|
||||
|
||||
export type TokenUsageViewPreset = "off" | "summary" | "per_turn" | "debug";
|
||||
|
||||
export interface TokenDebugStep {
|
||||
id: string;
|
||||
messageId: string;
|
||||
label: string;
|
||||
secondaryLabels: string[];
|
||||
usage: TokenUsage | null;
|
||||
sharedAttribution: boolean;
|
||||
}
|
||||
|
||||
type TokenUsageAttributionAction =
|
||||
| {
|
||||
kind: "todo_start" | "todo_complete" | "todo_update" | "todo_remove";
|
||||
content?: string;
|
||||
tool_call_id?: string;
|
||||
}
|
||||
| {
|
||||
kind: "subagent";
|
||||
description?: string | null;
|
||||
subagent_type?: string | null;
|
||||
tool_call_id?: string;
|
||||
}
|
||||
| {
|
||||
kind: "search";
|
||||
query?: string | null;
|
||||
tool_name?: string | null;
|
||||
tool_call_id?: string;
|
||||
}
|
||||
| {
|
||||
kind: "present_files" | "clarification";
|
||||
tool_call_id?: string;
|
||||
}
|
||||
| {
|
||||
kind: "tool";
|
||||
tool_name?: string | null;
|
||||
description?: string | null;
|
||||
tool_call_id?: string;
|
||||
};
|
||||
|
||||
interface TokenUsageAttribution {
|
||||
version?: number;
|
||||
kind?:
|
||||
| "thinking"
|
||||
| "final_answer"
|
||||
| "tool_batch"
|
||||
| "todo_update"
|
||||
| "subagent_dispatch";
|
||||
shared_attribution?: boolean;
|
||||
tool_call_ids?: string[];
|
||||
actions?: TokenUsageAttributionAction[];
|
||||
}
|
||||
|
||||
// Precise write_todos labels come from the backend attribution payload.
|
||||
// The frontend fallback intentionally stays generic so we do not duplicate
|
||||
// backend/packages/harness/deerflow/agents/middlewares/token_usage_middleware.py
|
||||
//::_build_todo_actions and risk the two diffing algorithms drifting apart.
|
||||
|
||||
export function getTokenUsageViewPreset(
|
||||
preferences: TokenUsagePreferences,
|
||||
): TokenUsageViewPreset {
|
||||
if (!preferences.headerTotal && preferences.inlineMode === "off") {
|
||||
return "off";
|
||||
}
|
||||
if (preferences.headerTotal && preferences.inlineMode === "off") {
|
||||
return "summary";
|
||||
}
|
||||
if (preferences.inlineMode === "step_debug") {
|
||||
return "debug";
|
||||
}
|
||||
return "per_turn";
|
||||
}
|
||||
|
||||
export function tokenUsagePreferencesFromPreset(
|
||||
preset: TokenUsageViewPreset,
|
||||
): TokenUsagePreferences {
|
||||
switch (preset) {
|
||||
case "off":
|
||||
return { headerTotal: false, inlineMode: "off" };
|
||||
case "summary":
|
||||
return { headerTotal: true, inlineMode: "off" };
|
||||
case "debug":
|
||||
return { headerTotal: true, inlineMode: "step_debug" };
|
||||
case "per_turn":
|
||||
default:
|
||||
return { headerTotal: true, inlineMode: "per_turn" };
|
||||
}
|
||||
}
|
||||
|
||||
export function buildTokenDebugSteps(
|
||||
messages: Message[],
|
||||
t: Translations,
|
||||
): TokenDebugStep[] {
|
||||
const steps: TokenDebugStep[] = [];
|
||||
|
||||
for (const [index, message] of messages.entries()) {
|
||||
if (message.type !== "ai") {
|
||||
continue;
|
||||
}
|
||||
|
||||
const usage = getUsageMetadata(message);
|
||||
const attribution = getTokenUsageAttribution(message);
|
||||
const actionLabels: string[] = [];
|
||||
|
||||
if (attribution) {
|
||||
actionLabels.push(...buildActionLabelsFromAttribution(attribution, t));
|
||||
|
||||
if (actionLabels.length === 0) {
|
||||
if (attribution.kind === "final_answer") {
|
||||
actionLabels.push(t.tokenUsage.finalAnswer);
|
||||
} else if (attribution.kind === "thinking") {
|
||||
actionLabels.push(t.common.thinking);
|
||||
}
|
||||
}
|
||||
|
||||
if (actionLabels.length > 0) {
|
||||
const sharedAttribution =
|
||||
attribution.shared_attribution ?? actionLabels.length > 1;
|
||||
steps.push({
|
||||
id: message.id ?? `token-step-${index}`,
|
||||
messageId: message.id ?? `token-step-${index}`,
|
||||
label:
|
||||
sharedAttribution && actionLabels.length > 1
|
||||
? t.tokenUsage.stepTotal
|
||||
: actionLabels[0]!,
|
||||
secondaryLabels:
|
||||
sharedAttribution && actionLabels.length > 1 ? actionLabels : [],
|
||||
usage,
|
||||
sharedAttribution,
|
||||
});
|
||||
continue;
|
||||
}
|
||||
}
|
||||
|
||||
for (const toolCall of message.tool_calls ?? []) {
|
||||
const toolArgs = (toolCall.args ?? {}) as Record<string, unknown>;
|
||||
|
||||
if (toolCall.name === "write_todos") {
|
||||
actionLabels.push(t.toolCalls.writeTodos);
|
||||
continue;
|
||||
}
|
||||
|
||||
actionLabels.push(
|
||||
describeToolCall(
|
||||
{
|
||||
name: toolCall.name,
|
||||
args: toolArgs,
|
||||
},
|
||||
t,
|
||||
),
|
||||
);
|
||||
}
|
||||
|
||||
if (actionLabels.length === 0) {
|
||||
if (hasContent(message)) {
|
||||
actionLabels.push(t.tokenUsage.finalAnswer);
|
||||
} else {
|
||||
actionLabels.push(t.common.thinking);
|
||||
}
|
||||
}
|
||||
|
||||
steps.push({
|
||||
id: message.id ?? `token-step-${index}`,
|
||||
messageId: message.id ?? `token-step-${index}`,
|
||||
label:
|
||||
actionLabels.length === 1 ? actionLabels[0]! : t.tokenUsage.stepTotal,
|
||||
secondaryLabels: actionLabels.length > 1 ? actionLabels : [],
|
||||
usage,
|
||||
sharedAttribution: actionLabels.length > 1,
|
||||
});
|
||||
}
|
||||
|
||||
return steps;
|
||||
}
|
||||
|
||||
function getTokenUsageAttribution(
|
||||
message: Message,
|
||||
): TokenUsageAttribution | null {
|
||||
if (message.type !== "ai") {
|
||||
return null;
|
||||
}
|
||||
|
||||
const additionalKwargs = message.additional_kwargs;
|
||||
if (!additionalKwargs || typeof additionalKwargs !== "object") {
|
||||
return null;
|
||||
}
|
||||
|
||||
const attribution = (additionalKwargs as Record<string, unknown>)
|
||||
.token_usage_attribution;
|
||||
const normalized = normalizeTokenUsageAttribution(attribution);
|
||||
if (!normalized) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return normalized;
|
||||
}
|
||||
|
||||
function buildActionLabelsFromAttribution(
|
||||
attribution: TokenUsageAttribution,
|
||||
t: Translations,
|
||||
): string[] {
|
||||
return (attribution.actions ?? [])
|
||||
.map((action) => describeAttributionAction(action, t))
|
||||
.filter((label): label is string => !!label);
|
||||
}
|
||||
|
||||
function describeAttributionAction(
|
||||
action: TokenUsageAttributionAction,
|
||||
t: Translations,
|
||||
): string | null {
|
||||
switch (action.kind) {
|
||||
case "todo_start":
|
||||
return action.content
|
||||
? t.tokenUsage.startTodo(action.content)
|
||||
: t.toolCalls.writeTodos;
|
||||
case "todo_complete":
|
||||
return action.content
|
||||
? t.tokenUsage.completeTodo(action.content)
|
||||
: t.toolCalls.writeTodos;
|
||||
case "todo_update":
|
||||
return action.content
|
||||
? t.tokenUsage.updateTodo(action.content)
|
||||
: t.toolCalls.writeTodos;
|
||||
case "todo_remove":
|
||||
return action.content
|
||||
? t.tokenUsage.removeTodo(action.content)
|
||||
: t.toolCalls.writeTodos;
|
||||
case "subagent":
|
||||
return t.tokenUsage.subagent(action.description ?? t.subtasks.subtask);
|
||||
case "search":
|
||||
if (action.query) {
|
||||
return t.toolCalls.searchFor(action.query);
|
||||
}
|
||||
return t.toolCalls.useTool(action.tool_name ?? "search");
|
||||
case "present_files":
|
||||
return t.toolCalls.presentFiles;
|
||||
case "clarification":
|
||||
return t.toolCalls.needYourHelp;
|
||||
case "tool":
|
||||
return describeToolCall(
|
||||
{
|
||||
name: action.tool_name ?? "tool",
|
||||
args: action.description ? { description: action.description } : {},
|
||||
},
|
||||
t,
|
||||
);
|
||||
default:
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
function describeToolCall(
|
||||
toolCall: {
|
||||
name: string;
|
||||
args: Record<string, unknown>;
|
||||
},
|
||||
t: Translations,
|
||||
): string {
|
||||
if (toolCall.name === "task") {
|
||||
const description =
|
||||
typeof toolCall.args.description === "string"
|
||||
? toolCall.args.description
|
||||
: t.subtasks.subtask;
|
||||
return t.tokenUsage.subagent(description);
|
||||
}
|
||||
|
||||
if (
|
||||
(toolCall.name === "web_search" || toolCall.name === "image_search") &&
|
||||
typeof toolCall.args.query === "string"
|
||||
) {
|
||||
return t.toolCalls.searchFor(toolCall.args.query);
|
||||
}
|
||||
|
||||
if (toolCall.name === "web_fetch") {
|
||||
return t.toolCalls.viewWebPage;
|
||||
}
|
||||
|
||||
if (toolCall.name === "present_files") {
|
||||
return t.toolCalls.presentFiles;
|
||||
}
|
||||
|
||||
if (toolCall.name === "ask_clarification") {
|
||||
return t.toolCalls.needYourHelp;
|
||||
}
|
||||
|
||||
if (typeof toolCall.args.description === "string") {
|
||||
return toolCall.args.description;
|
||||
}
|
||||
|
||||
return t.toolCalls.useTool(toolCall.name);
|
||||
}
|
||||
|
||||
function normalizeTokenUsageAttribution(
|
||||
value: unknown,
|
||||
): TokenUsageAttribution | null {
|
||||
const record = asRecord(value);
|
||||
if (!record) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const rawActions = record.actions;
|
||||
if (rawActions !== undefined && !Array.isArray(rawActions)) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return {
|
||||
// Versioning is additive for now: the frontend should ignore unknown
|
||||
// fields and fall back when required fields become incompatible.
|
||||
version: typeof record.version === "number" ? record.version : undefined,
|
||||
kind: isTokenUsageAttributionKind(record.kind) ? record.kind : undefined,
|
||||
shared_attribution:
|
||||
typeof record.shared_attribution === "boolean"
|
||||
? record.shared_attribution
|
||||
: undefined,
|
||||
tool_call_ids: Array.isArray(record.tool_call_ids)
|
||||
? record.tool_call_ids.filter(
|
||||
(toolCallId): toolCallId is string =>
|
||||
typeof toolCallId === "string" && toolCallId.trim().length > 0,
|
||||
)
|
||||
: undefined,
|
||||
actions: Array.isArray(rawActions)
|
||||
? rawActions
|
||||
.map((action) => normalizeTokenUsageAttributionAction(action))
|
||||
.filter(
|
||||
(action): action is TokenUsageAttributionAction => action !== null,
|
||||
)
|
||||
: undefined,
|
||||
};
|
||||
}
|
||||
|
||||
function normalizeTokenUsageAttributionAction(
|
||||
value: unknown,
|
||||
): TokenUsageAttributionAction | null {
|
||||
const record = asRecord(value);
|
||||
if (!record) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const kind = record.kind;
|
||||
if (
|
||||
kind !== "todo_start" &&
|
||||
kind !== "todo_complete" &&
|
||||
kind !== "todo_update" &&
|
||||
kind !== "todo_remove" &&
|
||||
kind !== "subagent" &&
|
||||
kind !== "search" &&
|
||||
kind !== "present_files" &&
|
||||
kind !== "clarification" &&
|
||||
kind !== "tool"
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const content = readString(record.content);
|
||||
const toolCallId = readString(record.tool_call_id);
|
||||
|
||||
switch (kind) {
|
||||
case "todo_start":
|
||||
case "todo_complete":
|
||||
case "todo_update":
|
||||
case "todo_remove":
|
||||
return {
|
||||
kind,
|
||||
content,
|
||||
tool_call_id: toolCallId,
|
||||
};
|
||||
case "subagent":
|
||||
return {
|
||||
kind,
|
||||
description: readString(record.description),
|
||||
subagent_type: readString(record.subagent_type),
|
||||
tool_call_id: toolCallId,
|
||||
};
|
||||
case "search":
|
||||
return {
|
||||
kind,
|
||||
query: readString(record.query),
|
||||
tool_name: readString(record.tool_name),
|
||||
tool_call_id: toolCallId,
|
||||
};
|
||||
case "present_files":
|
||||
case "clarification":
|
||||
return {
|
||||
kind,
|
||||
tool_call_id: toolCallId,
|
||||
};
|
||||
case "tool":
|
||||
return {
|
||||
kind,
|
||||
tool_name: readString(record.tool_name),
|
||||
description: readString(record.description),
|
||||
tool_call_id: toolCallId,
|
||||
};
|
||||
default:
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
function asRecord(value: unknown): Record<string, unknown> | null {
|
||||
if (!value || typeof value !== "object" || Array.isArray(value)) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return value as Record<string, unknown>;
|
||||
}
|
||||
|
||||
function readString(value: unknown): string | undefined {
|
||||
if (typeof value !== "string") {
|
||||
return undefined;
|
||||
}
|
||||
|
||||
const normalized = value.trim();
|
||||
return normalized.length > 0 ? normalized : undefined;
|
||||
}
|
||||
|
||||
function isTokenUsageAttributionKind(
|
||||
value: unknown,
|
||||
): value is NonNullable<TokenUsageAttribution["kind"]> {
|
||||
return (
|
||||
value === "thinking" ||
|
||||
value === "final_answer" ||
|
||||
value === "tool_batch" ||
|
||||
value === "todo_update" ||
|
||||
value === "subagent_dispatch"
|
||||
);
|
||||
}
|
||||
@@ -18,7 +18,7 @@ interface AssistantClarificationGroup extends GenericMessageGroup<"assistant:cla
|
||||
|
||||
interface AssistantSubagentGroup extends GenericMessageGroup<"assistant:subagent"> {}
|
||||
|
||||
type MessageGroup =
|
||||
export type MessageGroup =
|
||||
| HumanMessageGroup
|
||||
| AssistantProcessingGroup
|
||||
| AssistantMessageGroup
|
||||
@@ -26,10 +26,7 @@ type MessageGroup =
|
||||
| AssistantClarificationGroup
|
||||
| AssistantSubagentGroup;
|
||||
|
||||
export function groupMessages<T>(
|
||||
messages: Message[],
|
||||
mapper: (group: MessageGroup) => T,
|
||||
): T[] {
|
||||
export function getMessageGroups(messages: Message[]): MessageGroup[] {
|
||||
if (messages.length === 0) {
|
||||
return [];
|
||||
}
|
||||
@@ -124,11 +121,52 @@ export function groupMessages<T>(
|
||||
}
|
||||
}
|
||||
|
||||
return groups
|
||||
return groups;
|
||||
}
|
||||
|
||||
export function groupMessages<T>(
|
||||
messages: Message[],
|
||||
mapper: (group: MessageGroup) => T,
|
||||
): T[] {
|
||||
return getMessageGroups(messages)
|
||||
.map(mapper)
|
||||
.filter((result) => result !== undefined && result !== null) as T[];
|
||||
}
|
||||
|
||||
export function getAssistantTurnUsageMessages(groups: MessageGroup[]) {
|
||||
const usageMessagesByGroupIndex: Array<Message[] | null> = Array.from(
|
||||
{ length: groups.length },
|
||||
() => null,
|
||||
);
|
||||
|
||||
let turnStartIndex: number | null = null;
|
||||
|
||||
for (const [index, group] of groups.entries()) {
|
||||
if (group.type === "human") {
|
||||
turnStartIndex = null;
|
||||
continue;
|
||||
}
|
||||
|
||||
turnStartIndex ??= index;
|
||||
|
||||
const nextGroup = groups[index + 1];
|
||||
const isTurnEnd = !nextGroup || nextGroup.type === "human";
|
||||
|
||||
if (!isTurnEnd) {
|
||||
continue;
|
||||
}
|
||||
|
||||
usageMessagesByGroupIndex[index] = groups
|
||||
.slice(turnStartIndex, index + 1)
|
||||
.flatMap((currentGroup) => currentGroup.messages)
|
||||
.filter((message) => message.type === "ai");
|
||||
|
||||
turnStartIndex = null;
|
||||
}
|
||||
|
||||
return usageMessagesByGroupIndex;
|
||||
}
|
||||
|
||||
export function extractTextFromMessage(message: Message) {
|
||||
if (typeof message.content === "string") {
|
||||
return (
|
||||
|
||||
Reference in New Issue
Block a user