LangChain.js中如何追踪链运行的Token消耗与成本？

LangChain JS目前没有Python版get_openai_callback()的直接对应实现，但可以通过自定义CallbackHandler监听LLM调用事件，结合OpenAI API返回的Token使用数据，实现每次链运行的成本统计。以下是针对ConversationalRetrievalQAChain（带BufferMemory和向量检索）的具体实现方案：

核心思路 #

1. 利用LangChain的回调系统，监听LLM调用的结束事件（handleLLMEnd），从OpenAI响应中提取官方返回的Token使用量（prompt_tokens、completion_tokens），数据准确性更高。
2. 预设OpenAI各模型的定价规则，实时计算单次链运行的总成本。
3. 针对ConversationalRetrievalQAChain的多LLM调用场景（如问题重写、最终回答生成），自动累计所有调用的Token消耗。

实现代码 #

1. 自定义Token成本追踪回调类 #

const { CallbackHandler } = require("langchain/callbacks");

class TokenCostTracker extends CallbackHandler {
  constructor() {
    super();
    // 初始化统计数据
    this.totalPromptTokens = 0;
    this.totalCompletionTokens = 0;
    this.totalCost = 0;
    // 按OpenAI官方定价配置（根据使用的模型调整）
    this.modelPricing = {
      "gpt-3.5-turbo": {
        promptPer1k: 0.0015, // 每1000个Prompt Token的美元成本
        completionPer1k: 0.002 // 每1000个Completion Token的美元成本
      },
      "gpt-3.5-turbo-16k": {
        promptPer1k: 0.003,
        completionPer1k: 0.004
      },
      "gpt-4": {
        promptPer1k: 0.03,
        completionPer1k: 0.06
      }
    };
  }

  // 监听LLM调用结束事件，提取Token使用数据并计算成本
  async handleLLMEnd(output) {
    const usage = output.generations[0][0].generationInfo?.usage;
    const model = output.llmOutput?.model;
    if (!usage || !model || !this.modelPricing[model]) return;

    // 累计Token数量
    this.totalPromptTokens += usage.prompt_tokens;
    this.totalCompletionTokens += usage.completion_tokens;

    // 计算本次调用成本并累计
    const promptCost = (usage.prompt_tokens / 1000) * this.modelPricing[model].promptPer1k;
    const completionCost = (usage.completion_tokens / 1000) * this.modelPricing[model].completionPer1k;
    this.totalCost += promptCost + completionCost;
  }

  // 重置统计数据，用于下一次链运行
  reset() {
    this.totalPromptTokens = 0;
    this.totalCompletionTokens = 0;
    this.totalCost = 0;
  }
}

2. 在ConversationalRetrievalQAChain中使用回调 #

const { ConversationalRetrievalQAChain, OpenAI } = require("langchain");
const { BufferMemory } = require("langchain/memory");

// 假设你已初始化好向量检索器
// const retriever = yourVectorStore.asRetriever();

async function runChainWithCostTracking(userQuery, chatMemory) {
  const llm = new OpenAI({ temperature: 0 });
  const costTracker = new TokenCostTracker();

  // 创建带回调的ConversationalRetrievalQAChain
  const qaChain = ConversationalRetrievalQAChain.fromLLM(
    llm,
    retriever,
    {
      memory: chatMemory,
      callbacks: [costTracker] // 注入自定义回调
    }
  );

  // 重置统计，确保每次运行独立统计
  costTracker.reset();
  const chainResult = await qaChain.call({ question: userQuery });

  // 输出本次运行的统计结果
  console.log("=== 本次链运行Token与成本统计 ===");
  console.log(`Prompt Token总量: ${costTracker.totalPromptTokens}`);
  console.log(`Completion Token总量: ${costTracker.totalCompletionTokens}`);
  console.log(`总成本(美元): ${costTracker.totalCost.toFixed(6)}`);

  return chainResult;
}

// 使用示例
const chatMemory = new BufferMemory({ memoryKey: "chat_history", returnMessages: true });
runChainWithCostTracking("请解释向量检索的工作原理", chatMemory);

补充说明 #

• Token数据准确性：优先使用OpenAI API返回的usage数据，比本地计算更精准，因为OpenAI的Token计算规则包含特殊字符、分词逻辑等细节。
• 模型定价更新：若OpenAI调整定价，需同步更新modelPricing中的配置。
• 本地Token计算 fallback：若API未返回usage（如使用自定义部署模型），可使用@dqbd/tiktoken库手动计算Token数，在handleLLMStart中计算Prompt Token，handleLLMEnd中计算Completion Token。

内容的提问来源于stack exchange，提问作者Daniel Limón