continue
continue copied to clipboard
perf(llm): Optimize pruneLines functions in countTokens
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:
- The prompt is split into lines.
- The token count for each line is calculated once upfront and stored in an array (
lineTokens). - The total initial token count is calculated by summing
lineTokensand adding the count for necessary newline characters (\n). - A
whileloop iterates as long as thetotalTokensexceedsmaxTokens. - Inside the loop, instead of removing elements from the
linesarray, an index pointer (startorend) is adjusted. - The pre-calculated token count for the line being (conceptually) removed, along with its corresponding newline token, is subtracted from
totalTokens. - After the loop,
Array.prototype.slice()(an O(n) operation performed only once) is used with the finalstartorendindex to extract the desired lines. - 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
- [x] 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
[ For new or modified features, provide step-by-step testing instructions to validate the intended behavior of the change, including any relevant tests to run. ]