@@ -7,12 +7,15 @@ import assert from 'assert'
7
7
import * as FakeTimers from '@sinonjs/fake-timers'
8
8
import * as vscode from 'vscode'
9
9
import * as sinon from 'sinon'
10
+ import * as os from 'os'
10
11
import * as crossFile from 'aws-core-vscode/codewhisperer'
11
12
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'
13
14
import { toTextEditor } from 'aws-core-vscode/test'
14
15
import { LspController } from 'aws-core-vscode/amazonq'
15
16
17
+ const newLine = os . EOL
18
+
16
19
describe ( 'supplementalContextUtil' , function ( ) {
17
20
let testFolder : TestFolder
18
21
let clock : FakeTimers . InstalledClock
@@ -83,4 +86,180 @@ describe('supplementalContextUtil', function () {
83
86
} )
84
87
} )
85
88
} )
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
+ } )
86
256
} )
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
+ }
0 commit comments