From 871f75c1dbac01c91467c25b2a526ae2c338ce5d Mon Sep 17 00:00:00 2001 From: Osama Ansar Date: Wed, 28 Oct 2026 16:01:41 +0500 Subject: [PATCH] fix: keep text when the truncate size is 0 in context shift `truncateTextAndRoundToWords` used `text.slice(0, -truncateSize)`, which is an empty string when `truncateSize` is 0, so truncating nothing from the end deleted the whole text (the `truncateStart` branch already returned the text as is). A context shift that clamped the number of characters to remove to 0 therefore dropped chat history items it meant to keep. Return the text as is when there is nothing to truncate (including negative sizes), and keep the remaining items of a `LlamaText` once the requested size has been fully consumed. --- src/utils/truncateTextAndRoundToWords.ts | 12 +++ ...irstSystemChatContextShiftStrategy.test.ts | 40 ++++++++++ .../utils/truncateTextAndRoundToWords.test.ts | 80 +++++++++++++++++++ 3 files changed, 132 insertions(+) create mode 100644 test/standalone/utils/eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy.test.ts create mode 100644 test/standalone/utils/truncateTextAndRoundToWords.test.ts diff --git a/src/utils/truncateTextAndRoundToWords.ts b/src/utils/truncateTextAndRoundToWords.ts index 26288ef1..6cb23a0b 100644 --- a/src/utils/truncateTextAndRoundToWords.ts +++ b/src/utils/truncateTextAndRoundToWords.ts @@ -13,6 +13,10 @@ const truncatePrefix = "..."; export function truncateTextAndRoundToWords( text: string, truncateSize: number, maxRound: number = 6, truncateStart: boolean = false ): string { + // nothing to truncate. Note that `text.slice(0, -0)` is an empty string, so this has to be handled explicitly + if (truncateSize <= 0) + return text; + if (truncateStart) { const res = text.slice(truncateSize); @@ -79,6 +83,10 @@ export function truncateLlamaTextAndRoundToWords( if (value == null) continue; + // everything that had to be truncated was already truncated, so keep the rest as is + if (truncateSize <= 0) + return LlamaText(llamaText.values.slice(i)); + if (typeof value === "string") { if (value.length > truncateSize) { return LlamaText([ @@ -109,6 +117,10 @@ export function truncateLlamaTextAndRoundToWords( if (value == null) continue; + // everything that had to be truncated was already truncated, so keep the rest as is + if (truncateSize <= 0) + return LlamaText(llamaText.values.slice(0, i + 1)); + if (typeof value === "string") { if (value.length > truncateSize) { return LlamaText([ diff --git a/test/standalone/utils/eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy.test.ts b/test/standalone/utils/eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy.test.ts new file mode 100644 index 00000000..f190cb9e --- /dev/null +++ b/test/standalone/utils/eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy.test.ts @@ -0,0 +1,40 @@ +import {describe, expect, test} from "vitest"; +import {GeneralChatWrapper} from "../../../src/chatWrappers/GeneralChatWrapper.js"; +import { + eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy +} from "../../../src/evaluator/LlamaChat/utils/contextShiftStrategies/eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy.js"; +import {ChatHistoryItem, Tokenizer} from "../../../src/types.js"; + + +// a fake tokenizer that counts roughly one token per word, so the test doesn't need a model +const tokenizer = ((text: string) => { + return Array.from(text.matchAll(/\S+\s*/g)).map((_match, index) => index + 1); +}) as unknown as Tokenizer; +(tokenizer as unknown as {detokenize: () => string}).detokenize = () => ""; + +const words = (count: number, prefix: string) => Array.from({length: count}, (_, i) => prefix + i).join(" "); + +describe("eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy", () => { + test("does not drop a user message that was meant to be kept entirely", async () => { + const firstUserMessage = words(60, "ua"); + const chatHistory: ChatHistoryItem[] = [ + {type: "system", text: "You are a helpful assistant."}, + {type: "user", text: firstUserMessage}, + {type: "user", text: words(60, "ub")}, + {type: "user", text: words(60, "uc")}, + {type: "model", response: [words(300, "mc")]} + ]; + + const res = await eraseFirstResponseAndKeepFirstSystemChatContextShiftStrategy({ + chatHistory, + maxTokensCount: 400, + tokenizer, + chatWrapper: new GeneralChatWrapper() + }); + + const userMessages = res.chatHistory.filter((item) => item.type === "user"); + expect(res.chatHistory[0]).to.eql(chatHistory[0]); + expect(userMessages.length).toBeGreaterThanOrEqual(2); + expect(userMessages[0]).to.eql({type: "user", text: firstUserMessage}); + }); +}); diff --git a/test/standalone/utils/truncateTextAndRoundToWords.test.ts b/test/standalone/utils/truncateTextAndRoundToWords.test.ts new file mode 100644 index 00000000..ce4ebad0 --- /dev/null +++ b/test/standalone/utils/truncateTextAndRoundToWords.test.ts @@ -0,0 +1,80 @@ +import {describe, expect, test} from "vitest"; +import {truncateLlamaTextAndRoundToWords, truncateTextAndRoundToWords} from "../../../src/utils/truncateTextAndRoundToWords.js"; +import {LlamaText, SpecialToken, SpecialTokensText} from "../../../src/utils/LlamaText.js"; + + +describe("utils", () => { + describe("truncateTextAndRoundToWords", () => { + test("truncate size of 0 keeps the whole text", () => { + expect(truncateTextAndRoundToWords("hello world", 0, undefined, false)).toBe("hello world"); + expect(truncateTextAndRoundToWords("hello world", 0, undefined, true)).toBe("hello world"); + expect(truncateTextAndRoundToWords("hello world", 0)).toBe("hello world"); + }); + + test("negative truncate size keeps the whole text", () => { + expect(truncateTextAndRoundToWords("hello world", -3, undefined, false)).toBe("hello world"); + expect(truncateTextAndRoundToWords("hello world", -3, undefined, true)).toBe("hello world"); + }); + + test("empty text with truncate size of 0", () => { + expect(truncateTextAndRoundToWords("", 0, undefined, false)).toBe(""); + expect(truncateTextAndRoundToWords("", 0, undefined, true)).toBe(""); + }); + + test("truncating the entire text results in an empty string", () => { + expect(truncateTextAndRoundToWords("hello world", 11, undefined, false)).toBe(""); + expect(truncateTextAndRoundToWords("hello world", 11, undefined, true)).toBe(""); + expect(truncateTextAndRoundToWords("hello world", 50, undefined, false)).toBe(""); + expect(truncateTextAndRoundToWords("hello world", 50, undefined, true)).toBe(""); + }); + + test("truncate on a word boundary", () => { + expect(truncateTextAndRoundToWords("hello world", 6, undefined, false)).toBe("hello"); + expect(truncateTextAndRoundToWords("hello world", 6, undefined, true)).toBe("world"); + }); + }); + + describe("truncateLlamaTextAndRoundToWords", () => { + test("truncate size of 0 or less keeps the whole text", () => { + const text = LlamaText(["abc", new SpecialTokensText(""), "def"]); + + expect(truncateLlamaTextAndRoundToWords(text, 0, undefined, false).values).to.eql(text.values); + expect(truncateLlamaTextAndRoundToWords(text, 0, undefined, true).values).to.eql(text.values); + expect(truncateLlamaTextAndRoundToWords(text, -3, undefined, false).values).to.eql(text.values); + expect(truncateLlamaTextAndRoundToWords(text, -3, undefined, true).values).to.eql(text.values); + }); + + test("keeps the remaining items whole when the truncate size is consumed exactly at an item boundary (from end)", () => { + const text = LlamaText(["abc", new SpecialTokensText(""), "def"]); + + expect(truncateLlamaTextAndRoundToWords(text, 3, undefined, false).values) + .to.eql([...LlamaText(["abc", new SpecialTokensText("")]).values]); + }); + + test("keeps the remaining items whole when the truncate size is consumed exactly at an item boundary (from start)", () => { + const text = LlamaText(["abc", new SpecialTokensText(""), "def"]); + + expect(truncateLlamaTextAndRoundToWords(text, 3, undefined, true).values) + .to.eql([...LlamaText([new SpecialTokensText(""), "def"]).values]); + }); + + test("keeps a special token that follows an exactly consumed string", () => { + const token = new SpecialToken("BOS"); + const text = LlamaText(["abc", token, "def"]); + + expect(truncateLlamaTextAndRoundToWords(text, 3, undefined, false).values) + .to.eql([...LlamaText(["abc", token]).values]); + expect(truncateLlamaTextAndRoundToWords(text, 3, undefined, true).values) + .to.eql([...LlamaText([token, "def"]).values]); + }); + + test("still removes the items that fall within the truncate size", () => { + const text = LlamaText(["abc", new SpecialTokensText(""), "def"]); + + expect(truncateLlamaTextAndRoundToWords(text, 6, undefined, false).values).to.eql(["abc"]); + expect(truncateLlamaTextAndRoundToWords(text, 6, undefined, true).values).to.eql(["def"]); + expect(truncateLlamaTextAndRoundToWords(text, 100, undefined, false).values).to.eql([]); + expect(truncateLlamaTextAndRoundToWords(text, 100, undefined, true).values).to.eql([]); + }); + }); +});