Merge pull request #5310 from 0x23d11/fix-4947

perf(llm): Optimize pruneLines functions in countTokens

Author: sestinj

core/llm/countTokens.test.ts

@@ -46,6 +46,39 @@ describe.skip("pruneLinesFromTop", () => {
     const pruned = pruneLinesFromTop(prompt, 5, "gpt-4");
     expect(pruned.split("\n").length).toBeLessThan(prompt.split("\n").length);
   });
+
+  it("should return the original prompt if it's within max tokens", () => {
+    const prompt = "Line 1\nLine 2";
+    const pruned = pruneLinesFromTop(prompt, 10, "gpt-4");
+    expect(pruned).toEqual(prompt);
+  });
+
+  it("should return an empty string if maxTokens is 0", () => {
+    const prompt = "Line 1\nLine 2\nLine 3\nLine 4";
+    const pruned = pruneLinesFromTop(prompt, 0, "gpt-4");
+    expect(pruned).toEqual("");
+  });
+
+  it("should handle an empty prompt string", () => {
+    const prompt = "";
+    const pruned = pruneLinesFromTop(prompt, 5, "gpt-4");
+    expect(pruned).toEqual("");
+  });
+
+  it("should handle a prompt with a single line that exceeds maxTokens", () => {
+    const prompt =
+      "This is a single long line that will exceed the token limit";
+    const pruned = pruneLinesFromTop(prompt, 5, "gpt-4");
+
+    expect(pruned).toEqual("");
+  });
+
+  it("should correctly prune when all lines together exceed maxTokens but individual lines do not", () => {
+    const prompt = "L1\nL2\nL3\nL4";
+
+    const pruned = pruneLinesFromTop(prompt, 5, "gpt-4");
+    expect(pruned).toEqual("L3\nL4");
+  });
 });
 
 describe.skip("pruneLinesFromBottom", () => {
@@ -54,6 +87,39 @@ describe.skip("pruneLinesFromBottom", () => {
     const pruned = pruneLinesFromBottom(prompt, 5, "gpt-4");
     expect(pruned.split("\n").length).toBeLessThan(prompt.split("\n").length);
   });
+
+  it("should return the original prompt if it's within max tokens", () => {
+    const prompt = "Line 1\nLine 2";
+    const pruned = pruneLinesFromBottom(prompt, 10, "gpt-4");
+    expect(pruned).toEqual(prompt);
+  });
+
+  it("should return an empty string if maxTokens is 0", () => {
+    const prompt = "Line 1\nLine 2\nLine 3\nLine 4";
+    const pruned = pruneLinesFromBottom(prompt, 0, "gpt-4");
+    expect(pruned).toEqual("");
+  });
+
+  it("should handle an empty prompt string", () => {
+    const prompt = "";
+    const pruned = pruneLinesFromBottom(prompt, 5, "gpt-4");
+    expect(pruned).toEqual("");
+  });
+
+  it("should handle a prompt with a single line that exceeds maxTokens", () => {
+    const prompt =
+      "This is a single long line that will exceed the token limit";
+    const pruned = pruneLinesFromBottom(prompt, 5, "gpt-4");
+
+    expect(pruned).toEqual("");
+  });
+
+  it("should correctly prune when all lines together exceed maxTokens but individual lines do not", () => {
+    const prompt = "L1\nL2\nL3\nL4";
+
+    const pruned = pruneLinesFromBottom(prompt, 5, "gpt-4");
+    expect(pruned).toEqual("L1\nL2");
+  });
 });
 
 describe.skip("pruneRawPromptFromTop", () => {

core/llm/countTokens.ts

@@ -215,27 +215,55 @@ function pruneLinesFromTop(
   maxTokens: number,
   modelName: string,
 ): string {
-  let totalTokens = countTokens(prompt, modelName);
   const lines = prompt.split("\n");
-  while (totalTokens > maxTokens && lines.length > 0) {
-    totalTokens -= countTokens(lines.shift()!, modelName);
+  // Preprocess tokens for all lines and cache them.
+  const lineTokens = lines.map((line) => countTokens(line, modelName));
+  let totalTokens = lineTokens.reduce((sum, tokens) => sum + tokens, 0);
+  let start = 0;
+  let currentLines = lines.length;
+
+  // Calculate initial token count including newlines
+  totalTokens += Math.max(0, currentLines - 1); // Add tokens for joining newlines
+
+  // Using indexes instead of array modifications.
+  // Remove lines from the top until the token count is within the limit.
+  while (totalTokens > maxTokens && start < currentLines) {
+    totalTokens -= lineTokens[start];
+    // Decrement token count for the removed line and its preceding/joining newline (if not the last line)
+    if (currentLines - start > 1) {
+      totalTokens--;
+    }
+    start++;
   }
 
-  return lines.join("\n");
+  return lines.slice(start).join("\n");
 }
 
 function pruneLinesFromBottom(
   prompt: string,
   maxTokens: number,
   modelName: string,
 ): string {
-  let totalTokens = countTokens(prompt, modelName);
   const lines = prompt.split("\n");
-  while (totalTokens > maxTokens && lines.length > 0) {
-    totalTokens -= countTokens(lines.pop()!, modelName);
+  const lineTokens = lines.map((line) => countTokens(line, modelName));
+  let totalTokens = lineTokens.reduce((sum, tokens) => sum + tokens, 0);
+  let end = lines.length;
+
+  // Calculate initial token count including newlines
+  totalTokens += Math.max(0, end - 1); // Add tokens for joining newlines
+
+  // Reverse traversal to avoid array modification
+  // Remove lines from the bottom until the token count is within the limit.
+  while (totalTokens > maxTokens && end > 0) {
+    end--;
+    totalTokens -= lineTokens[end];
+    // Decrement token count for the removed line and its following/joining newline (if not the first line)
+    if (end > 0) {
+      totalTokens--;
+    }
   }
 
-  return lines.join("\n");
+  return lines.slice(0, end).join("\n");
 }
 
 function pruneStringFromBottom(
@@ -452,5 +480,5 @@ export {
   pruneLinesFromTop,
   pruneRawPromptFromTop,
   pruneStringFromBottom,
-  pruneStringFromTop,
+  pruneStringFromTop
 };