症状
老胡茶室的智器间上线了一个新的小工具 Token Counter,不过我们在测试的时候发现它的 tokens 数量计算似乎并不准确。
我们使用了 Langchain.js 的生态,这部分具体的代码示例如下:
import {
ChatGoogleGenerativeAI,
type GoogleGenerativeAIChatInput,
} from "@langchain/google-genai";
import { ChatVertexAI, type VertexAIInput } from "@langchain/google-vertexai";
export function newChatModel(params: unknown): BaseChatModel {
switch (LLM_PROVIDER) {
case "google":
return new ChatGoogleGenerativeAI({
...(params as GoogleGenerativeAIChatInput),
apiKey: GOOGLE_API_KEY,
model: LLM_NAME,
});
case "vertexai": {
return new ChatVertexAI({
...(params as VertexAIInput),
model: LLM_NAME,
authOptions: {
projectId: projectId,
credentials: credentials,
},
});
}
default:
throw new Error(`Unsupported LLM provider: ${LLM_PROVIDER}`);
}
}
const model = newChatModel({});
await model.getNumTokens(
"品茗之余,论道技术与人生;涵盖 AI、区块链、软件研发与创业思辨之道。"
);
我们发现即使使用同一个 model gemini-2.0-flash
,使用 Gemini API 接口 (即 ChatGoogleGenerativeAI
) 和 Vertex AI 接口 (即 ChatVertexAI
)时,发现它们的 tokens 数量计算结果差异很大。我们使用了 getNumTokens
方法来计算 tokens 数量。
当使用 ChatGoogleGenerativeAI
时,返回的 tokens 数量是 70。
而使用 ChatVertexAI
时,返回的 tokens 数量是 9,并且可以在终端看到有这样的错误堆栈:
Failed to calculate number of tokens, falling back to approximate count Error: Unknown model
at getEncodingNameForModel (file:///.../.pnpm/js-tiktoken@1.0.20/node_modules/js-tiktoken/dist/chunk-ZDNLBERF.js:273:13)
at encodingForModel (file:///.../.pnpm/@langchain+core@0.3.55_openai@4.98.0_ws@8.18.2_zod@3.24.4_/node_modules/@langchain/core/dist/utils/tiktoken.js:19:24)
at ChatVertexAI.getNumTokens (file:///.../.pnpm/@langchain+core@0.3.55_openai@4.98.0_ws@8.18.2_zod@3.24.4_/node_modules/@langchain/core/dist/language_models/base.js:197:40)
于是怀着好奇的心态,我们开始了排错之旅。
原因
我们首先查看了 getNumTokens
方法的实现,发现它的实现代码部分如下:
export abstract class BaseLanguageModel<...> {
// ... 省略其他代码
async getNumTokens(content: MessageContent) {
// TODO: Figure out correct value.
if (typeof content !== "string") {
return 0;
}
// fallback to approximate calculation if tiktoken is not available
let numTokens = Math.ceil(content.length / 4);
if (!this._encoding) {
try {
this._encoding = await encodingForModel(
"modelName" in this
? getModelNameForTiktoken(this.modelName as string)
: "gpt2"
);
} catch (error) {
console.warn(
"Failed to calculate number of tokens, falling back to approximate count",
error
);
}
}
if (this._encoding) {
try {
numTokens = this._encoding.encode(content).length;
} catch (error) {
console.warn(
"Failed to calculate number of tokens, falling back to approximate count",
error
);
}
}
return numTokens;
}
// ... 省略其他代码
}
所有的 Chat Model 都继承自 BaseLanguageModel
,根据代码分析,我们可以知道 getNumTokens
方法的实现逻辑大致:
- fallback 一个近似值,计算方法为字符串长度除以 4 向上取整。
- 如果当前子类有
modelName
属性,则使用getModelNameForTiktoken
方法来获取模型名称。 - 如果当前子类没有
modelName
属性,则使用gpt2
作为模型名称。 - 调用
tiktoken
的方法来计算 tokens 数量。 - 如果出错,则返回 fallback 值。
我们发现 langchain.js 在计算 tokens 数量时,使用了 js-tiktoken
这个库来计算 tokens 数量,而这个库只支持 OpenAI 的模型名称,具体支持列表参见 这里,对于不支持的模型,则直接抛出异常:
throw new Error("Unknown model");
明白了这些,我们就可以分析 ChatGoogleGenerativeAI
和 ChatVertexAI
的实现了。
ChatGoogleGenerativeAI
直接继承自 BaseLanguageModel
,并没有实现 getNumTokens
方法,因此会直接调用父类的方法:
export class ChatGoogleGenerativeAI
extends BaseChatModel<GoogleGenerativeAIChatCallOptions, AIMessageChunk>
implements GoogleGenerativeAIChatInput
// ...
它也没有定义 modelName
属性,因此在调用 getNumTokens
方法时,会使用 gpt2
作为模型名称来计算 tokens 数量,得到 70
这个结果。
而 ChatVertexAI
并不直接继承自 BaseLanguageModel
,而是经过了一系列继承链之后才间接继承了这个类:
export class ChatVertexAI extends ChatGoogle {
//...
}
// ...
export declare class ChatGoogle
extends ChatGoogleBase<GoogleAuthOptions>
implements ChatGoogleInput {
// ...
}
// ...
export abstract class ChatGoogleBase<AuthOptions>
extends BaseChatModel<GoogleAIBaseLanguageModelCallOptions, AIMessageChunk>
implements ChatGoogleBaseInput<AuthOptions>
{
// ...
modelName = "gemini-pro";
// ...
}
它定义了 modelName
属性,并且在 getNumTokens
方法中使用了 getModelNameForTiktoken
方法来获取模型名称,由于不支持 gemini 模型,因此会抛出异常,导致 getNumTokens
方法会使用 fallback 结果,最终返回了 9
。
总结
经过排查,我们发现 getNumTokens
方法不能正确计算 tokens 数量,甚至对于非 OpenAI 的模型计算结果也不准确 (实际在 Vertex AI 官方的计算器中,我们可以算出真实的 tokens 数量是 26)。
Langchain.js 在抽象基类中定义了 getNumTokens
方法,并且使用了 tiktoken
来计算 tokens 数量,可能是期望后续继承的子类可以自己实现这个方法来计算 tokens 数量,但是很明显,后续的开发者偷懒了,并没有实现这个方法,而是直接使用了父类的方法,导致了这个问题。
如果想要真正获取准确的 tokens 数量,我们应该调用模型提供的 count tokens 的接口,而不是使用 getNumTokens
方法来计算 tokens 数量。例如,Gemini 就提供了一个 Count Tokens API 来计算 tokens 数量。
另外就是,按照 Langchain.js 的文档,在 response 的元数据中,可能也会返回 tokens 数量,如果想要实现类似于按照 tokens 计费的需求,可能也可以使用这个元数据来计算 tokens 数量。