greatmengqi
61 lines
1.6 KiB
TypeScript
61 lines
1.6 KiB
TypeScript
import type { Message } from "@langchain/langgraph-sdk";
|
|
|
|
export interface TokenUsage {
|
|
inputTokens: number;
|
|
outputTokens: number;
|
|
totalTokens: number;
|
|
}
|
|
|
|
/**
|
|
* 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 {
|
|
if (message.type !== "ai") {
|
|
return null;
|
|
}
|
|
const usage = (message as Record<string, unknown>).usage_metadata as
|
|
| { input_tokens?: number; output_tokens?: number; total_tokens?: number }
|
|
| undefined;
|
|
if (!usage) {
|
|
return null;
|
|
}
|
|
return {
|
|
inputTokens: usage.input_tokens ?? 0,
|
|
outputTokens: usage.output_tokens ?? 0,
|
|
totalTokens: usage.total_tokens ?? 0,
|
|
};
|
|
}
|
|
|
|
/**
|
|
* Accumulate token usage across all AI messages in a thread.
|
|
*/
|
|
export function accumulateUsage(messages: Message[]): TokenUsage | null {
|
|
const cumulative: TokenUsage = {
|
|
inputTokens: 0,
|
|
outputTokens: 0,
|
|
totalTokens: 0,
|
|
};
|
|
let hasUsage = false;
|
|
for (const message of messages) {
|
|
const usage = getUsageMetadata(message);
|
|
if (usage) {
|
|
hasUsage = true;
|
|
cumulative.inputTokens += usage.inputTokens;
|
|
cumulative.outputTokens += usage.outputTokens;
|
|
cumulative.totalTokens += usage.totalTokens;
|
|
}
|
|
}
|
|
return hasUsage ? cumulative : null;
|
|
}
|
|
|
|
/**
|
|
* Format a token count for display: 1234 -> "1,234", 12345 -> "12.3K"
|
|
*/
|
|
export function formatTokenCount(count: number): string {
|
|
if (count < 10_000) {
|
|
return count.toLocaleString();
|
|
}
|
|
return `${(count / 1000).toFixed(1)}K`;
|
|
}
|