Skip to content

Commit 25ec92a

Browse files
ulugbeknaulugbekna
authored
nes: feat: add a setting to enforce max token count for edit window (#1764)
* nes: refactor: easier representation for tagged current file * nes: feat: add a setting to enforce max token count for edit window
1 parent 20f9418 commit 25ec92a

File tree

5 files changed

+51
-27
lines changed

5 files changed

+51
-27
lines changed

‎src/extension/xtab/common/promptCrafting.ts‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ export class PromptPieces {
2626
public readonly areaAroundEditWindowLinesRange: OffsetRange,
2727
public readonly activeDoc: StatelessNextEditDocument,
2828
public readonly xtabHistory: readonly IXtabHistoryEntry[],
29-
public readonly currentFileContent: string,
29+
public readonly taggedCurrentDocLines: readonly string[],
3030
public readonly areaAroundCodeToEdit: string,
3131
public readonly langCtx: LanguageContextResponse | undefined,
3232
public readonly computeTokens: (s: string) => number,
@@ -37,7 +37,8 @@ export class PromptPieces {
3737

3838
export function getUserPrompt(promptPieces: PromptPieces): string {
3939

40-
const { activeDoc, xtabHistory, currentFileContent, areaAroundCodeToEdit, langCtx, computeTokens, opts } = promptPieces;
40+
const { activeDoc, xtabHistory, taggedCurrentDocLines, areaAroundCodeToEdit, langCtx, computeTokens, opts } = promptPieces;
41+
const currentFileContent = taggedCurrentDocLines.join('\n');
4142

4243
const { codeSnippets: recentlyViewedCodeSnippets, documents: docsInPrompt } = getRecentCodeSnippets(activeDoc, xtabHistory, langCtx, computeTokens, opts);
4344

@@ -408,7 +409,7 @@ export function buildCodeSnippetsUsingPagedClipping(
408409
return { snippets: snippets.reverse(), docsInPrompt };
409410
}
410411

411-
function countTokensForLines(page: string[], computeTokens: (s: string) => number): number {
412+
export function countTokensForLines(page: string[], computeTokens: (s: string) => number): number {
412413
return page.reduce((sum, line) => sum + computeTokens(line) + 1 /* \n */, 0);
413414
}
414415

@@ -578,7 +579,7 @@ export function createTaggedCurrentFileContentUsingPagedClipping(
578579
computeTokens: (s: string) => number,
579580
pageSize: number,
580581
opts: CurrentFileOptions
581-
): Result<{ taggedCurrentFileContent: string; nLines: number }, 'outOfBudget'> {
582+
): Result<readonly string[], 'outOfBudget'> {
582583

583584
const r = clipPreservingRange(
584585
currentDocLines,
@@ -600,5 +601,5 @@ export function createTaggedCurrentFileContentUsingPagedClipping(
600601
...currentDocLines.slice(areaAroundEditWindowLinesRange.endExclusive, clippedRange.endExclusive),
601602
];
602603

603-
return Result.ok({ taggedCurrentFileContent: taggedCurrentFileContent.join('\n'), nLines: taggedCurrentFileContent.length });
604+
return Result.ok(taggedCurrentFileContent);
604605
}

‎src/extension/xtab/node/xtabProvider.ts‎

Lines changed: 13 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ import { Position as VscodePosition } from '../../../vscodeTypes';
5454
import { Delayer, DelaySession } from '../../inlineEdits/common/delayer';
5555
import { getOrDeduceSelectionFromLastEdit } from '../../inlineEdits/common/nearbyCursorInlineEditProvider';
5656
import { IgnoreImportChangesAspect } from '../../inlineEdits/node/importFiltering';
57-
import { createTaggedCurrentFileContentUsingPagedClipping, getUserPrompt, N_LINES_ABOVE, N_LINES_AS_CONTEXT, N_LINES_BELOW, PromptPieces } from '../common/promptCrafting';
57+
import { countTokensForLines, createTaggedCurrentFileContentUsingPagedClipping, getUserPrompt, N_LINES_ABOVE, N_LINES_AS_CONTEXT, N_LINES_BELOW, PromptPieces } from '../common/promptCrafting';
5858
import { nes41Miniv3SystemPrompt, simplifiedPrompt, systemPromptTemplate, unifiedModelSystemPrompt, xtab275SystemPrompt } from '../common/systemMessages';
5959
import { PromptTags } from '../common/tags';
6060
import { CurrentDocument } from '../common/xtabCurrentDocument';
@@ -271,6 +271,11 @@ export class XtabProvider implements IStatelessNextEditProvider {
271271

272272
const editWindowLines = currentDocument.lines.slice(editWindowLinesRange.start, editWindowLinesRange.endExclusive);
273273

274+
const editWindowTokenLimit = this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabEditWindowMaxTokens, this.expService);
275+
if (editWindowTokenLimit !== undefined && countTokensForLines(editWindowLines, XtabProvider.computeTokens) > editWindowTokenLimit) {
276+
return Result.error(new NoNextEditReason.PromptTooLarge('editWindow'));
277+
}
278+
274279
// Expected: editWindow.substring(activeDocument.documentAfterEdits.value) === editWindowLines.join('\n')
275280

276281
const doesIncludeCursorTag = editWindowLines.some(line => line.includes(PromptTags.CURSOR));
@@ -289,9 +294,9 @@ export class XtabProvider implements IStatelessNextEditProvider {
289294
return Result.error(new NoNextEditReason.PromptTooLarge('currentFile'));
290295
}
291296

292-
const { taggedCurrentFileR: { taggedCurrentFileContent, nLines: nLinesCurrentFile }, areaAroundCodeToEdit } = taggedCurrentFileContentResult.val;
297+
const { taggedCurrentDocLines, areaAroundCodeToEdit } = taggedCurrentFileContentResult.val;
293298

294-
telemetryBuilder.setNLinesOfCurrentFileInPrompt(nLinesCurrentFile);
299+
telemetryBuilder.setNLinesOfCurrentFileInPrompt(taggedCurrentDocLines.length);
295300

296301
const langCtx = await this.getAndProcessLanguageContext(
297302
request,
@@ -313,7 +318,7 @@ export class XtabProvider implements IStatelessNextEditProvider {
313318
areaAroundEditWindowLinesRange,
314319
activeDocument,
315320
request.xtabEditHistory,
316-
taggedCurrentFileContent,
321+
taggedCurrentDocLines,
317322
areaAroundCodeToEdit,
318323
langCtx,
319324
XtabProvider.computeTokens,
@@ -434,8 +439,8 @@ export class XtabProvider implements IStatelessNextEditProvider {
434439
promptOptions.currentFile,
435440
);
436441

437-
return taggedCurrentFileContentResult.map(taggedCurrentFileR => ({
438-
taggedCurrentFileR,
442+
return taggedCurrentFileContentResult.map(taggedCurrentDocLines => ({
443+
taggedCurrentDocLines,
439444
areaAroundCodeToEdit,
440445
}));
441446
}
@@ -1165,15 +1170,15 @@ export class XtabProvider implements IStatelessNextEditProvider {
11651170
return Result.fromString(currentFileContentR.err);
11661171
}
11671172

1168-
const { taggedCurrentFileR: { taggedCurrentFileContent }, areaAroundCodeToEdit } = currentFileContentR.val;
1173+
const { taggedCurrentDocLines, areaAroundCodeToEdit } = currentFileContentR.val;
11691174

11701175
const newPromptPieces = new PromptPieces(
11711176
promptPieces.currentDocument,
11721177
promptPieces.editWindowLinesRange,
11731178
promptPieces.areaAroundEditWindowLinesRange,
11741179
promptPieces.activeDoc,
11751180
promptPieces.xtabHistory,
1176-
taggedCurrentFileContent,
1181+
taggedCurrentDocLines,
11771182
areaAroundCodeToEdit,
11781183
promptPieces.langCtx,
11791184
XtabProvider.computeTokens,

‎src/extension/xtab/test/common/promptCrafting.spec.ts‎

Lines changed: 30 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import { assert, expect, suite, test } from 'vitest';
77
import { DocumentId } from '../../../../platform/inlineEdits/common/dataTypes/documentId';
88
import { CurrentFileOptions, DEFAULT_OPTIONS } from '../../../../platform/inlineEdits/common/dataTypes/xtabPromptOptions';
9+
import { Result } from '../../../../util/common/result';
910
import { OffsetRange } from '../../../../util/vs/editor/common/core/ranges/offsetRange';
1011
import { StringText } from '../../../../util/vs/editor/common/core/text/abstractText';
1112
import { buildCodeSnippetsUsingPagedClipping, createTaggedCurrentFileContentUsingPagedClipping } from '../../common/promptCrafting';
@@ -97,6 +98,25 @@ suite('Paged clipping - current file', () => {
9798

9899
const opts: CurrentFileOptions = DEFAULT_OPTIONS.currentFile;
99100

101+
function createTaggedFile(
102+
currentDocLines: string[],
103+
areaAroundCodeToEdit: string,
104+
areaAroundEditWindowLinesRange: OffsetRange,
105+
computeTokens: (s: string) => number,
106+
pageSize: number,
107+
opts: CurrentFileOptions,
108+
): Result<string, 'outOfBudget'> {
109+
110+
return createTaggedCurrentFileContentUsingPagedClipping(
111+
currentDocLines,
112+
areaAroundCodeToEdit,
113+
areaAroundEditWindowLinesRange,
114+
computeTokens,
115+
pageSize,
116+
opts
117+
).map(taggedCurrentFileContent => taggedCurrentFileContent.join('\n'));
118+
}
119+
100120
test('unlim budget - includes whole context', () => {
101121

102122
const docLines = nLines(40);
@@ -113,7 +133,7 @@ suite('Paged clipping - current file', () => {
113133
</area_around_code_to_edit>
114134
`.trim();
115135

116-
const result = createTaggedCurrentFileContentUsingPagedClipping(
136+
const result = createTaggedFile(
117137
docLines.getLines(),
118138
areaAroundCodeToEdit,
119139
new OffsetRange(21, 26),
@@ -122,7 +142,7 @@ suite('Paged clipping - current file', () => {
122142
{ ...opts, maxTokens: 2000 }
123143
);
124144
assert(result.isOk());
125-
const { taggedCurrentFileContent } = result.val;
145+
const taggedCurrentFileContent = result.val;
126146

127147
expect(taggedCurrentFileContent).toMatchInlineSnapshot(`
128148
"1
@@ -189,19 +209,18 @@ suite('Paged clipping - current file', () => {
189209
</area_around_code_to_edit>
190210
`.trim();
191211

192-
const result = createTaggedCurrentFileContentUsingPagedClipping(
212+
const result = createTaggedFile(
193213
docLines.getLines(),
194214
areaAroundCodeToEdit,
195215
new OffsetRange(21, 26),
196216
computeTokens,
197217
10,
198218
{ ...opts, maxTokens: 20 },
199219
);
200-
expect(result).toMatchInlineSnapshot(`
201-
ResultOk {
202-
"val": {
203-
"nLines": 6,
204-
"taggedCurrentFileContent": "21
220+
assert(result.isOk());
221+
const taggedCurrentFileContent = result.val;
222+
expect(taggedCurrentFileContent).toMatchInlineSnapshot(`
223+
"21
205224
<area_around_code_to_edit>
206225
22
207226
23
@@ -214,9 +233,7 @@ suite('Paged clipping - current file', () => {
214233
27
215234
28
216235
29
217-
30",
218-
},
219-
}
236+
30"
220237
`);
221238
});
222239

@@ -236,7 +253,7 @@ suite('Paged clipping - current file', () => {
236253
</a>
237254
`.trim();
238255

239-
const result = createTaggedCurrentFileContentUsingPagedClipping(
256+
const result = createTaggedFile(
240257
docLines.getLines(),
241258
areaAroundCodeToEdit,
242259
new OffsetRange(10, 14),
@@ -245,7 +262,7 @@ suite('Paged clipping - current file', () => {
245262
{ ...opts, maxTokens: 50 }
246263
);
247264
assert(result.isOk());
248-
const { taggedCurrentFileContent } = result.val;
265+
const taggedCurrentFileContent = result.val;
249266

250267
expect(taggedCurrentFileContent).toMatchInlineSnapshot(`
251268
"<a>

‎src/platform/configuration/common/configurationService.ts‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -683,6 +683,7 @@ export namespace ConfigKey {
683683
export const InlineEditsXtabProviderEmitFastCursorLineChange = defineExpSetting<boolean>('chat.advanced.inlineEdits.xtabProvider.emitFastCursorLineChange', true, INTERNAL_RESTRICTED);
684684
export const InlineEditsXtabIncludeViewedFiles = defineExpSetting<boolean>('chat.advanced.inlineEdits.xtabProvider.includeViewedFiles', xtabPromptOptions.DEFAULT_OPTIONS.recentlyViewedDocuments.includeViewedFiles, INTERNAL_RESTRICTED);
685685
export const InlineEditsXtabPageSize = defineExpSetting<number>('chat.advanced.inlineEdits.xtabProvider.pageSize', xtabPromptOptions.DEFAULT_OPTIONS.pagedClipping.pageSize, INTERNAL_RESTRICTED);
686+
export const InlineEditsXtabEditWindowMaxTokens = defineExpSetting<number | undefined>('chat.advanced.inlineEdits.xtabProvider.editWindowMaxTokens', undefined, INTERNAL_RESTRICTED);
686687
export const InlineEditsXtabIncludeTagsInCurrentFile = defineExpSetting<boolean>('chat.advanced.inlineEdits.xtabProvider.includeTagsInCurrentFile', xtabPromptOptions.DEFAULT_OPTIONS.currentFile.includeTags, INTERNAL_RESTRICTED);
687688
export const InlineEditsXtabCurrentFileMaxTokens = defineExpSetting<number>('chat.advanced.inlineEdits.xtabProvider.currentFileMaxTokens', xtabPromptOptions.DEFAULT_OPTIONS.currentFile.maxTokens, INTERNAL_RESTRICTED);
688689
export const InlineEditsXtabPrioritizeAboveCursor = defineExpSetting<boolean>('chat.advanced.inlineEdits.xtabProvider.currentFile.prioritizeAboveCursor', xtabPromptOptions.DEFAULT_OPTIONS.currentFile.prioritizeAboveCursor, INTERNAL_RESTRICTED);

‎src/platform/inlineEdits/common/statelessNextEditProvider.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -213,7 +213,7 @@ export namespace NoNextEditReason {
213213
}
214214
export class PromptTooLarge {
215215
public readonly kind = 'promptTooLarge';
216-
constructor(public readonly message: 'currentFile' | 'final') {
216+
constructor(public readonly message: 'editWindow' | 'currentFile' | 'final') {
217217
}
218218
}
219219
export class Uncategorized {

0 commit comments

Comments
 (0)