fix: keep text when truncating zero characters in context shift

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([

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});
+    });
+});

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("<x>"), "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("<x>"), "def"]);
+
+            expect(truncateLlamaTextAndRoundToWords(text, 3, undefined, false).values)
+                .to.eql([...LlamaText(["abc", new SpecialTokensText("<x>")]).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("<x>"), "def"]);
+
+            expect(truncateLlamaTextAndRoundToWords(text, 3, undefined, true).values)
+                .to.eql([...LlamaText([new SpecialTokensText("<x>"), "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("<x>"), "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([]);
+        });
+    });
+});