feat: show token usage per assistant response (#2270)

* feat: show token usage per assistant response

* fix: align client models response with token usage

* fix: address token usage review feedback

* docs: clarify token usage config example

---------

Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
YuJitang 2026-04-16 08:56:49 +08:00 committed by GitHub
parent 0e16a7fe55
commit 105db00987
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
17 changed files with 271 additions and 50 deletions

View File

@ -17,10 +17,17 @@ class ModelResponse(BaseModel):
supports_reasoning_effort: bool = Field(default=False, description="Whether model supports reasoning effort")
class TokenUsageResponse(BaseModel):
"""Token usage display configuration."""
enabled: bool = Field(default=False, description="Whether token usage display is enabled")
class ModelsListResponse(BaseModel):
"""Response model for listing all models."""
models: list[ModelResponse]
token_usage: TokenUsageResponse
@router.get(
@ -36,7 +43,7 @@ async def list_models() -> ModelsListResponse:
excluding sensitive fields like API keys and internal configuration.
Returns:
A list of all configured models with their metadata.
A list of all configured models with their metadata and token usage display settings.
Example Response:
```json
@ -44,17 +51,24 @@ async def list_models() -> ModelsListResponse:
"models": [
{
"name": "gpt-4",
"model": "gpt-4",
"display_name": "GPT-4",
"description": "OpenAI GPT-4 model",
"supports_thinking": false
"supports_thinking": false,
"supports_reasoning_effort": false
},
{
"name": "claude-3-opus",
"model": "claude-3-opus",
"display_name": "Claude 3 Opus",
"description": "Anthropic Claude 3 Opus model",
"supports_thinking": true
"supports_thinking": true,
"supports_reasoning_effort": false
}
]
],
"token_usage": {
"enabled": true
}
}
```
"""
@ -70,7 +84,10 @@ async def list_models() -> ModelsListResponse:
)
for model in config.models
]
return ModelsListResponse(models=models)
return ModelsListResponse(
models=models,
token_usage=TokenUsageResponse(enabled=config.token_usage.enabled),
)
@router.get(

View File

@ -722,6 +722,10 @@ class DeerFlowClient:
Dict with "models" key containing list of model info dicts,
matching the Gateway API ``ModelsListResponse`` schema.
"""
token_usage_enabled = getattr(getattr(self._app_config, "token_usage", None), "enabled", False)
if not isinstance(token_usage_enabled, bool):
token_usage_enabled = False
return {
"models": [
{
@ -733,7 +737,8 @@ class DeerFlowClient:
"supports_reasoning_effort": getattr(model, "supports_reasoning_effort", False),
}
for model in self._app_config.models
]
],
"token_usage": {"enabled": token_usage_enabled},
}
def list_skills(self, enabled_only: bool = False) -> dict:

View File

@ -38,6 +38,7 @@ def mock_app_config():
config = MagicMock()
config.models = [model]
config.token_usage.enabled = False
return config
@ -107,6 +108,7 @@ class TestConfigQueries:
def test_list_models(self, client):
result = client.list_models()
assert "models" in result
assert result["token_usage"] == {"enabled": False}
assert len(result["models"]) == 1
assert result["models"][0]["name"] == "test-model"
# Verify Gateway-aligned fields are present
@ -2196,7 +2198,9 @@ class TestGatewayConformance:
model.display_name = "Test Model"
model.description = "A test model"
model.supports_thinking = False
model.supports_reasoning_effort = False
mock_app_config.models = [model]
mock_app_config.token_usage.enabled = True
with patch("deerflow.client.get_app_config", return_value=mock_app_config):
client = DeerFlowClient()
@ -2206,6 +2210,7 @@ class TestGatewayConformance:
assert len(parsed.models) == 1
assert parsed.models[0].name == "test-model"
assert parsed.models[0].model == "gpt-test"
assert parsed.token_usage.enabled is True
def test_get_model(self, mock_app_config):
model = MagicMock()

View File

@ -21,10 +21,11 @@ config_version: 7
log_level: info
# ============================================================================
# Token Usage Tracking
# Token Usage
# ============================================================================
# Track LLM token usage per model call (input/output/total tokens)
# Logs at info level via TokenUsageMiddleware
# Enable token usage collection and display.
# When enabled, DeerFlow records input/output/total tokens per model call
# and shows usage metadata in the workspace UI when providers return it.
token_usage:
enabled: false

View File

@ -23,6 +23,7 @@ import { TokenUsageIndicator } from "@/components/workspace/token-usage-indicato
import { Tooltip } from "@/components/workspace/tooltip";
import { useAgent } from "@/core/agents";
import { useI18n } from "@/core/i18n/hooks";
import { useModels } from "@/core/models/hooks";
import { useNotification } from "@/core/notification/hooks";
import { useThreadSettings } from "@/core/settings";
import { useThreadStream } from "@/core/threads/hooks";
@ -44,6 +45,7 @@ export default function AgentChatPage() {
const { threadId, setThreadId, isNewThread, setIsNewThread } =
useThreadChat();
const [settings, setSettings] = useThreadSettings(threadId);
const { tokenUsageEnabled } = useModels();
const { showNotification } = useNotification();
const [thread, sendMessage] = useThreadStream({
@ -128,7 +130,10 @@ export default function AgentChatPage() {
<PlusSquare /> {t.agents.newChat}
</Button>
</Tooltip>
<TokenUsageIndicator messages={thread.messages} />
<TokenUsageIndicator
enabled={tokenUsageEnabled}
messages={thread.messages}
/>
<ExportTrigger threadId={threadId} />
<ArtifactTrigger />
</div>
@ -141,6 +146,7 @@ export default function AgentChatPage() {
threadId={threadId}
thread={thread}
paddingBottom={messageListPaddingBottom}
tokenUsageEnabled={tokenUsageEnabled}
/>
</div>

View File

@ -22,6 +22,7 @@ import { TodoList } from "@/components/workspace/todo-list";
import { TokenUsageIndicator } from "@/components/workspace/token-usage-indicator";
import { Welcome } from "@/components/workspace/welcome";
import { useI18n } from "@/core/i18n/hooks";
import { useModels } from "@/core/models/hooks";
import { useNotification } from "@/core/notification/hooks";
import { useThreadSettings } from "@/core/settings";
import { useThreadStream } from "@/core/threads/hooks";
@ -36,6 +37,7 @@ export default function ChatPage() {
useThreadChat();
const [settings, setSettings] = useThreadSettings(threadId);
const [mounted, setMounted] = useState(false);
const { tokenUsageEnabled } = useModels();
useSpecificChatMode();
useEffect(() => {
@ -103,7 +105,10 @@ export default function ChatPage() {
<ThreadTitle threadId={threadId} thread={thread} />
</div>
<div className="flex items-center gap-2">
<TokenUsageIndicator messages={thread.messages} />
<TokenUsageIndicator
enabled={tokenUsageEnabled}
messages={thread.messages}
/>
<ExportTrigger threadId={threadId} />
<ArtifactTrigger />
</div>
@ -115,6 +120,7 @@ export default function ChatPage() {
threadId={threadId}
thread={thread}
paddingBottom={messageListPaddingBottom}
tokenUsageEnabled={tokenUsageEnabled}
/>
</div>
<div className="absolute right-0 bottom-0 left-0 z-30 flex justify-center px-4">

View File

@ -38,17 +38,20 @@ import { cn } from "@/lib/utils";
import { CopyButton } from "../copy-button";
import { MarkdownContent } from "./markdown-content";
import { MessageTokenUsage } from "./message-token-usage";
export function MessageListItem({
className,
message,
isLoading,
threadId,
tokenUsageEnabled = false,
}: {
className?: string;
message: Message;
isLoading?: boolean;
threadId: string;
tokenUsageEnabled?: boolean;
}) {
const isHuman = message.type === "human";
return (
@ -61,6 +64,7 @@ export function MessageListItem({
message={message}
isLoading={isLoading}
threadId={threadId}
tokenUsageEnabled={tokenUsageEnabled}
/>
{!isLoading && (
<MessageToolbar
@ -119,11 +123,13 @@ function MessageContent_({
message,
isLoading = false,
threadId,
tokenUsageEnabled = false,
}: {
className?: string;
message: Message;
isLoading?: boolean;
threadId: string;
tokenUsageEnabled?: boolean;
}) {
const rehypePlugins = useRehypeSplitWordsIntoSpans(isLoading);
const isHuman = message.type === "human";
@ -201,6 +207,11 @@ function MessageContent_({
<ReasoningTrigger />
<ReasoningContent>{reasoningContent}</ReasoningContent>
</Reasoning>
<MessageTokenUsage
enabled={tokenUsageEnabled}
isLoading={isLoading}
message={message}
/>
</AIElementMessageContent>
);
}
@ -238,6 +249,11 @@ function MessageContent_({
className="my-3"
components={components}
/>
<MessageTokenUsage
enabled={tokenUsageEnabled}
isLoading={isLoading}
message={message}
/>
</AIElementMessageContent>
);
}

View File

@ -13,6 +13,7 @@ import {
hasContent,
hasPresentFiles,
hasReasoning,
hasToolCalls,
} from "@/core/messages/utils";
import { useRehypeSplitWordsIntoSpans } from "@/core/rehype";
import type { Subtask } from "@/core/tasks";
@ -26,6 +27,7 @@ import { StreamingIndicator } from "../streaming-indicator";
import { MarkdownContent } from "./markdown-content";
import { MessageGroup } from "./message-group";
import { MessageListItem } from "./message-list-item";
import { MessageTokenUsageList } from "./message-token-usage";
import { MessageListSkeleton } from "./skeleton";
import { SubtaskCard } from "./subtask-card";
@ -37,11 +39,13 @@ export function MessageList({
threadId,
thread,
paddingBottom = MESSAGE_LIST_DEFAULT_PADDING_BOTTOM,
tokenUsageEnabled = false,
}: {
className?: string;
threadId: string;
thread: BaseStream<AgentThreadState>;
paddingBottom?: number;
tokenUsageEnabled?: boolean;
}) {
const { t } = useI18n();
const rehypePlugins = useRehypeSplitWordsIntoSpans(thread.isLoading);
@ -64,6 +68,7 @@ export function MessageList({
message={msg}
isLoading={thread.isLoading}
threadId={threadId}
tokenUsageEnabled={tokenUsageEnabled}
/>
);
});
@ -71,12 +76,18 @@ export function MessageList({
const message = group.messages[0];
if (message && hasContent(message)) {
return (
<MarkdownContent
key={group.id}
content={extractContentFromMessage(message)}
isLoading={thread.isLoading}
rehypePlugins={rehypePlugins}
/>
<div key={group.id} className="w-full">
<MarkdownContent
content={extractContentFromMessage(message)}
isLoading={thread.isLoading}
rehypePlugins={rehypePlugins}
/>
<MessageTokenUsageList
enabled={tokenUsageEnabled}
isLoading={thread.isLoading}
messages={group.messages}
/>
</div>
);
}
return null;
@ -99,6 +110,11 @@ export function MessageList({
/>
)}
<ArtifactFileList files={files} threadId={threadId} />
<MessageTokenUsageList
enabled={tokenUsageEnabled}
isLoading={thread.isLoading}
messages={group.messages}
/>
</div>
);
} else if (group.type === "assistant:subagent") {
@ -191,15 +207,31 @@ export function MessageList({
className="relative z-1 flex flex-col gap-2"
>
{results}
<MessageTokenUsageList
enabled={tokenUsageEnabled}
isLoading={thread.isLoading}
messages={group.messages}
/>
</div>
);
}
const tokenUsageMessages = group.messages.filter(
(message) =>
message.type === "ai" &&
(hasToolCalls(message) ? true : !hasContent(message)),
);
return (
<MessageGroup
key={"group-" + group.id}
messages={group.messages}
isLoading={thread.isLoading}
/>
<div key={"group-" + group.id} className="w-full">
<MessageGroup
messages={group.messages}
isLoading={thread.isLoading}
/>
<MessageTokenUsageList
enabled={tokenUsageEnabled}
isLoading={thread.isLoading}
messages={tokenUsageMessages}
/>
</div>
);
})}
{thread.isLoading && <StreamingIndicator className="my-4" />}

View File

@ -0,0 +1,91 @@
import type { Message } from "@langchain/langgraph-sdk";
import { CoinsIcon } from "lucide-react";
import { useI18n } from "@/core/i18n/hooks";
import { formatTokenCount, getUsageMetadata } from "@/core/messages/usage";
import { cn } from "@/lib/utils";
export function MessageTokenUsage({
className,
enabled = false,
isLoading = false,
message,
}: {
className?: string;
enabled?: boolean;
isLoading?: boolean;
message: Message;
}) {
const { t } = useI18n();
if (!enabled || isLoading || message.type !== "ai") {
return null;
}
const usage = getUsageMetadata(message);
return (
<div
className={cn(
"text-muted-foreground border-border/60 mt-1 flex flex-wrap items-center gap-x-3 gap-y-1 border-t pt-2 text-[11px]",
className,
)}
>
<span className="inline-flex items-center gap-1 font-medium">
<CoinsIcon className="size-3" />
{t.tokenUsage.label}
</span>
{usage ? (
<>
<span>
{t.tokenUsage.input}: {formatTokenCount(usage.inputTokens)}
</span>
<span>
{t.tokenUsage.output}: {formatTokenCount(usage.outputTokens)}
</span>
<span className="font-medium">
{t.tokenUsage.total}: {formatTokenCount(usage.totalTokens)}
</span>
</>
) : (
<span>{t.tokenUsage.unavailableShort}</span>
)}
</div>
);
}
export function MessageTokenUsageList({
className,
enabled = false,
isLoading = false,
messages,
}: {
className?: string;
enabled?: boolean;
isLoading?: boolean;
messages: Message[];
}) {
if (!enabled || isLoading) {
return null;
}
const aiMessages = messages.filter((message) => message.type === "ai");
if (aiMessages.length === 0) {
return null;
}
return (
<>
{aiMessages.map((message, index) => (
<MessageTokenUsage
className={className}
enabled={enabled}
isLoading={isLoading}
key={message.id ?? index}
message={message}
/>
))}
</>
);
}

View File

@ -15,18 +15,20 @@ import { cn } from "@/lib/utils";
interface TokenUsageIndicatorProps {
messages: Message[];
enabled?: boolean;
className?: string;
}
export function TokenUsageIndicator({
messages,
enabled = false,
className,
}: TokenUsageIndicatorProps) {
const { t } = useI18n();
const usage = useMemo(() => accumulateUsage(messages), [messages]);
if (!usage) {
if (!enabled) {
return null;
}
@ -36,37 +38,49 @@ export function TokenUsageIndicator({
<button
type="button"
className={cn(
"text-muted-foreground flex cursor-default items-center gap-1 text-xs",
"text-muted-foreground bg-background/70 flex cursor-default items-center gap-1.5 rounded-full border px-2 py-1 text-xs",
!usage && "opacity-60",
className,
)}
>
<CoinsIcon size={14} />
<span>{formatTokenCount(usage.totalTokens)}</span>
<span>{t.tokenUsage.label}</span>
<span className="font-mono">
{usage ? formatTokenCount(usage.totalTokens) : "-"}
</span>
</button>
</TooltipTrigger>
<TooltipContent side="bottom" align="end">
<div className="space-y-1 text-xs">
<div className="font-medium">{t.tokenUsage.title}</div>
<div className="flex justify-between gap-4">
<span>{t.tokenUsage.input}</span>
<span className="font-mono">
{formatTokenCount(usage.inputTokens)}
</span>
</div>
<div className="flex justify-between gap-4">
<span>{t.tokenUsage.output}</span>
<span className="font-mono">
{formatTokenCount(usage.outputTokens)}
</span>
</div>
<div className="border-t pt-1">
<div className="flex justify-between gap-4">
<span>{t.tokenUsage.total}</span>
<span className="font-mono font-medium">
{formatTokenCount(usage.totalTokens)}
</span>
{usage ? (
<>
<div className="flex justify-between gap-4">
<span>{t.tokenUsage.input}</span>
<span className="font-mono">
{formatTokenCount(usage.inputTokens)}
</span>
</div>
<div className="flex justify-between gap-4">
<span>{t.tokenUsage.output}</span>
<span className="font-mono">
{formatTokenCount(usage.outputTokens)}
</span>
</div>
<div className="border-t pt-1">
<div className="flex justify-between gap-4">
<span>{t.tokenUsage.total}</span>
<span className="font-mono font-medium">
{formatTokenCount(usage.totalTokens)}
</span>
</div>
</div>
</>
) : (
<div className="text-muted-foreground max-w-56">
{t.tokenUsage.unavailable}
</div>
</div>
)}
</div>
</TooltipContent>
</Tooltip>

View File

@ -298,9 +298,13 @@ export const enUS: Translations = {
// Token Usage
tokenUsage: {
title: "Token Usage",
label: "Tokens",
input: "Input",
output: "Output",
total: "Total",
unavailable:
"No token usage yet. Usage appears only after a successful model response when the provider returns usage_metadata.",
unavailableShort: "No usage returned",
},
// Shortcuts

View File

@ -229,9 +229,12 @@ export interface Translations {
// Token Usage
tokenUsage: {
title: string;
label: string;
input: string;
output: string;
total: string;
unavailable: string;
unavailableShort: string;
};
// Shortcuts

View File

@ -284,9 +284,13 @@ export const zhCN: Translations = {
// Token Usage
tokenUsage: {
title: "Token 用量",
label: "Tokens",
input: "输入",
output: "输出",
total: "总计",
unavailable:
"暂无 Token 用量。只有模型成功返回且供应商提供 usage_metadata 时才会显示。",
unavailableShort: "未返回用量",
},
// Shortcuts

View File

@ -10,7 +10,7 @@ export interface TokenUsage {
* Extract usage_metadata from an AI message if present.
* The field is added by the backend (PR #1218) but not typed in the SDK.
*/
function getUsageMetadata(message: Message): TokenUsage | null {
export function getUsageMetadata(message: Message): TokenUsage | null {
if (message.type !== "ai") {
return null;
}

View File

@ -1,9 +1,12 @@
import { getBackendBaseURL } from "../config";
import type { Model } from "./types";
import type { ModelsResponse } from "./types";
export async function loadModels() {
export async function loadModels(): Promise<ModelsResponse> {
const res = await fetch(`${getBackendBaseURL()}/api/models`);
const { models } = (await res.json()) as { models: Model[] };
return models;
const data = (await res.json()) as Partial<ModelsResponse>;
return {
models: data.models ?? [],
token_usage: data.token_usage ?? { enabled: false },
};
}

View File

@ -9,5 +9,10 @@ export function useModels({ enabled = true }: { enabled?: boolean } = {}) {
enabled,
refetchOnWindowFocus: false,
});
return { models: data ?? [], isLoading, error };
return {
models: data?.models ?? [],
tokenUsageEnabled: data?.token_usage.enabled ?? false,
isLoading,
error,
};
}

View File

@ -7,3 +7,12 @@ export interface Model {
supports_thinking?: boolean;
supports_reasoning_effort?: boolean;
}
export interface TokenUsageSettings {
enabled: boolean;
}
export interface ModelsResponse {
models: Model[];
token_usage: TokenUsageSettings;
}