调用OpenAI Chat API前计算Tokens：对象数组处理与超限截断

用户的聊天机器人代码如下：

import { Configuration, OpenAIApi } from "openai"
import { readFile } from './readFile.js'

// Config OpenAI API
const configuration = new Configuration({
    organization: "xyx......",
    apiKey: "abc.......",
});

// OpenAI API instance
export const openai = new OpenAIApi(configuration);


const generateAnswer = async (conversation, userMessage) => {
    try {
        const dataset = await readFile();
        const dataFeed = { role: 'system', content: dataset };
        const prompt = conversation ? [...conversation?.messages, dataFeed, userMessage] : [dataFeed, userMessage];
        const completion = await openai.createChatCompletion({
            model: "gpt-3.5-turbo",
            messages: prompt
        })

        const aiMessage = completion.data.choices[0].message;
        console.log(completion.data.usage)
        return aiMessage
    } catch (e) {
        console.log(e)
    }
}
export { generateAnswer };

需求：调用OpenAI Chat API前，计算对话的Tokens总量，若超出模型限制则截断prompt。尝试过用gpt3 encoder，但因为prompt是对象数组而非字符串，不知道如何操作。

1. 使用官方推荐的tiktoken库计算Token #

OpenAI官方推荐用tiktoken计算Chat模型的Token数，比gpt3-encoder更贴合最新模型的Token规则，计算精度更高。

先安装依赖：

npm install tiktoken

2. 编写Token计算函数 #

Chat模型的消息数组（messages）Token计算需要覆盖每个消息的role、content以及格式开销，以下是适配gpt-3.5-turbo的计算逻辑：

import { encoding_for_model } from 'tiktoken';

// 获取对应模型的编码器
const encoder = encoding_for_model('gpt-3.5-turbo');

// 计算messages数组的总Token数
function calculateMessageTokens(messages) {
    let totalTokens = 0;
    for (const message of messages) {
        // 每个消息的基础格式开销：4个Token
        totalTokens += 4;
        // 计算role字段的Token数
        totalTokens += encoder.encode(message.role).length;
        // 计算content字段的Token数
        totalTokens += encoder.encode(message.content).length;
    }
    // 对话结尾的额外开销：2个Token
    totalTokens += 2;
    return totalTokens;
}

3. 实现Prompt截断逻辑 #

核心原则：保留业务信息dataFeed（system角色）和最新用户消息，只从最早的历史对话消息开始截断，直到总Token数符合模型限制。

修改原generateAnswer函数，加入Token校验和截断逻辑：

const generateAnswer = async (conversation, userMessage) => {
    try {
        const dataset = await readFile();
        const dataFeed = { role: 'system', content: dataset };
        // 构建初始消息数组
        let prompt = conversation ? [...conversation.messages, dataFeed, userMessage] : [dataFeed, userMessage];
        
        // 模型参数：gpt-3.5-turbo总Token上限4096，预留1000个Token给输出内容
        const MAX_INPUT_TOKENS = 4096 - 1000;
        let totalTokens = calculateMessageTokens(prompt);
        
        // 超出限制时循环截断历史消息
        while (totalTokens > MAX_INPUT_TOKENS) {
            // 找到第一个可删除的历史消息（排除system消息和最新用户消息）
            const historyIndex = prompt.findIndex(msg => msg.role !== 'system' && msg !== userMessage);
            if (historyIndex === -1) break; // 只剩必要消息，无法再截断
            
            // 删除最早的历史消息
            prompt.splice(historyIndex, 1);
            // 重新计算Token数
            totalTokens = calculateMessageTokens(prompt);
        }

        const completion = await openai.createChatCompletion({
            model: "gpt-3.5-turbo",
            messages: prompt
        })

        const aiMessage = completion.data.choices[0].message;
        console.log(completion.data.usage)
        return aiMessage
    } catch (e) {
        console.log(e)
    }
}

4. 注意事项 #

• 不同模型的总Token上限不同：比如gpt-3.5-turbo-16k支持16384个Token，需要根据使用的模型调整MAX_INPUT_TOKENS。
• 预留的输出Token数可按需调整：如果需要更长的AI回复，可以减少预留数，但不能超过模型总Token限制。
• 若dataFeed本身的Token数就超过输入上限，需要单独对业务数据集做截断处理。

内容的提问来源于stack exchange，提问作者Sorab