ai-utils.js
Version:
Build AI applications, chatbots, and agents with JavaScript and TypeScript.
51 lines (50 loc) • 2.22 kB
JavaScript
;
Object.defineProperty(exports, "__esModule", { value: true });
exports.trimChatPrompt = void 0;
const validateChatPrompt_js_1 = require("./validateChatPrompt.cjs");
/**
* Keeps only the most recent messages in the prompt, while leaving enough space for the completion.
*
* It will remove user-ai message pairs that don't fit. The result is always a valid chat prompt.
*
* When the minimal chat prompt (system message + last user message) is already too long, it will only
* return this minimal chat prompt.
*/
async function trimChatPrompt({ prompt, model, tokenLimit = model.contextWindowSize - model.maxCompletionTokens, }) {
(0, validateChatPrompt_js_1.validateChatPrompt)(prompt);
const startsWithSystemMessage = "system" in prompt[0];
const systemMessage = startsWithSystemMessage ? [prompt[0]] : [];
let messages = [];
// add the last message (final user message) to the prompt
messages.push(prompt[prompt.length - 1]);
// check if the minimal prompt is already too long
const promptTokenCount = await model.countPromptTokens([
...systemMessage,
...messages,
]);
// the minimal chat prompt is already over the token limit and cannot be trimmed further:
if (promptTokenCount > tokenLimit) {
return [...systemMessage, prompt[prompt.length - 1]];
}
// inner messages
const innerMessages = prompt.slice(startsWithSystemMessage ? 1 : 0, -1);
// taking always a pair of user-message and ai-message from the end, moving backwards
for (let i = innerMessages.length - 1; i >= 0; i -= 2) {
const aiMessage = innerMessages[i];
const userMessage = innerMessages[i - 1];
// create a temporary array and check if it fits within the token limit
const tokenCount = await model.countPromptTokens([
...systemMessage,
userMessage,
aiMessage,
...messages,
]);
if (tokenCount > tokenLimit) {
break;
}
// if it fits, add the messages to the messages array
messages = [userMessage, aiMessage, ...messages];
}
return [...systemMessage, ...messages];
}
exports.trimChatPrompt = trimChatPrompt;