mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-05-24 08:55:59 +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
@@ -2,6 +2,7 @@ import type { Message } from "@langchain/langgraph-sdk";
|
||||
import {
|
||||
BookOpenTextIcon,
|
||||
ChevronUp,
|
||||
CoinsIcon,
|
||||
FolderOpenIcon,
|
||||
GlobeIcon,
|
||||
LightbulbIcon,
|
||||
@@ -24,6 +25,8 @@ import {
|
||||
import { CodeBlock } from "@/components/ai-elements/code-block";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { useI18n } from "@/core/i18n/hooks";
|
||||
import { formatTokenCount } from "@/core/messages/usage";
|
||||
import type { TokenDebugStep } from "@/core/messages/usage-model";
|
||||
import {
|
||||
extractReasoningContentFromMessage,
|
||||
findToolCallResult,
|
||||
@@ -43,10 +46,14 @@ export function MessageGroup({
|
||||
className,
|
||||
messages,
|
||||
isLoading = false,
|
||||
tokenDebugSteps = [],
|
||||
showTokenDebugSummaries = false,
|
||||
}: {
|
||||
className?: string;
|
||||
messages: Message[];
|
||||
isLoading?: boolean;
|
||||
tokenDebugSteps?: TokenDebugStep[];
|
||||
showTokenDebugSummaries?: boolean;
|
||||
}) {
|
||||
const { t } = useI18n();
|
||||
const [showAbove, setShowAbove] = useState(
|
||||
@@ -56,6 +63,28 @@ export function MessageGroup({
|
||||
env.NEXT_PUBLIC_STATIC_WEBSITE_ONLY === "true",
|
||||
);
|
||||
const steps = useMemo(() => convertToSteps(messages), [messages]);
|
||||
const debugStepByMessageId = useMemo(
|
||||
() =>
|
||||
new Map(
|
||||
tokenDebugSteps.map(
|
||||
(step) => [step.messageId || step.id, step] as const,
|
||||
),
|
||||
),
|
||||
[tokenDebugSteps],
|
||||
);
|
||||
const toolCallCountByMessageId = useMemo(() => {
|
||||
const counts = new Map<string, number>();
|
||||
|
||||
for (const step of steps) {
|
||||
if (step.type !== "toolCall" || !step.messageId) {
|
||||
continue;
|
||||
}
|
||||
|
||||
counts.set(step.messageId, (counts.get(step.messageId) ?? 0) + 1);
|
||||
}
|
||||
|
||||
return counts;
|
||||
}, [steps]);
|
||||
const lastToolCallStep = useMemo(() => {
|
||||
const filteredSteps = steps.filter((step) => step.type === "toolCall");
|
||||
return filteredSteps[filteredSteps.length - 1];
|
||||
@@ -77,6 +106,125 @@ export function MessageGroup({
|
||||
}
|
||||
}, [lastToolCallStep, steps]);
|
||||
const rehypePlugins = useRehypeSplitWordsIntoSpans(isLoading);
|
||||
const firstEligibleDebugSummaryStepIndexByMessageId = useMemo(() => {
|
||||
const firstIndices = new Map<string, number>();
|
||||
|
||||
if (!showTokenDebugSummaries) {
|
||||
return firstIndices;
|
||||
}
|
||||
|
||||
for (const [index, step] of steps.entries()) {
|
||||
const messageId = step.messageId;
|
||||
if (!messageId || firstIndices.has(messageId)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const debugStep = debugStepByMessageId.get(messageId);
|
||||
if (!debugStep) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const toolCallCount = toolCallCountByMessageId.get(messageId) ?? 0;
|
||||
if (!debugStep.sharedAttribution && toolCallCount > 0) {
|
||||
continue;
|
||||
}
|
||||
if (
|
||||
!debugStep.sharedAttribution &&
|
||||
toolCallCount === 0 &&
|
||||
debugStep.label === t.common.thinking &&
|
||||
debugStep.secondaryLabels.length === 0
|
||||
) {
|
||||
continue;
|
||||
}
|
||||
|
||||
firstIndices.set(messageId, index);
|
||||
}
|
||||
|
||||
return firstIndices;
|
||||
}, [
|
||||
debugStepByMessageId,
|
||||
showTokenDebugSummaries,
|
||||
steps,
|
||||
t.common.thinking,
|
||||
toolCallCountByMessageId,
|
||||
]);
|
||||
|
||||
const renderDebugSummary = (
|
||||
messageId: string | undefined,
|
||||
stepIndex: number,
|
||||
) => {
|
||||
if (!showTokenDebugSummaries || !messageId) {
|
||||
return null;
|
||||
}
|
||||
|
||||
const debugStep = debugStepByMessageId.get(messageId);
|
||||
if (!debugStep) {
|
||||
return null;
|
||||
}
|
||||
if (
|
||||
firstEligibleDebugSummaryStepIndexByMessageId.get(messageId) !== stepIndex
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return (
|
||||
<ChainOfThoughtStep
|
||||
key={`token-debug-${messageId}`}
|
||||
icon={CoinsIcon}
|
||||
label={
|
||||
<DebugStepLabel
|
||||
label={debugStep.label}
|
||||
token={formatDebugToken(debugStep, t)}
|
||||
/>
|
||||
}
|
||||
description={
|
||||
debugStep.sharedAttribution
|
||||
? t.tokenUsage.sharedAttribution
|
||||
: undefined
|
||||
}
|
||||
>
|
||||
{debugStep.secondaryLabels.length > 0 && (
|
||||
<ChainOfThoughtSearchResults>
|
||||
{debugStep.secondaryLabels.map((label, index) => (
|
||||
<ChainOfThoughtSearchResult
|
||||
key={`${debugStep.id}-${index}-${label}`}
|
||||
>
|
||||
{label}
|
||||
</ChainOfThoughtSearchResult>
|
||||
))}
|
||||
</ChainOfThoughtSearchResults>
|
||||
)}
|
||||
</ChainOfThoughtStep>
|
||||
);
|
||||
};
|
||||
|
||||
const renderToolCall = (
|
||||
step: CoTToolCallStep,
|
||||
options?: { isLast?: boolean },
|
||||
) => {
|
||||
const debugStep =
|
||||
showTokenDebugSummaries && step.messageId
|
||||
? debugStepByMessageId.get(step.messageId)
|
||||
: undefined;
|
||||
|
||||
return (
|
||||
<ToolCall
|
||||
key={step.id}
|
||||
{...step}
|
||||
isLast={options?.isLast}
|
||||
isLoading={isLoading}
|
||||
tokenDebugStep={
|
||||
debugStep && !debugStep.sharedAttribution ? debugStep : undefined
|
||||
}
|
||||
/>
|
||||
);
|
||||
};
|
||||
|
||||
const lastReasoningDebugStep =
|
||||
showTokenDebugSummaries && lastReasoningStep?.messageId
|
||||
? debugStepByMessageId.get(lastReasoningStep.messageId)
|
||||
: undefined;
|
||||
|
||||
return (
|
||||
<ChainOfThought
|
||||
className={cn("w-full gap-2 rounded-lg border p-0.5", className)}
|
||||
@@ -111,36 +259,46 @@ export function MessageGroup({
|
||||
{lastToolCallStep && (
|
||||
<ChainOfThoughtContent className="px-4 pb-2">
|
||||
{showAbove &&
|
||||
aboveLastToolCallSteps.map((step) =>
|
||||
step.type === "reasoning" ? (
|
||||
<ChainOfThoughtStep
|
||||
key={step.id}
|
||||
label={
|
||||
<MarkdownContent
|
||||
content={step.reasoning ?? ""}
|
||||
isLoading={isLoading}
|
||||
rehypePlugins={rehypePlugins}
|
||||
/>
|
||||
}
|
||||
></ChainOfThoughtStep>
|
||||
) : (
|
||||
<ToolCall key={step.id} {...step} isLoading={isLoading} />
|
||||
),
|
||||
)}
|
||||
aboveLastToolCallSteps.flatMap((step) => {
|
||||
const stepIndex = steps.indexOf(step);
|
||||
if (step.type === "reasoning") {
|
||||
return [
|
||||
renderDebugSummary(step.messageId, stepIndex),
|
||||
<ChainOfThoughtStep
|
||||
key={step.id}
|
||||
label={
|
||||
<MarkdownContent
|
||||
content={step.reasoning ?? ""}
|
||||
isLoading={isLoading}
|
||||
rehypePlugins={rehypePlugins}
|
||||
/>
|
||||
}
|
||||
></ChainOfThoughtStep>,
|
||||
];
|
||||
}
|
||||
|
||||
return [
|
||||
renderDebugSummary(step.messageId, stepIndex),
|
||||
renderToolCall(step),
|
||||
];
|
||||
})}
|
||||
{renderDebugSummary(
|
||||
lastToolCallStep.messageId,
|
||||
steps.indexOf(lastToolCallStep),
|
||||
)}
|
||||
{lastToolCallStep && (
|
||||
<FlipDisplay uniqueKey={lastToolCallStep.id ?? ""}>
|
||||
<ToolCall
|
||||
key={lastToolCallStep.id}
|
||||
{...lastToolCallStep}
|
||||
isLast={true}
|
||||
isLoading={isLoading}
|
||||
/>
|
||||
{renderToolCall(lastToolCallStep, { isLast: true })}
|
||||
</FlipDisplay>
|
||||
)}
|
||||
</ChainOfThoughtContent>
|
||||
)}
|
||||
{lastReasoningStep && (
|
||||
<>
|
||||
{renderDebugSummary(
|
||||
lastReasoningStep.messageId,
|
||||
steps.indexOf(lastReasoningStep),
|
||||
)}
|
||||
<Button
|
||||
key={lastReasoningStep.id}
|
||||
className="w-full items-start justify-start text-left"
|
||||
@@ -150,7 +308,22 @@ export function MessageGroup({
|
||||
<div className="flex w-full items-center justify-between">
|
||||
<ChainOfThoughtStep
|
||||
className="font-normal"
|
||||
label={t.common.thinking}
|
||||
label={
|
||||
<DebugStepLabel
|
||||
label={t.common.thinking}
|
||||
token={shouldInlineThinkingToken({
|
||||
debugStep: lastReasoningDebugStep,
|
||||
toolCallCount: lastReasoningStep.messageId
|
||||
? (toolCallCountByMessageId.get(
|
||||
lastReasoningStep.messageId,
|
||||
) ?? 0)
|
||||
: 0,
|
||||
enabled: showTokenDebugSummaries,
|
||||
thinkingLabel: t.common.thinking,
|
||||
t,
|
||||
})}
|
||||
/>
|
||||
}
|
||||
icon={LightbulbIcon}
|
||||
></ChainOfThoughtStep>
|
||||
<div>
|
||||
@@ -183,6 +356,60 @@ export function MessageGroup({
|
||||
);
|
||||
}
|
||||
|
||||
function formatDebugToken(
|
||||
debugStep: TokenDebugStep,
|
||||
t: ReturnType<typeof useI18n>["t"],
|
||||
) {
|
||||
return debugStep.usage
|
||||
? `${formatTokenCount(debugStep.usage.totalTokens)} ${t.tokenUsage.label}`
|
||||
: t.tokenUsage.unavailableShort;
|
||||
}
|
||||
|
||||
function shouldInlineThinkingToken({
|
||||
debugStep,
|
||||
toolCallCount,
|
||||
enabled,
|
||||
thinkingLabel,
|
||||
t,
|
||||
}: {
|
||||
debugStep?: TokenDebugStep;
|
||||
toolCallCount: number;
|
||||
enabled: boolean;
|
||||
thinkingLabel: string;
|
||||
t: ReturnType<typeof useI18n>["t"];
|
||||
}) {
|
||||
if (
|
||||
!enabled ||
|
||||
!debugStep ||
|
||||
debugStep.sharedAttribution ||
|
||||
toolCallCount > 0 ||
|
||||
debugStep.label !== thinkingLabel
|
||||
) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return formatDebugToken(debugStep, t);
|
||||
}
|
||||
|
||||
function DebugStepLabel({
|
||||
label,
|
||||
token,
|
||||
}: {
|
||||
label: React.ReactNode;
|
||||
token?: string | null;
|
||||
}) {
|
||||
return (
|
||||
<div className="flex items-center justify-between gap-3">
|
||||
<div className="min-w-0 flex-1">{label}</div>
|
||||
{token ? (
|
||||
<div className="text-muted-foreground shrink-0 font-mono text-[11px]">
|
||||
{token}
|
||||
</div>
|
||||
) : null}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
function ToolCall({
|
||||
id,
|
||||
messageId,
|
||||
@@ -191,6 +418,7 @@ function ToolCall({
|
||||
result,
|
||||
isLast = false,
|
||||
isLoading = false,
|
||||
tokenDebugStep,
|
||||
}: {
|
||||
id?: string;
|
||||
messageId?: string;
|
||||
@@ -199,10 +427,20 @@ function ToolCall({
|
||||
result?: string | Record<string, unknown>;
|
||||
isLast?: boolean;
|
||||
isLoading?: boolean;
|
||||
tokenDebugStep?: TokenDebugStep;
|
||||
}) {
|
||||
const { t } = useI18n();
|
||||
const { setOpen, autoOpen, autoSelect, selectedArtifact, select } =
|
||||
useArtifacts();
|
||||
const tokenLabel = tokenDebugStep
|
||||
? formatDebugToken(tokenDebugStep, t)
|
||||
: null;
|
||||
const resolveLabel = (fallback: React.ReactNode) =>
|
||||
tokenDebugStep ? (
|
||||
<DebugStepLabel label={tokenDebugStep.label} token={tokenLabel} />
|
||||
) : (
|
||||
fallback
|
||||
);
|
||||
|
||||
if (name === "web_search") {
|
||||
let label: React.ReactNode = t.toolCalls.searchForRelatedInfo;
|
||||
@@ -210,7 +448,11 @@ function ToolCall({
|
||||
label = t.toolCalls.searchOnWebFor(args.query);
|
||||
}
|
||||
return (
|
||||
<ChainOfThoughtStep key={id} label={label} icon={SearchIcon}>
|
||||
<ChainOfThoughtStep
|
||||
key={id}
|
||||
label={resolveLabel(label)}
|
||||
icon={SearchIcon}
|
||||
>
|
||||
{Array.isArray(result) && (
|
||||
<ChainOfThoughtSearchResults>
|
||||
{result.map((item) => (
|
||||
@@ -240,7 +482,11 @@ function ToolCall({
|
||||
}
|
||||
)?.results;
|
||||
return (
|
||||
<ChainOfThoughtStep key={id} label={label} icon={SearchIcon}>
|
||||
<ChainOfThoughtStep
|
||||
key={id}
|
||||
label={resolveLabel(label)}
|
||||
icon={SearchIcon}
|
||||
>
|
||||
{Array.isArray(results) && (
|
||||
<ChainOfThoughtSearchResults>
|
||||
{Array.isArray(results) &&
|
||||
@@ -280,7 +526,7 @@ function ToolCall({
|
||||
return (
|
||||
<ChainOfThoughtStep
|
||||
key={id}
|
||||
label={t.toolCalls.viewWebPage}
|
||||
label={resolveLabel(t.toolCalls.viewWebPage)}
|
||||
icon={GlobeIcon}
|
||||
>
|
||||
<ChainOfThoughtSearchResult>
|
||||
@@ -305,7 +551,11 @@ function ToolCall({
|
||||
}
|
||||
const path: string | undefined = (args as { path: string })?.path;
|
||||
return (
|
||||
<ChainOfThoughtStep key={id} label={description} icon={FolderOpenIcon}>
|
||||
<ChainOfThoughtStep
|
||||
key={id}
|
||||
label={resolveLabel(description)}
|
||||
icon={FolderOpenIcon}
|
||||
>
|
||||
{path && (
|
||||
<ChainOfThoughtSearchResult className="cursor-pointer">
|
||||
{path}
|
||||
@@ -321,7 +571,11 @@ function ToolCall({
|
||||
}
|
||||
const { path } = args as { path: string; content: string };
|
||||
return (
|
||||
<ChainOfThoughtStep key={id} label={description} icon={BookOpenTextIcon}>
|
||||
<ChainOfThoughtStep
|
||||
key={id}
|
||||
label={resolveLabel(description)}
|
||||
icon={BookOpenTextIcon}
|
||||
>
|
||||
{path && (
|
||||
<ChainOfThoughtSearchResult className="cursor-pointer">
|
||||
{path}
|
||||
@@ -353,7 +607,7 @@ function ToolCall({
|
||||
<ChainOfThoughtStep
|
||||
key={id}
|
||||
className="cursor-pointer"
|
||||
label={description}
|
||||
label={resolveLabel(description)}
|
||||
icon={NotebookPenIcon}
|
||||
onClick={() => {
|
||||
select(
|
||||
@@ -375,13 +629,19 @@ function ToolCall({
|
||||
const description: string | undefined = (args as { description: string })
|
||||
?.description;
|
||||
if (!description) {
|
||||
return t.toolCalls.executeCommand;
|
||||
return (
|
||||
<ChainOfThoughtStep
|
||||
key={id}
|
||||
label={resolveLabel(t.toolCalls.executeCommand)}
|
||||
icon={SquareTerminalIcon}
|
||||
/>
|
||||
);
|
||||
}
|
||||
const command: string | undefined = (args as { command: string })?.command;
|
||||
return (
|
||||
<ChainOfThoughtStep
|
||||
key={id}
|
||||
label={description}
|
||||
label={resolveLabel(description)}
|
||||
icon={SquareTerminalIcon}
|
||||
>
|
||||
{command && (
|
||||
@@ -398,7 +658,7 @@ function ToolCall({
|
||||
return (
|
||||
<ChainOfThoughtStep
|
||||
key={id}
|
||||
label={t.toolCalls.needYourHelp}
|
||||
label={resolveLabel(t.toolCalls.needYourHelp)}
|
||||
icon={MessageCircleQuestionMarkIcon}
|
||||
></ChainOfThoughtStep>
|
||||
);
|
||||
@@ -406,7 +666,7 @@ function ToolCall({
|
||||
return (
|
||||
<ChainOfThoughtStep
|
||||
key={id}
|
||||
label={t.toolCalls.writeTodos}
|
||||
label={resolveLabel(t.toolCalls.writeTodos)}
|
||||
icon={ListTodoIcon}
|
||||
></ChainOfThoughtStep>
|
||||
);
|
||||
@@ -416,7 +676,7 @@ function ToolCall({
|
||||
return (
|
||||
<ChainOfThoughtStep
|
||||
key={id}
|
||||
label={description ?? t.toolCalls.useTool(name)}
|
||||
label={resolveLabel(description ?? t.toolCalls.useTool(name))}
|
||||
icon={WrenchIcon}
|
||||
></ChainOfThoughtStep>
|
||||
);
|
||||
|
||||
@@ -50,7 +50,6 @@ import { cn } from "@/lib/utils";
|
||||
import { CopyButton } from "../copy-button";
|
||||
|
||||
import { MarkdownContent } from "./markdown-content";
|
||||
import { MessageTokenUsage } from "./message-token-usage";
|
||||
|
||||
function FeedbackButtons({
|
||||
threadId,
|
||||
@@ -121,20 +120,20 @@ function FeedbackButtons({
|
||||
|
||||
export function MessageListItem({
|
||||
className,
|
||||
threadId,
|
||||
message,
|
||||
isLoading,
|
||||
tokenUsageEnabled = false,
|
||||
feedback,
|
||||
runId,
|
||||
threadId,
|
||||
showCopyButton = true,
|
||||
}: {
|
||||
className?: string;
|
||||
message: Message;
|
||||
isLoading?: boolean;
|
||||
threadId: string;
|
||||
tokenUsageEnabled?: boolean;
|
||||
feedback?: FeedbackData | null;
|
||||
runId?: string;
|
||||
showCopyButton?: boolean;
|
||||
}) {
|
||||
const isHuman = message.type === "human";
|
||||
return (
|
||||
@@ -147,16 +146,17 @@ export function MessageListItem({
|
||||
message={message}
|
||||
isLoading={isLoading}
|
||||
threadId={threadId}
|
||||
tokenUsageEnabled={tokenUsageEnabled}
|
||||
/>
|
||||
{!isLoading && (
|
||||
{!isLoading && showCopyButton && (
|
||||
<MessageToolbar
|
||||
className={cn(
|
||||
isHuman ? "-bottom-9 justify-end" : "-bottom-8",
|
||||
"absolute right-0 left-0 z-20",
|
||||
isHuman
|
||||
? "absolute right-0 -bottom-9 left-0 justify-end"
|
||||
: "absolute right-0 bottom-0 left-0",
|
||||
"z-20 opacity-0 transition-opacity delay-200 duration-300 group-hover/conversation-message:opacity-100",
|
||||
)}
|
||||
>
|
||||
<div className="flex gap-1">
|
||||
<div className="pointer-events-auto flex gap-1">
|
||||
<CopyButton
|
||||
clipboardData={
|
||||
extractContentFromMessage(message) ??
|
||||
@@ -213,13 +213,11 @@ 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";
|
||||
@@ -297,11 +295,6 @@ function MessageContent_({
|
||||
<ReasoningTrigger />
|
||||
<ReasoningContent>{reasoningContent}</ReasoningContent>
|
||||
</Reasoning>
|
||||
<MessageTokenUsage
|
||||
enabled={tokenUsageEnabled}
|
||||
isLoading={isLoading}
|
||||
message={message}
|
||||
/>
|
||||
</AIElementMessageContent>
|
||||
);
|
||||
}
|
||||
@@ -339,11 +332,6 @@ function MessageContent_({
|
||||
className="my-3"
|
||||
components={components}
|
||||
/>
|
||||
<MessageTokenUsage
|
||||
enabled={tokenUsageEnabled}
|
||||
isLoading={isLoading}
|
||||
message={message}
|
||||
/>
|
||||
</AIElementMessageContent>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
import type { Message } from "@langchain/langgraph-sdk";
|
||||
import type { BaseStream } from "@langchain/langgraph-sdk/react";
|
||||
import { ChevronUpIcon, Loader2Icon } from "lucide-react";
|
||||
import { useCallback, useEffect, useRef } from "react";
|
||||
import { useCallback, useEffect, useMemo, useRef } from "react";
|
||||
|
||||
import {
|
||||
Conversation,
|
||||
@@ -8,15 +9,20 @@ import {
|
||||
} from "@/components/ai-elements/conversation";
|
||||
import { Button } from "@/components/ui/button";
|
||||
import { useI18n } from "@/core/i18n/hooks";
|
||||
import {
|
||||
buildTokenDebugSteps,
|
||||
type TokenUsageInlineMode,
|
||||
} from "@/core/messages/usage-model";
|
||||
import {
|
||||
extractContentFromMessage,
|
||||
extractPresentFilesFromMessage,
|
||||
extractReasoningContentFromMessage,
|
||||
extractTextFromMessage,
|
||||
groupMessages,
|
||||
getAssistantTurnUsageMessages,
|
||||
getMessageGroups,
|
||||
hasContent,
|
||||
hasPresentFiles,
|
||||
hasReasoning,
|
||||
hasToolCalls,
|
||||
} from "@/core/messages/utils";
|
||||
import { useRehypeSplitWordsIntoSpans } from "@/core/rehype";
|
||||
import type { Subtask } from "@/core/tasks";
|
||||
@@ -25,12 +31,16 @@ import type { AgentThreadState } from "@/core/threads";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
import { ArtifactFileList } from "../artifacts/artifact-file-list";
|
||||
import { CopyButton } from "../copy-button";
|
||||
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 {
|
||||
MessageTokenUsageDebugList,
|
||||
MessageTokenUsageList,
|
||||
} from "./message-token-usage";
|
||||
import { MessageListSkeleton } from "./skeleton";
|
||||
import { SubtaskCard } from "./subtask-card";
|
||||
|
||||
@@ -149,7 +159,7 @@ export function MessageList({
|
||||
threadId,
|
||||
thread,
|
||||
paddingBottom = MESSAGE_LIST_DEFAULT_PADDING_BOTTOM,
|
||||
tokenUsageEnabled = false,
|
||||
tokenUsageInlineMode = "off",
|
||||
hasMoreHistory,
|
||||
loadMoreHistory,
|
||||
isHistoryLoading,
|
||||
@@ -158,7 +168,7 @@ export function MessageList({
|
||||
threadId: string;
|
||||
thread: BaseStream<AgentThreadState>;
|
||||
paddingBottom?: number;
|
||||
tokenUsageEnabled?: boolean;
|
||||
tokenUsageInlineMode?: TokenUsageInlineMode;
|
||||
hasMoreHistory?: boolean;
|
||||
loadMoreHistory?: () => void;
|
||||
isHistoryLoading?: boolean;
|
||||
@@ -167,10 +177,85 @@ export function MessageList({
|
||||
const rehypePlugins = useRehypeSplitWordsIntoSpans(thread.isLoading);
|
||||
const updateSubtask = useUpdateSubtask();
|
||||
const messages = thread.messages;
|
||||
const groupedMessages = getMessageGroups(messages);
|
||||
const turnUsageMessagesByGroupIndex =
|
||||
getAssistantTurnUsageMessages(groupedMessages);
|
||||
const tokenDebugSteps = useMemo(
|
||||
() => buildTokenDebugSteps(messages, t),
|
||||
[messages, t],
|
||||
);
|
||||
|
||||
const renderAssistantCopyButton = useCallback((messages: Message[]) => {
|
||||
const clipboardData = [...messages]
|
||||
.reverse()
|
||||
.filter((message) => message.type === "ai")
|
||||
.map((message) => {
|
||||
const content = extractContentFromMessage(message);
|
||||
return content ?? extractReasoningContentFromMessage(message) ?? "";
|
||||
})
|
||||
.find((content) => content.length > 0);
|
||||
|
||||
if (!clipboardData) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return (
|
||||
<div className="mt-2 flex justify-start opacity-0 transition-opacity delay-200 duration-300 group-hover/assistant-turn:opacity-100">
|
||||
<CopyButton clipboardData={clipboardData} />
|
||||
</div>
|
||||
);
|
||||
}, []);
|
||||
|
||||
const renderTokenUsage = useCallback(
|
||||
({
|
||||
messages,
|
||||
turnUsageMessages,
|
||||
inlineDebug = true,
|
||||
debugMessageIds,
|
||||
}: {
|
||||
messages: Message[];
|
||||
turnUsageMessages?: Message[] | null;
|
||||
inlineDebug?: boolean;
|
||||
debugMessageIds?: string[];
|
||||
}) => {
|
||||
if (tokenUsageInlineMode === "per_turn") {
|
||||
return (
|
||||
<MessageTokenUsageList
|
||||
enabled={true}
|
||||
isLoading={thread.isLoading}
|
||||
messages={turnUsageMessages ?? []}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
if (tokenUsageInlineMode === "step_debug" && inlineDebug) {
|
||||
const messageIds = new Set(
|
||||
debugMessageIds ??
|
||||
messages
|
||||
.filter((message) => message.type === "ai")
|
||||
.map((message) => message.id)
|
||||
.filter((id): id is string => typeof id === "string"),
|
||||
);
|
||||
return (
|
||||
<MessageTokenUsageDebugList
|
||||
enabled={true}
|
||||
isLoading={thread.isLoading}
|
||||
steps={tokenDebugSteps.filter((step) =>
|
||||
messageIds.has(step.messageId),
|
||||
)}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
return null;
|
||||
},
|
||||
[thread.isLoading, tokenDebugSteps, tokenUsageInlineMode],
|
||||
);
|
||||
|
||||
if (thread.isThreadLoading && messages.length === 0) {
|
||||
return <MessageListSkeleton />;
|
||||
}
|
||||
|
||||
return (
|
||||
<Conversation
|
||||
className={cn("flex size-full flex-col justify-center", className)}
|
||||
@@ -181,19 +266,37 @@ export function MessageList({
|
||||
hasMore={hasMoreHistory}
|
||||
loadMore={loadMoreHistory}
|
||||
/>
|
||||
{groupMessages(messages, (group) => {
|
||||
{groupedMessages.map((group, groupIndex) => {
|
||||
const turnUsageMessages = turnUsageMessagesByGroupIndex[groupIndex];
|
||||
|
||||
if (group.type === "human" || group.type === "assistant") {
|
||||
return group.messages.map((msg) => {
|
||||
return (
|
||||
<MessageListItem
|
||||
key={`${group.id}/${msg.id}`}
|
||||
threadId={threadId}
|
||||
message={msg}
|
||||
isLoading={thread.isLoading}
|
||||
tokenUsageEnabled={tokenUsageEnabled}
|
||||
/>
|
||||
);
|
||||
});
|
||||
return (
|
||||
<div
|
||||
key={group.id}
|
||||
className={cn(
|
||||
"w-full",
|
||||
group.type === "assistant" && "group/assistant-turn",
|
||||
)}
|
||||
>
|
||||
{group.messages.map((msg) => {
|
||||
return (
|
||||
<MessageListItem
|
||||
key={`${group.id}/${msg.id}`}
|
||||
message={msg}
|
||||
isLoading={thread.isLoading}
|
||||
threadId={threadId}
|
||||
showCopyButton={group.type !== "assistant"}
|
||||
/>
|
||||
);
|
||||
})}
|
||||
{renderTokenUsage({
|
||||
messages: group.messages,
|
||||
turnUsageMessages,
|
||||
})}
|
||||
{group.type === "assistant" &&
|
||||
renderAssistantCopyButton(group.messages)}
|
||||
</div>
|
||||
);
|
||||
} else if (group.type === "assistant:clarification") {
|
||||
const message = group.messages[0];
|
||||
if (message && hasContent(message)) {
|
||||
@@ -204,11 +307,10 @@ export function MessageList({
|
||||
isLoading={thread.isLoading}
|
||||
rehypePlugins={rehypePlugins}
|
||||
/>
|
||||
<MessageTokenUsageList
|
||||
enabled={tokenUsageEnabled}
|
||||
isLoading={thread.isLoading}
|
||||
messages={group.messages}
|
||||
/>
|
||||
{renderTokenUsage({
|
||||
messages: group.messages,
|
||||
turnUsageMessages,
|
||||
})}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
@@ -232,11 +334,10 @@ export function MessageList({
|
||||
/>
|
||||
)}
|
||||
<ArtifactFileList files={files} threadId={threadId} />
|
||||
<MessageTokenUsageList
|
||||
enabled={tokenUsageEnabled}
|
||||
isLoading={thread.isLoading}
|
||||
messages={group.messages}
|
||||
/>
|
||||
{renderTokenUsage({
|
||||
messages: group.messages,
|
||||
turnUsageMessages,
|
||||
})}
|
||||
</div>
|
||||
);
|
||||
} else if (group.type === "assistant:subagent") {
|
||||
@@ -289,7 +390,19 @@ export function MessageList({
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const results: React.ReactNode[] = [];
|
||||
const subagentDebugMessageIds: string[] = [];
|
||||
if (tasks.size > 0) {
|
||||
results.push(
|
||||
<div
|
||||
key="subtask-count"
|
||||
className="text-muted-foreground pt-2 text-sm font-normal"
|
||||
>
|
||||
{t.subtasks.executing(tasks.size)}
|
||||
</div>,
|
||||
);
|
||||
}
|
||||
for (const message of group.messages.filter(
|
||||
(message) => message.type === "ai",
|
||||
)) {
|
||||
@@ -299,17 +412,17 @@ export function MessageList({
|
||||
key={"thinking-group-" + message.id}
|
||||
messages={[message]}
|
||||
isLoading={thread.isLoading}
|
||||
tokenDebugSteps={tokenDebugSteps.filter(
|
||||
(step) => step.messageId === message.id,
|
||||
)}
|
||||
showTokenDebugSummaries={
|
||||
tokenUsageInlineMode === "step_debug"
|
||||
}
|
||||
/>,
|
||||
);
|
||||
} else if (message.id) {
|
||||
subagentDebugMessageIds.push(message.id);
|
||||
}
|
||||
results.push(
|
||||
<div
|
||||
key="subtask-count"
|
||||
className="text-muted-foreground font-norma pt-2 text-sm"
|
||||
>
|
||||
{t.subtasks.executing(tasks.size)}
|
||||
</div>,
|
||||
);
|
||||
const taskIds = message.tool_calls
|
||||
?.filter((toolCall) => toolCall.name === "task")
|
||||
.map((toolCall) => toolCall.id);
|
||||
@@ -329,30 +442,31 @@ export function MessageList({
|
||||
className="relative z-1 flex flex-col gap-2"
|
||||
>
|
||||
{results}
|
||||
<MessageTokenUsageList
|
||||
enabled={tokenUsageEnabled}
|
||||
isLoading={thread.isLoading}
|
||||
messages={group.messages}
|
||||
/>
|
||||
{renderTokenUsage({
|
||||
messages: group.messages,
|
||||
turnUsageMessages,
|
||||
debugMessageIds: subagentDebugMessageIds,
|
||||
})}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
const tokenUsageMessages = group.messages.filter(
|
||||
(message) =>
|
||||
message.type === "ai" &&
|
||||
(hasToolCalls(message) ? true : !hasContent(message)),
|
||||
);
|
||||
return (
|
||||
<div key={"group-" + group.id} className="w-full">
|
||||
<MessageGroup
|
||||
messages={group.messages}
|
||||
isLoading={thread.isLoading}
|
||||
tokenDebugSteps={tokenDebugSteps.filter((step) =>
|
||||
group.messages.some(
|
||||
(message) => message.id === step.messageId,
|
||||
),
|
||||
)}
|
||||
showTokenDebugSummaries={tokenUsageInlineMode === "step_debug"}
|
||||
/>
|
||||
<MessageTokenUsageList
|
||||
enabled={tokenUsageEnabled}
|
||||
isLoading={thread.isLoading}
|
||||
messages={tokenUsageMessages}
|
||||
/>
|
||||
{renderTokenUsage({
|
||||
messages: group.messages,
|
||||
turnUsageMessages,
|
||||
inlineDebug: false,
|
||||
})}
|
||||
</div>
|
||||
);
|
||||
})}
|
||||
|
||||
@@ -1,29 +1,27 @@
|
||||
import type { Message } from "@langchain/langgraph-sdk";
|
||||
import { CoinsIcon } from "lucide-react";
|
||||
|
||||
import { Badge } from "@/components/ui/badge";
|
||||
import { useI18n } from "@/core/i18n/hooks";
|
||||
import { formatTokenCount, getUsageMetadata } from "@/core/messages/usage";
|
||||
import { accumulateUsage, formatTokenCount } from "@/core/messages/usage";
|
||||
import type { TokenDebugStep } from "@/core/messages/usage-model";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
export function MessageTokenUsage({
|
||||
function TokenUsageSummary({
|
||||
className,
|
||||
enabled = false,
|
||||
isLoading = false,
|
||||
message,
|
||||
inputTokens,
|
||||
outputTokens,
|
||||
totalTokens,
|
||||
unavailable = false,
|
||||
}: {
|
||||
className?: string;
|
||||
enabled?: boolean;
|
||||
isLoading?: boolean;
|
||||
message: Message;
|
||||
inputTokens?: number;
|
||||
outputTokens?: number;
|
||||
totalTokens?: number;
|
||||
unavailable?: boolean;
|
||||
}) {
|
||||
const { t } = useI18n();
|
||||
|
||||
if (!enabled || isLoading || message.type !== "ai") {
|
||||
return null;
|
||||
}
|
||||
|
||||
const usage = getUsageMetadata(message);
|
||||
|
||||
return (
|
||||
<div
|
||||
className={cn(
|
||||
@@ -35,16 +33,16 @@ export function MessageTokenUsage({
|
||||
<CoinsIcon className="size-3" />
|
||||
{t.tokenUsage.label}
|
||||
</span>
|
||||
{usage ? (
|
||||
{!unavailable ? (
|
||||
<>
|
||||
<span>
|
||||
{t.tokenUsage.input}: {formatTokenCount(usage.inputTokens)}
|
||||
{t.tokenUsage.input}: {formatTokenCount(inputTokens ?? 0)}
|
||||
</span>
|
||||
<span>
|
||||
{t.tokenUsage.output}: {formatTokenCount(usage.outputTokens)}
|
||||
{t.tokenUsage.output}: {formatTokenCount(outputTokens ?? 0)}
|
||||
</span>
|
||||
<span className="font-medium">
|
||||
{t.tokenUsage.total}: {formatTokenCount(usage.totalTokens)}
|
||||
{t.tokenUsage.total}: {formatTokenCount(totalTokens ?? 0)}
|
||||
</span>
|
||||
</>
|
||||
) : (
|
||||
@@ -75,17 +73,93 @@ export function MessageTokenUsageList({
|
||||
return null;
|
||||
}
|
||||
|
||||
const usage = accumulateUsage(aiMessages);
|
||||
|
||||
return (
|
||||
<>
|
||||
{aiMessages.map((message, index) => (
|
||||
<MessageTokenUsage
|
||||
className={className}
|
||||
enabled={enabled}
|
||||
isLoading={isLoading}
|
||||
key={message.id ?? index}
|
||||
message={message}
|
||||
/>
|
||||
))}
|
||||
</>
|
||||
<TokenUsageSummary
|
||||
className={className}
|
||||
inputTokens={usage?.inputTokens}
|
||||
outputTokens={usage?.outputTokens}
|
||||
totalTokens={usage?.totalTokens}
|
||||
unavailable={!usage}
|
||||
/>
|
||||
);
|
||||
}
|
||||
|
||||
export function MessageTokenUsageDebugList({
|
||||
className,
|
||||
enabled = false,
|
||||
isLoading = false,
|
||||
steps,
|
||||
}: {
|
||||
className?: string;
|
||||
enabled?: boolean;
|
||||
isLoading?: boolean;
|
||||
steps: TokenDebugStep[];
|
||||
}) {
|
||||
const { t } = useI18n();
|
||||
|
||||
if (!enabled || isLoading) {
|
||||
return null;
|
||||
}
|
||||
|
||||
if (steps.length === 0) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return (
|
||||
<div className={cn("border-border/60 mt-1 border-t pt-2", className)}>
|
||||
<div className="space-y-2">
|
||||
{steps.map((step) => (
|
||||
<div
|
||||
key={step.id}
|
||||
className="bg-muted/30 border-border/50 flex items-start justify-between gap-3 rounded-md border px-3 py-2"
|
||||
>
|
||||
<div className="min-w-0 flex-1 space-y-1">
|
||||
<div className="text-foreground flex items-center gap-2 text-xs font-medium">
|
||||
<CoinsIcon className="text-muted-foreground size-3" />
|
||||
<span className="truncate">{step.label}</span>
|
||||
</div>
|
||||
{step.secondaryLabels.length > 0 && (
|
||||
<div className="flex flex-wrap gap-1.5">
|
||||
{step.secondaryLabels.map((label, index) => (
|
||||
<Badge
|
||||
key={`${step.id}-${index}-${label}`}
|
||||
className="px-1.5 py-0 text-[10px] font-normal"
|
||||
variant="secondary"
|
||||
>
|
||||
{label}
|
||||
</Badge>
|
||||
))}
|
||||
</div>
|
||||
)}
|
||||
{step.sharedAttribution && (
|
||||
<div className="text-muted-foreground text-[11px]">
|
||||
{t.tokenUsage.sharedAttribution}
|
||||
</div>
|
||||
)}
|
||||
<div className="text-muted-foreground text-[11px]">
|
||||
{step.usage ? (
|
||||
<>
|
||||
{t.tokenUsage.input}:{" "}
|
||||
{formatTokenCount(step.usage.inputTokens)}
|
||||
{" · "}
|
||||
{t.tokenUsage.output}:{" "}
|
||||
{formatTokenCount(step.usage.outputTokens)}
|
||||
</>
|
||||
) : (
|
||||
t.tokenUsage.unavailableShort
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
<Badge className="shrink-0 font-mono" variant="outline">
|
||||
{step.usage
|
||||
? `${formatTokenCount(step.usage.totalTokens)} ${t.tokenUsage.label}`
|
||||
: t.tokenUsage.unavailableShort}
|
||||
</Badge>
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
@@ -1,60 +1,81 @@
|
||||
"use client";
|
||||
|
||||
import type { Message } from "@langchain/langgraph-sdk";
|
||||
import { CoinsIcon } from "lucide-react";
|
||||
import { ChevronDownIcon, CoinsIcon } from "lucide-react";
|
||||
import { useMemo } from "react";
|
||||
|
||||
import { Button } from "@/components/ui/button";
|
||||
import {
|
||||
Tooltip,
|
||||
TooltipContent,
|
||||
TooltipTrigger,
|
||||
} from "@/components/ui/tooltip";
|
||||
DropdownMenu,
|
||||
DropdownMenuContent,
|
||||
DropdownMenuLabel,
|
||||
DropdownMenuRadioGroup,
|
||||
DropdownMenuRadioItem,
|
||||
DropdownMenuSeparator,
|
||||
DropdownMenuTrigger,
|
||||
} from "@/components/ui/dropdown-menu";
|
||||
import { useI18n } from "@/core/i18n/hooks";
|
||||
import { accumulateUsage, formatTokenCount } from "@/core/messages/usage";
|
||||
import {
|
||||
getTokenUsageViewPreset,
|
||||
tokenUsagePreferencesFromPreset,
|
||||
type TokenUsagePreferences,
|
||||
type TokenUsageViewPreset,
|
||||
} from "@/core/messages/usage-model";
|
||||
import { cn } from "@/lib/utils";
|
||||
|
||||
interface TokenUsageIndicatorProps {
|
||||
messages: Message[];
|
||||
enabled?: boolean;
|
||||
preferences: TokenUsagePreferences;
|
||||
onPreferencesChange: (preferences: TokenUsagePreferences) => void;
|
||||
className?: string;
|
||||
}
|
||||
|
||||
export function TokenUsageIndicator({
|
||||
messages,
|
||||
enabled = false,
|
||||
preferences,
|
||||
onPreferencesChange,
|
||||
className,
|
||||
}: TokenUsageIndicatorProps) {
|
||||
const { t } = useI18n();
|
||||
|
||||
const usage = useMemo(() => accumulateUsage(messages), [messages]);
|
||||
const preset = getTokenUsageViewPreset(preferences);
|
||||
|
||||
if (!enabled) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return (
|
||||
<Tooltip delayDuration={200}>
|
||||
<TooltipTrigger asChild>
|
||||
<button
|
||||
<DropdownMenu>
|
||||
<DropdownMenuTrigger asChild>
|
||||
<Button
|
||||
type="button"
|
||||
variant="ghost"
|
||||
className={cn(
|
||||
"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",
|
||||
"text-muted-foreground bg-background/70 hover:bg-background/90 flex h-auto items-center gap-1.5 rounded-full border px-2 py-1 text-xs font-normal",
|
||||
className,
|
||||
)}
|
||||
>
|
||||
<CoinsIcon size={14} />
|
||||
<span>{t.tokenUsage.label}</span>
|
||||
<span className="font-mono">
|
||||
{usage ? formatTokenCount(usage.totalTokens) : "-"}
|
||||
{preferences.headerTotal
|
||||
? usage
|
||||
? formatTokenCount(usage.totalTokens)
|
||||
: "-"
|
||||
: t.tokenUsage.presets[presetKeyToTranslationKey(preset)]}
|
||||
</span>
|
||||
</button>
|
||||
</TooltipTrigger>
|
||||
<TooltipContent side="bottom" align="end">
|
||||
<div className="space-y-1 text-xs">
|
||||
<div className="font-medium">{t.tokenUsage.title}</div>
|
||||
<ChevronDownIcon className="size-3" />
|
||||
</Button>
|
||||
</DropdownMenuTrigger>
|
||||
<DropdownMenuContent side="bottom" align="end" className="w-80">
|
||||
<DropdownMenuLabel>{t.tokenUsage.title}</DropdownMenuLabel>
|
||||
<div className="px-2 py-1 text-xs">
|
||||
{usage ? (
|
||||
<>
|
||||
<div className="space-y-1">
|
||||
<div className="flex justify-between gap-4">
|
||||
<span>{t.tokenUsage.input}</span>
|
||||
<span className="font-mono">
|
||||
@@ -75,14 +96,53 @@ export function TokenUsageIndicator({
|
||||
</span>
|
||||
</div>
|
||||
</div>
|
||||
</>
|
||||
</div>
|
||||
) : (
|
||||
<div className="text-muted-foreground max-w-56">
|
||||
<div className="text-muted-foreground">
|
||||
{t.tokenUsage.unavailable}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</TooltipContent>
|
||||
</Tooltip>
|
||||
<DropdownMenuSeparator />
|
||||
<DropdownMenuLabel>{t.tokenUsage.view}</DropdownMenuLabel>
|
||||
<DropdownMenuRadioGroup
|
||||
value={preset}
|
||||
onValueChange={(value) =>
|
||||
onPreferencesChange(
|
||||
tokenUsagePreferencesFromPreset(value as TokenUsageViewPreset),
|
||||
)
|
||||
}
|
||||
>
|
||||
{(
|
||||
["off", "summary", "per_turn", "debug"] as TokenUsageViewPreset[]
|
||||
).map((value) => {
|
||||
const translationKey = presetKeyToTranslationKey(value);
|
||||
return (
|
||||
<DropdownMenuRadioItem key={value} value={value}>
|
||||
<div className="grid gap-0.5">
|
||||
<span>{t.tokenUsage.presets[translationKey]}</span>
|
||||
<span className="text-muted-foreground text-xs">
|
||||
{t.tokenUsage.presetDescriptions[translationKey]}
|
||||
</span>
|
||||
</div>
|
||||
</DropdownMenuRadioItem>
|
||||
);
|
||||
})}
|
||||
</DropdownMenuRadioGroup>
|
||||
<DropdownMenuSeparator />
|
||||
<div className="text-muted-foreground px-2 py-2 text-xs leading-relaxed">
|
||||
{t.tokenUsage.note}
|
||||
</div>
|
||||
</DropdownMenuContent>
|
||||
</DropdownMenu>
|
||||
);
|
||||
}
|
||||
|
||||
function presetKeyToTranslationKey(preset: TokenUsageViewPreset) {
|
||||
switch (preset) {
|
||||
case "per_turn":
|
||||
return "perTurn" as const;
|
||||
default:
|
||||
return preset;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user