Skip to content

Commit 070c36d

Browse files
Will-ShaoHuaWill-ShaoHua
authored
fix(inline-completion): potential inline completion failure due to input validation exception of supplemental context (aws#6758)
## Problem ERROR: `at 'supplementalContexts' failed to satisfy constraint: Member must have length less than or equal to 5` ERROR: `supplementalContexts.1.member.content' failed to satisfy constraint: Member must have length less than or equal to 10240` ## Solution Jebtrains PR aws/aws-toolkit-jetbrains#5466 * Requirement * - Maximum 5 supplemental context. * - Each chunk can't exceed 10240 characters * - Sum of all chunks can't exceed 20480 characters --- - Treat all work as PUBLIC. Private `feature/x` branches will not be squash-merged at release time. - Your code changes must meet the guidelines in [CONTRIBUTING.md](https://github.com/aws/aws-toolkit-vscode/blob/master/CONTRIBUTING.md#guidelines). - License: I confirm that my contribution is made under the terms of the Apache 2.0 license.
1 parent 6e367e9 commit 070c36d

File tree

4 files changed

+254
-2
lines changed

4 files changed

+254
-2
lines changed

‎packages/amazonq/.changes/next-release/Bug Fix-bde47093-ee34-4802-a556-d1e5dc526f4d.json‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
{
2+
"type": "Bug Fix",
3+
"description": "Fix inline completion failure due to context length exceeding the threshold"
4+
}

‎packages/amazonq/test/unit/codewhisperer/util/supplemetalContextUtil.test.ts‎

Lines changed: 180 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,15 @@ import assert from 'assert'
77
import * as FakeTimers from '@sinonjs/fake-timers'
88
import * as vscode from 'vscode'
99
import * as sinon from 'sinon'
10+
import * as os from 'os'
1011
import * as crossFile from 'aws-core-vscode/codewhisperer'
1112
import { TestFolder, assertTabCount, installFakeClock } from 'aws-core-vscode/test'
12-
import { FeatureConfigProvider } from 'aws-core-vscode/codewhisperer'
13+
import { CodeWhispererSupplementalContext, FeatureConfigProvider } from 'aws-core-vscode/codewhisperer'
1314
import { toTextEditor } from 'aws-core-vscode/test'
1415
import { LspController } from 'aws-core-vscode/amazonq'
1516

17+
const newLine = os.EOL
18+
1619
describe('supplementalContextUtil', function () {
1720
let testFolder: TestFolder
1821
let clock: FakeTimers.InstalledClock
@@ -83,4 +86,180 @@ describe('supplementalContextUtil', function () {
8386
})
8487
})
8588
})
89+
90+
describe('truncation', function () {
91+
it('truncate context should do nothing if everything fits in constraint', function () {
92+
const chunkA: crossFile.CodeWhispererSupplementalContextItem = {
93+
content: 'a',
94+
filePath: 'a.java',
95+
score: 0,
96+
}
97+
const chunkB: crossFile.CodeWhispererSupplementalContextItem = {
98+
content: 'b',
99+
filePath: 'b.java',
100+
score: 1,
101+
}
102+
const chunks = [chunkA, chunkB]
103+
104+
const supplementalContext: CodeWhispererSupplementalContext = {
105+
isUtg: false,
106+
isProcessTimeout: false,
107+
supplementalContextItems: chunks,
108+
contentsLength: 25000,
109+
latency: 0,
110+
strategy: 'codemap',
111+
}
112+
113+
const actual = crossFile.truncateSuppelementalContext(supplementalContext)
114+
assert.strictEqual(actual.supplementalContextItems.length, 2)
115+
assert.strictEqual(actual.supplementalContextItems[0].content, 'a')
116+
assert.strictEqual(actual.supplementalContextItems[1].content, 'b')
117+
})
118+
119+
it('truncateLineByLine should drop the last line if max length is greater than threshold', function () {
120+
const input =
121+
repeatString('a', 11) +
122+
newLine +
123+
repeatString('b', 11) +
124+
newLine +
125+
repeatString('c', 11) +
126+
newLine +
127+
repeatString('d', 11) +
128+
newLine +
129+
repeatString('e', 11)
130+
131+
assert.ok(input.length > 50)
132+
const actual = crossFile.truncateLineByLine(input, 50)
133+
assert.ok(actual.length <= 50)
134+
135+
const input2 = repeatString(`b${newLine}`, 10)
136+
const actual2 = crossFile.truncateLineByLine(input2, 8)
137+
assert.ok(actual2.length <= 8)
138+
})
139+
140+
it('truncation context should make context length per item lte 10240 cap', function () {
141+
const chunkA: crossFile.CodeWhispererSupplementalContextItem = {
142+
content: repeatString(`a${newLine}`, 4000),
143+
filePath: 'a.java',
144+
score: 0,
145+
}
146+
const chunkB: crossFile.CodeWhispererSupplementalContextItem = {
147+
content: repeatString(`b${newLine}`, 6000),
148+
filePath: 'b.java',
149+
score: 1,
150+
}
151+
const chunkC: crossFile.CodeWhispererSupplementalContextItem = {
152+
content: repeatString(`c${newLine}`, 1000),
153+
filePath: 'c.java',
154+
score: 2,
155+
}
156+
const chunkD: crossFile.CodeWhispererSupplementalContextItem = {
157+
content: repeatString(`d${newLine}`, 1500),
158+
filePath: 'd.java',
159+
score: 3,
160+
}
161+
162+
assert.ok(
163+
chunkA.content.length + chunkB.content.length + chunkC.content.length + chunkD.content.length > 20480
164+
)
165+
166+
const supplementalContext: CodeWhispererSupplementalContext = {
167+
isUtg: false,
168+
isProcessTimeout: false,
169+
supplementalContextItems: [chunkA, chunkB, chunkC, chunkD],
170+
contentsLength: 25000,
171+
latency: 0,
172+
strategy: 'codemap',
173+
}
174+
175+
const actual = crossFile.truncateSuppelementalContext(supplementalContext)
176+
assert.strictEqual(actual.supplementalContextItems.length, 3)
177+
assert.ok(actual.contentsLength <= 20480)
178+
assert.strictEqual(actual.strategy, 'codemap')
179+
})
180+
181+
it('truncate context should make context items lte 5', function () {
182+
const chunkA: crossFile.CodeWhispererSupplementalContextItem = {
183+
content: 'a',
184+
filePath: 'a.java',
185+
score: 0,
186+
}
187+
const chunkB: crossFile.CodeWhispererSupplementalContextItem = {
188+
content: 'b',
189+
filePath: 'b.java',
190+
score: 1,
191+
}
192+
const chunkC: crossFile.CodeWhispererSupplementalContextItem = {
193+
content: 'c',
194+
filePath: 'c.java',
195+
score: 2,
196+
}
197+
const chunkD: crossFile.CodeWhispererSupplementalContextItem = {
198+
content: 'd',
199+
filePath: 'd.java',
200+
score: 3,
201+
}
202+
const chunkE: crossFile.CodeWhispererSupplementalContextItem = {
203+
content: 'e',
204+
filePath: 'e.java',
205+
score: 4,
206+
}
207+
const chunkF: crossFile.CodeWhispererSupplementalContextItem = {
208+
content: 'f',
209+
filePath: 'f.java',
210+
score: 5,
211+
}
212+
const chunkG: crossFile.CodeWhispererSupplementalContextItem = {
213+
content: 'g',
214+
filePath: 'g.java',
215+
score: 6,
216+
}
217+
const chunks = [chunkA, chunkB, chunkC, chunkD, chunkE, chunkF, chunkG]
218+
219+
assert.strictEqual(chunks.length, 7)
220+
221+
const supplementalContext: CodeWhispererSupplementalContext = {
222+
isUtg: false,
223+
isProcessTimeout: false,
224+
supplementalContextItems: chunks,
225+
contentsLength: 25000,
226+
latency: 0,
227+
strategy: 'codemap',
228+
}
229+
230+
const actual = crossFile.truncateSuppelementalContext(supplementalContext)
231+
assert.strictEqual(actual.supplementalContextItems.length, 5)
232+
})
233+
234+
describe('truncate line by line', function () {
235+
it('should return empty if empty string is provided', function () {
236+
const input = ''
237+
const actual = crossFile.truncateLineByLine(input, 50)
238+
assert.strictEqual(actual, '')
239+
})
240+
241+
it('should return empty if 0 max length is provided', function () {
242+
const input = 'aaaaa'
243+
const actual = crossFile.truncateLineByLine(input, 0)
244+
assert.strictEqual(actual, '')
245+
})
246+
247+
it('should flip the value if negative max length is provided', function () {
248+
const input = `aaaaa${newLine}bbbbb`
249+
const actual = crossFile.truncateLineByLine(input, -6)
250+
const expected = crossFile.truncateLineByLine(input, 6)
251+
assert.strictEqual(actual, expected)
252+
assert.strictEqual(actual, 'aaaaa')
253+
})
254+
})
255+
})
86256
})
257+
258+
function repeatString(s: string, n: number): string {
259+
let output = ''
260+
for (let i = 0; i < n; i++) {
261+
output += s
262+
}
263+
264+
return output
265+
}

‎packages/core/src/codewhisperer/models/constants.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -856,6 +856,8 @@ export const crossFileContextConfig = {
856856
topK: 3,
857857
numberOfLinesEachChunk: 50,
858858
maximumTotalLength: 20480,
859+
maxLengthEachChunk: 10240,
860+
maxContextCount: 5,
859861
}
860862

861863
export const utgConfig = {

‎packages/core/src/codewhisperer/util/supplementalContext/supplementalContextUtil.ts‎

Lines changed: 68 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ import { CancellationError } from '../../../shared/utilities/timeoutUtils'
1111
import { ToolkitError } from '../../../shared/errors'
1212
import { getLogger } from '../../../shared/logger/logger'
1313
import { CodeWhispererSupplementalContext } from '../../models/model'
14+
import * as os from 'os'
15+
import { crossFileContextConfig } from '../../models/constants'
1416

1517
export async function fetchSupplementalContext(
1618
editor: vscode.TextEditor,
@@ -36,7 +38,7 @@ export async function fetchSupplementalContext(
3638
return supplementalContextPromise
3739
.then((value) => {
3840
if (value) {
39-
return {
41+
const resBeforeTruncation = {
4042
isUtg: isUtg,
4143
isProcessTimeout: false,
4244
supplementalContextItems: value.supplementalContextItems.filter(
@@ -46,6 +48,8 @@ export async function fetchSupplementalContext(
4648
latency: performance.now() - timesBeforeFetching,
4749
strategy: value.strategy,
4850
}
51+
52+
return truncateSuppelementalContext(resBeforeTruncation)
4953
} else {
5054
return undefined
5155
}
@@ -68,3 +72,66 @@ export async function fetchSupplementalContext(
6872
}
6973
})
7074
}
75+
76+
/**
77+
* Requirement
78+
* - Maximum 5 supplemental context.
79+
* - Each chunk can't exceed 10240 characters
80+
* - Sum of all chunks can't exceed 20480 characters
81+
*/
82+
export function truncateSuppelementalContext(
83+
context: CodeWhispererSupplementalContext
84+
): CodeWhispererSupplementalContext {
85+
let c = context.supplementalContextItems.map((item) => {
86+
if (item.content.length > crossFileContextConfig.maxLengthEachChunk) {
87+
return {
88+
...item,
89+
content: truncateLineByLine(item.content, crossFileContextConfig.maxLengthEachChunk),
90+
}
91+
} else {
92+
return item
93+
}
94+
})
95+
96+
if (c.length > crossFileContextConfig.maxContextCount) {
97+
c = c.slice(0, crossFileContextConfig.maxContextCount)
98+
}
99+
100+
let curTotalLength = c.reduce((acc, cur) => {
101+
return acc + cur.content.length
102+
}, 0)
103+
while (curTotalLength >= 20480 && c.length - 1 >= 0) {
104+
const last = c[c.length - 1]
105+
c = c.slice(0, -1)
106+
curTotalLength -= last.content.length
107+
}
108+
109+
return {
110+
...context,
111+
supplementalContextItems: c,
112+
contentsLength: curTotalLength,
113+
}
114+
}
115+
116+
export function truncateLineByLine(input: string, l: number): string {
117+
const maxLength = l > 0 ? l : -1 * l
118+
if (input.length === 0) {
119+
return ''
120+
}
121+
122+
const shouldAddNewLineBack = input.endsWith(os.EOL)
123+
let lines = input.trim().split(os.EOL)
124+
let curLen = input.length
125+
while (curLen > maxLength && lines.length - 1 >= 0) {
126+
const last = lines[lines.length - 1]
127+
lines = lines.slice(0, -1)
128+
curLen -= last.length + 1
129+
}
130+
131+
const r = lines.join(os.EOL)
132+
if (shouldAddNewLineBack) {
133+
return r + os.EOL
134+
} else {
135+
return r
136+
}
137+
}

0 commit comments

Comments
 (0)