在LangChain.js中高效追踪OpenAI模型令牌消耗与成本

import { ChatOpenAI } from 'langchain/chat_models/openai';
import { HumanMessage } from 'langchain/schema';

// 定义全局变量来累积令牌使用量
let totalCompletionTokens = 0;
let totalPromptTokens = 0;
let totalExecutionTokens = 0;

// 实例化ChatOpenAI模型，并配置回调函数
const llm = new ChatOpenAI({
    // 配置回调数组
    callbacks: [
      {
        // handleLLMEnd在LLM调用结束后触发
        handleLLMEnd: (output, runId, parentRunId?, tags?) => {
          // 从output中安全地提取tokenUsage信息
          const { completionTokens, promptTokens, totalTokens } = output.llmOutput?.tokenUsage || {};

          // 累加令牌数量，确保处理undefined情况
          totalCompletionTokens += completionTokens ?? 0;
          totalPromptTokens += promptTokens ?? 0;
          totalExecutionTokens += totalTokens ?? 0;

          console.log(`--- LLM Run ID: ${runId} ---`);
          console.log(`Prompt Tokens: ${promptTokens ?? 0}`);
          console.log(`Completion Tokens: ${completionTokens ?? 0}`);
          console.log(`Total Tokens for this run: ${totalTokens ?? 0}`);
          console.log(`Current Accumulated Total Tokens: ${totalExecutionTokens}`);
        },
      },
    ],
    // 推荐使用特定的模型版本，例如 'gpt-3.5-turbo-0613'
    // 注意：某些模型或版本可能存在令牌统计不准确的问题
    modelName: 'gpt-3.5-turbo-0613', 
    // 其他OpenAI API参数，例如 temperature, openAIApiKey等
    temperature: 0.7,
    openAIApiKey: process.env.OPENAI_API_KEY, // 确保从环境变量获取API Key
});

// 示例：执行一次LLM调用
async function runExampleChat() {
    console.log("Starting LLM call...");
    const response = await llm.call([
        new HumanMessage("请用一句话介绍一下大型语言模型。"),
    ]);
    console.log("LLM Response:", response.content);
    console.log("\n--- Final Token Usage Summary ---");
    console.log(`Total Prompt Tokens Across All Runs: ${totalPromptTokens}`);
    console.log(`Total Completion Tokens Across All Runs: ${totalCompletionTokens}`);
    console.log(`Total Execution Tokens Across All Runs: ${totalExecutionTokens}`);

    // 重置计数器以便进行下一次独立测试（可选）
    totalCompletionTokens = 0;
    totalPromptTokens = 0;
    totalExecutionTokens = 0;
}

// 调用示例函数
runExampleChat().catch(console.error);

// 示例：执行另一次LLM调用
async function runAnotherExampleChat() {
    console.log("\nStarting another LLM call...");
    const response = await llm.call([
        new HumanMessage("解释一下云计算的优势。"),
    ]);
    console.log("LLM Response:", response.content);
    console.log("\n--- Final Token Usage Summary ---");
    console.log(`Total Prompt Tokens Across All Runs: ${totalPromptTokens}`);
    console.log(`Total Completion Tokens Across All Runs: ${totalCompletionTokens}`);
    console.log(`Total Execution Tokens Across All Runs: ${totalExecutionTokens}`);
}

// 可以在适当的时机调用 runAnotherExampleChat()
// setTimeout(() => runAnotherExampleChat().catch(console.error), 3000);