Skip to content

Commit d8bff7c

Browse files
ulugbeknaulugbekna
committed
nes: refactor: extract tagging current file into its own method
1 parent 885cde9 commit d8bff7c

File tree

2 files changed

+114
-44
lines changed

2 files changed

+114
-44
lines changed

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

Lines changed: 38 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -665,40 +665,62 @@ function expandRangeToPageRange(
665665
return { firstPageIdx, lastPageIdx, budgetLeft: tokenBudget };
666666
}
667667

668-
/**
669-
* @remark exported for testing
670-
*/
671-
export function createTaggedCurrentFileContentUsingPagedClipping(
672-
currentDocLines: string[],
673-
areaAroundCodeToEdit: string,
674-
areaAroundEditWindowLinesRange: OffsetRange,
668+
export function clipPreservingRange(
669+
docLines: string[],
670+
rangeToPreserve: OffsetRange,
675671
computeTokens: (s: string) => number,
676672
pageSize: number,
677-
opts: CurrentFileOptions
678-
): Result<{ taggedCurrentFileContent: string; nLines: number }, 'outOfBudget'> {
673+
opts: CurrentFileOptions,
674+
): Result<OffsetRange, 'outOfBudget'> {
679675

680-
// subtract budget consumed by areaAroundCodeToEdit
681-
const availableTokenBudget = opts.maxTokens - countTokensForLines(areaAroundCodeToEdit.split(/\r?\n/), computeTokens);
676+
// subtract budget consumed by rangeToPreserve
677+
const availableTokenBudget = opts.maxTokens - countTokensForLines(docLines.slice(rangeToPreserve.start, rangeToPreserve.endExclusive), computeTokens);
682678
if (availableTokenBudget < 0) {
683679
return Result.error('outOfBudget');
684680
}
685681

686682
const { firstPageIdx, lastPageIdx } = expandRangeToPageRange(
687-
currentDocLines,
688-
areaAroundEditWindowLinesRange,
683+
docLines,
684+
rangeToPreserve,
689685
pageSize,
690686
availableTokenBudget,
691687
computeTokens,
692688
opts.prioritizeAboveCursor,
693689
);
694690

695691
const linesOffsetStart = firstPageIdx * pageSize;
696-
const linesOffsetEnd = lastPageIdx * pageSize + pageSize;
692+
const linesOffsetEndExcl = lastPageIdx * pageSize + pageSize + 1 /* because excl */;
693+
694+
return Result.ok(new OffsetRange(linesOffsetStart, linesOffsetEndExcl));
695+
}
696+
697+
export function createTaggedCurrentFileContentUsingPagedClipping(
698+
currentDocLines: string[],
699+
areaAroundCodeToEdit: string,
700+
areaAroundEditWindowLinesRange: OffsetRange,
701+
computeTokens: (s: string) => number,
702+
pageSize: number,
703+
opts: CurrentFileOptions
704+
): Result<{ taggedCurrentFileContent: string; nLines: number }, 'outOfBudget'> {
705+
706+
const r = clipPreservingRange(
707+
currentDocLines,
708+
areaAroundEditWindowLinesRange,
709+
computeTokens,
710+
pageSize,
711+
opts
712+
);
713+
714+
if (r.isError()) {
715+
return Result.error('outOfBudget');
716+
}
717+
718+
const clippedRange = r.val;
697719

698720
const taggedCurrentFileContent = [
699-
...currentDocLines.slice(linesOffsetStart, areaAroundEditWindowLinesRange.start),
721+
...currentDocLines.slice(clippedRange.start, areaAroundEditWindowLinesRange.start),
700722
areaAroundCodeToEdit,
701-
...currentDocLines.slice(areaAroundEditWindowLinesRange.endExclusive, linesOffsetEnd),
723+
...currentDocLines.slice(areaAroundEditWindowLinesRange.endExclusive, clippedRange.endExclusive),
702724
];
703725

704726
return Result.ok({ taggedCurrentFileContent: taggedCurrentFileContent.join('\n'), nLines: taggedCurrentFileContent.length });

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

Lines changed: 76 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -43,6 +43,7 @@ import { Position } from '../../../util/vs/editor/common/core/position';
4343
import { Range } from '../../../util/vs/editor/common/core/range';
4444
import { LineRange } from '../../../util/vs/editor/common/core/ranges/lineRange';
4545
import { OffsetRange } from '../../../util/vs/editor/common/core/ranges/offsetRange';
46+
import { StringText } from '../../../util/vs/editor/common/core/text/abstractText';
4647
import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation';
4748
import { Position as VscodePosition } from '../../../vscodeTypes';
4849
import { Delayer, DelaySession } from '../../inlineEdits/common/delayer';
@@ -258,43 +259,22 @@ export class XtabProvider implements IStatelessNextEditProvider {
258259
const doesIncludeCursorTag = editWindowLines.some(line => line.includes(PromptTags.CURSOR));
259260
const shouldRemoveCursorTagFromResponse = !doesIncludeCursorTag; // we'd like to remove the tag only if the original edit-window didn't include the tag
260261

261-
const addCursorTagEdit = StringEdit.single(StringReplacement.insert(cursorOffset, PromptTags.CURSOR));
262-
const contentWithCursor = addCursorTagEdit.applyOnText(currentFileContent);
263-
const contentWithCursorLines = contentWithCursor.getLines();
264-
265-
const editWindowWithCursorLines = contentWithCursorLines.slice(editWindowLinesRange.start, editWindowLinesRange.endExclusive);
266-
267-
const areaAroundCodeToEdit = [
268-
PromptTags.AREA_AROUND.start,
269-
...contentWithCursorLines.slice(areaAroundEditWindowLinesRange.start, editWindowLinesRange.start),
270-
PromptTags.EDIT_WINDOW.start,
271-
...editWindowWithCursorLines,
272-
PromptTags.EDIT_WINDOW.end,
273-
...contentWithCursorLines.slice(editWindowLinesRange.endExclusive, areaAroundEditWindowLinesRange.endExclusive),
274-
PromptTags.AREA_AROUND.end
275-
].join('\n');
276-
277-
const areaAroundCodeToEditForCurrentFile = promptOptions.currentFile.includeTags
278-
? areaAroundCodeToEdit
279-
: [
280-
...contentWithCursorLines.slice(areaAroundEditWindowLinesRange.start, editWindowLinesRange.start),
281-
...editWindowLines,
282-
...contentWithCursorLines.slice(editWindowLinesRange.endExclusive, areaAroundEditWindowLinesRange.endExclusive),
283-
].join('\n');
284-
const taggedCurrentFileContentResult = createTaggedCurrentFileContentUsingPagedClipping(
262+
const taggedCurrentFileContentResult = this.constructTaggedFile(
263+
currentFileContent,
285264
currentFileContentLines,
286-
areaAroundCodeToEditForCurrentFile,
265+
cursorOffset,
266+
editWindowLinesRange,
287267
areaAroundEditWindowLinesRange,
268+
promptOptions,
288269
computeTokens,
289-
promptOptions.pagedClipping.pageSize,
290-
promptOptions.currentFile,
270+
{ includeLineNumbers: false }
291271
);
292272

293273
if (taggedCurrentFileContentResult.isError()) {
294274
return Result.error(new NoNextEditReason.PromptTooLarge('currentFile'));
295275
}
296276

297-
const { taggedCurrentFileContent, nLines: nLinesCurrentFile } = taggedCurrentFileContentResult.val;
277+
const { taggedCurrentFileR: { taggedCurrentFileContent, nLines: nLinesCurrentFile }, areaAroundCodeToEdit } = taggedCurrentFileContentResult.val;
298278

299279
telemetryBuilder.setNLinesOfCurrentFileInPrompt(nLinesCurrentFile);
300280

@@ -382,6 +362,74 @@ export class XtabProvider implements IStatelessNextEditProvider {
382362
return Result.ok<void>(undefined);
383363
}
384364

365+
private constructTaggedFile(
366+
currentFileContent: StringText,
367+
currentFileContentLines: string[],
368+
cursorOffset: number,
369+
editWindowLinesRange: OffsetRange,
370+
areaAroundEditWindowLinesRange: OffsetRange,
371+
promptOptions: ModelConfig,
372+
computeTokens: (s: string) => number,
373+
opts: {
374+
includeLineNumbers: boolean;
375+
}
376+
) {
377+
const contentWithCursorAsLinesOriginal = (() => {
378+
const addCursorTagEdit = StringEdit.single(StringReplacement.insert(cursorOffset, PromptTags.CURSOR));
379+
const contentWithCursor = addCursorTagEdit.applyOnText(currentFileContent);
380+
return contentWithCursor.getLines();
381+
})();
382+
383+
const addLineNumbers = (lines: string[]) => lines.map((line, idx) => `${idx + 1}| ${line}`);
384+
385+
const contentWithCursorAsLines = opts.includeLineNumbers
386+
? addLineNumbers(contentWithCursorAsLinesOriginal)
387+
: contentWithCursorAsLinesOriginal;
388+
389+
const editWindowWithCursorAsLines = contentWithCursorAsLines.slice(editWindowLinesRange.start, editWindowLinesRange.endExclusive);
390+
391+
const areaAroundCodeToEdit = [
392+
PromptTags.AREA_AROUND.start,
393+
...contentWithCursorAsLines.slice(areaAroundEditWindowLinesRange.start, editWindowLinesRange.start),
394+
PromptTags.EDIT_WINDOW.start,
395+
...editWindowWithCursorAsLines,
396+
PromptTags.EDIT_WINDOW.end,
397+
...contentWithCursorAsLines.slice(editWindowLinesRange.endExclusive, areaAroundEditWindowLinesRange.endExclusive),
398+
PromptTags.AREA_AROUND.end
399+
].join('\n');
400+
401+
currentFileContentLines = opts.includeLineNumbers
402+
? addLineNumbers(currentFileContentLines)
403+
: currentFileContentLines;
404+
405+
let areaAroundCodeToEditForCurrentFile: string;
406+
if (promptOptions.currentFile.includeTags) {
407+
areaAroundCodeToEditForCurrentFile = areaAroundCodeToEdit;
408+
} else {
409+
const editWindowLines = currentFileContentLines.slice(editWindowLinesRange.start, editWindowLinesRange.endExclusive);
410+
areaAroundCodeToEditForCurrentFile = [
411+
...contentWithCursorAsLines.slice(areaAroundEditWindowLinesRange.start, editWindowLinesRange.start),
412+
...editWindowLines,
413+
...contentWithCursorAsLines.slice(editWindowLinesRange.endExclusive, areaAroundEditWindowLinesRange.endExclusive),
414+
].join('\n');
415+
}
416+
417+
const taggedCurrentFileContentResult = createTaggedCurrentFileContentUsingPagedClipping(
418+
currentFileContentLines,
419+
areaAroundCodeToEditForCurrentFile,
420+
areaAroundEditWindowLinesRange,
421+
computeTokens,
422+
promptOptions.pagedClipping.pageSize,
423+
promptOptions.currentFile,
424+
);
425+
426+
return taggedCurrentFileContentResult.map(taggedCurrentFileR => ({
427+
taggedCurrentFileR,
428+
areaAroundCodeToEdit,
429+
}));
430+
}
431+
432+
385433
private async getLanguageContext(
386434
request: StatelessNextEditRequest,
387435
delaySession: DelaySession,

0 commit comments

Comments
 (0)