perf(llm): Optimize pruneLines functions in countTokens

[0x23d11]

Description

Closes #4947

This PR addresses issue #4947 by optimizing the performance of the pruneLinesFromTop and pruneLinesFromBottom functions in core/llm/countTokens.ts

Problem

The previous implementations used Array.prototype.shift() and Array.prototype.pop() within a while loop to remove lines from the beginning or end of a prompt until it fit within the token limit. These array modification methods have a time complexity of O(n) because they require shifting subsequent elements. When dealing with very long prompts (e.g., thousands of lines), repeatedly calling shift() or pop() becomes computationally expensive, leading to significant performance degradation.

Solution

The implemented solution refactors these functions to avoid costly array modifications within the loop:

1. The prompt is split into lines.
2. The token count for each line is calculated once upfront and stored in an array (lineTokens).
3. The total initial token count is calculated by summing lineTokens and adding the count for necessary newline characters (\n).
4. A while loop iterates as long as the totalTokens exceeds maxTokens.
5. Inside the loop, instead of removing elements from the lines array, an index pointer (start or end) is adjusted.
6. The pre-calculated token count for the line being (conceptually) removed, along with its corresponding newline token, is subtracted from totalTokens.
7. After the loop, Array.prototype.slice() (an O(n) operation performed only once) is used with the final start or end index to extract the desired lines.
8. The resulting lines are joined back into a string.

Benefits

This approach drastically reduces the computational complexity, especially for large prompts, as the expensive O(n) operations (shift/pop) inside the loop are replaced by cheap O(1) index increments/decrements. The token calculation per line and the final slice operation are performed only once.

Checklist

• I've read the contributing guide
• [] The relevant docs, if any, have been updated or created
• [] The relevant tests, if any, have been updated or created

Screenshots

[ For visual changes, include screenshots. Screen recordings are particularly helpful, and appreciated! ]

Testing instructions

Objective:

The primary goal of this PR is to optimize the performance of the pruneLinesFromTop and pruneLinesFromBottom utility functions. These functions are responsible for truncating a given string (prompt) by removing lines from the top or bottom, respectively, to ensure the resulting string does not exceed a specified maxTokens limit when processed by an LLM. This testing aims to verify that the optimized functions:
Correctly prune lines to meet the maxTokens constraint.
Maintain the existing logical behavior (i.e., correctly choose which lines to remove and which to keep).
Handle various edge cases appropriately.
Instructions:

1. Automated Unit Tests (Primary Validation):

The most direct way to test these changes is by running the unit tests specifically written for these functions. We've recently added comprehensive tests for various scenarios.
Prerequisites: Ensure your development environment is set up and you can run tests.
Execution:
Navigate to the root directory of the continue project in your terminal.
Run the following command to execute the test suite for countTokens.ts:

npm test -- core/llm/countTokens.test.ts

Expected Results:

• All tests within the describe("pruneLinesFromTop", ...) block should pass.
• All tests within the describe("pruneLinesFromBottom", ...) block should pass.

These tests cover:

• Basic pruning when the prompt exceeds maxTokens.
• Cases where the prompt is already within maxTokens (no pruning should occur).
• Edge cases such as an empty input prompt.
• Edge cases such as maxTokens being 0.
• Scenarios with single long lines that exceed maxTokens.
• Correct pruning of multi-line prompts to specific token counts (based on assumptions made in the test cases for token costs of lines and newlines).

[netlify]

✅ Deploy Preview for continuedev ready!

Name	Link
🔨 Latest commit	5543927
🔍 Latest deploy log	https://app.netlify.com/projects/continuedev/deploys/6835b604511a340008add33f
😎 Deploy Preview	https://deploy-preview-5310--continuedev.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.

---

To edit notification comments on pull requests, go to your Netlify project configuration.

[sestinj]

This is a fairly complex change that affects a lot of things, so if we are going to consider merging this I would ask that it come with very good testing. Are you able to add unit tests that cover edge cases and make sure that this hasn't regressed in any way?

[RomneyDa]

@0x23d11 before writing tests for this note that #5138 changes pruning logic quite a bit and will effect this.

EDIT: didn't look closely enough, looks like this is for pruning lines not messages

[RomneyDa]

@0x23d11 There are some tests in countTokens.test.ts that you could just flesh out a bit with more examples and then unskip!

[0x23d11]

@RomneyDa ok I've seen the tests, I'll write some more examples for the tests.

After that it should be good to go right? Considering all the new test additions are good enough

[sestinj]

@0x23d11 I'd love to merge this PR, please let me know if you have the chance to write some tests, or if you'd like any help!

[sestinj]

Just wanted to bump this. I'm happy to leave it open for a while, but tests are important here since it's a very core piece of logic, and relatively complex

[sestinj]

Hey @0x23d11, let me know if you have chance to look at this PR again. I'm probably going to close it if it becomes stale for longer. At this point just waiting on tests

[0x23d11]

Hi @sestinj sorry for the delay, I'll add the tests by this weekend

[sestinj]

@0x23d11 just wanted to bump this again. Let me know if you need any help

[github-actions]

All contributors have signed the CLA ✍️ ✅
_{Posted by the CLA Assistant Lite bot.}

[0x23d11]

I have read the CLA Document and I hereby sign the CLA

[0x23d11]

@sestinj please take a look, I've added the tests

[sestinj]

thanks for adding the tests, this looks great!