老胡茶室
老胡茶室
Beta

排错:Langchain getNumTokens 返回错误的 tokens 数

冯宇

症状

老胡茶室的智器间上线了一个新的小工具 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");

明白了这些,我们就可以分析 ChatGoogleGenerativeAIChatVertexAI 的实现了。

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 数量。