Skip to content

Commit 44e806b

Browse files
ulugbeknaulugbekna
authored
nes: feat: support next cursor line prediction (#1238)
* nes: feat: support next cursor line prediction not well tested due to 404 from the model * nes: fix: validate predicted line number isn't out of max # of lines in document * nes: fix: use custom url & api key only if set; otherwise, use proxy * nes: next cursor line: fix: exclude postscript
1 parent 0b1a979 commit 44e806b

File tree

4 files changed

+150
-6
lines changed

4 files changed

+150
-6
lines changed

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import { illegalArgument } from '../../../util/vs/base/common/errors';
1616
import { Schemas } from '../../../util/vs/base/common/network';
1717
import { OffsetRange } from '../../../util/vs/editor/common/core/ranges/offsetRange';
1818
import { StringText } from '../../../util/vs/editor/common/core/text/abstractText';
19+
import { CurrentDocument } from './xtabCurrentDocument';
1920

2021
export namespace PromptTags {
2122
export const CURSOR = "<|cursor|>";
@@ -147,6 +148,9 @@ export const xtab275SystemPrompt = `Predict the next code edit based on user con
147148

148149
export class PromptPieces {
149150
constructor(
151+
public readonly currentDocument: CurrentDocument,
152+
public readonly editWindowLinesRange: OffsetRange,
153+
public readonly areaAroundEditWindowLinesRange: OffsetRange,
150154
public readonly activeDoc: StatelessNextEditDocument,
151155
public readonly xtabHistory: readonly IXtabHistoryEntry[],
152156
public readonly currentFileContent: string,
@@ -172,7 +176,7 @@ export function getUserPrompt(promptPieces: PromptPieces): string {
172176

173177
const currentFilePath = toUniquePath(activeDoc.id, activeDoc.workspaceRoot?.path);
174178

175-
const postScript = getPostScript(opts.promptingStrategy, currentFilePath);
179+
const postScript = promptPieces.opts.includePostScript ? getPostScript(opts.promptingStrategy, currentFilePath) : '';
176180

177181
const mainPrompt = `${PromptTags.RECENT_FILES.start}
178182
${recentlyViewedCodeSnippets}

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

Lines changed: 139 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
* Licensed under the MIT License. See License.txt in the project root for license information.
44
*--------------------------------------------------------------------------------------------*/
55

6+
import { RequestType } from '@vscode/copilot-api';
67
import { Raw } from '@vscode/prompt-tsx';
78
import { ChatCompletionContentPartKind } from '@vscode/prompt-tsx/dist/base/output/rawTypes';
89
import { FetchStreamSource } from '../../../platform/chat/common/chatMLFetcher';
@@ -34,6 +35,7 @@ import { IWorkspaceService } from '../../../platform/workspace/common/workspaceS
3435
import { raceFilter } from '../../../util/common/async';
3536
import * as errors from '../../../util/common/errors';
3637
import { Result } from '../../../util/common/result';
38+
import { TokenizerType } from '../../../util/common/tokenizer';
3739
import { createTracer, ITracer } from '../../../util/common/tracing';
3840
import { AsyncIterableObject, DeferredPromise, raceTimeout, timeout } from '../../../util/vs/base/common/async';
3941
import { CancellationToken } from '../../../util/vs/base/common/cancellation';
@@ -287,6 +289,9 @@ export class XtabProvider implements IStatelessNextEditProvider {
287289
}
288290

289291
const promptPieces = new PromptPieces(
292+
currentDocument,
293+
editWindowLinesRange,
294+
areaAroundEditWindowLinesRange,
290295
activeDocument,
291296
request.xtabEditHistory,
292297
taggedCurrentFileContent,
@@ -332,6 +337,7 @@ export class XtabProvider implements IStatelessNextEditProvider {
332337
cursorOriginalLinesOffset,
333338
cursorLineOffset,
334339
editWindowLinesRange,
340+
promptPieces,
335341
prediction,
336342
{
337343
shouldRemoveCursorTagFromResponse,
@@ -351,7 +357,7 @@ export class XtabProvider implements IStatelessNextEditProvider {
351357
currentDocument: CurrentDocument,
352358
editWindowLinesRange: OffsetRange,
353359
areaAroundEditWindowLinesRange: OffsetRange,
354-
promptOptions: ModelConfig,
360+
promptOptions: xtabPromptOptions.PromptOptions,
355361
computeTokens: (s: string) => number,
356362
opts: {
357363
includeLineNumbers: boolean;
@@ -533,6 +539,7 @@ export class XtabProvider implements IStatelessNextEditProvider {
533539
cursorOriginalLinesOffset: number,
534540
cursorLineOffset: number, // cursor offset within the line it's in; 1-based
535541
editWindowLineRange: OffsetRange,
542+
promptPieces: PromptPieces,
536543
prediction: Prediction | undefined,
537544
opts: {
538545
promptingStrategy: xtabPromptOptions.PromptingStrategy | undefined;
@@ -682,7 +689,7 @@ export class XtabProvider implements IStatelessNextEditProvider {
682689
const trimmedLines = firstLine.value.trim();
683690

684691
if (trimmedLines === ResponseTags.NO_CHANGE.start) {
685-
this.pushNoSuggestionsOrRetry(request, editWindow, pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, opts.retryState);
692+
await this.pushNoSuggestionsOrRetry(request, editWindow, promptPieces, pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, opts.retryState);
686693
return;
687694
}
688695

@@ -810,7 +817,7 @@ export class XtabProvider implements IStatelessNextEditProvider {
810817
if (hadEdits) {
811818
pushEdit(Result.error(new NoNextEditReason.NoSuggestions(request.documentBeforeEdits, editWindow)));
812819
} else {
813-
this.pushNoSuggestionsOrRetry(request, editWindow, pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, opts.retryState);
820+
await this.pushNoSuggestionsOrRetry(request, editWindow, promptPieces, pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, opts.retryState);
814821
}
815822

816823
} catch (err) {
@@ -821,9 +828,10 @@ export class XtabProvider implements IStatelessNextEditProvider {
821828
})();
822829
}
823830

824-
private pushNoSuggestionsOrRetry(
831+
private async pushNoSuggestionsOrRetry(
825832
request: StatelessNextEditRequest,
826833
editWindow: OffsetRange,
834+
promptPieces: PromptPieces,
827835
pushEdit: PushEdit,
828836
delaySession: DelaySession,
829837
logContext: InlineEditRequestLogContext,
@@ -839,6 +847,17 @@ export class XtabProvider implements IStatelessNextEditProvider {
839847
return;
840848
}
841849

850+
// FIXME@ulugbekna: think out how it works with retrying logic
851+
if (this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsNextCursorPredictionEnabled, this.expService)) {
852+
// FIXME@ulugbekna: possibly convert from 1-based to 0-based
853+
const nextCursorLine = await this.predictNextCursorPosition(promptPieces);
854+
if (nextCursorLine) {
855+
this.tracer.trace(`Predicted next cursor line: ${nextCursorLine}`);
856+
this.doGetNextEditWithSelection(request, new Range(nextCursorLine, 1, nextCursorLine, 1), pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, RetryState.NotRetrying);
857+
return;
858+
}
859+
}
860+
842861
pushEdit(Result.error(new NoNextEditReason.NoSuggestions(request.documentBeforeEdits, editWindow)));
843862
return;
844863
}
@@ -980,7 +999,8 @@ export class XtabProvider implements IStatelessNextEditProvider {
980999
maxTokens: this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabDiffMaxTokens, this.expService),
9811000
onlyForDocsInPrompt: this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabDiffOnlyForDocsInPrompt, this.expService),
9821001
useRelativePaths: this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabDiffUseRelativePaths, this.expService),
983-
}
1002+
},
1003+
includePostScript: true,
9841004
};
9851005

9861006
const overridingModelConfig = this.configService.getConfig(ConfigKey.Internal.InlineEditsXtabProviderModelConfiguration);
@@ -999,6 +1019,120 @@ export class XtabProvider implements IStatelessNextEditProvider {
9991019
return sourcedModelConfig;
10001020
}
10011021

1022+
private async predictNextCursorPosition(promptPieces: PromptPieces) {
1023+
1024+
const tracer = this.tracer.sub('predictNextCursorPosition');
1025+
1026+
const systemMessage = 'Your task is to predict the next line number in the current file where the developer is most likely to make their next edit, using the provided context.';
1027+
1028+
const currentFileContentR = this.constructTaggedFile(
1029+
promptPieces.currentDocument,
1030+
promptPieces.editWindowLinesRange,
1031+
promptPieces.areaAroundEditWindowLinesRange,
1032+
promptPieces.opts,
1033+
XtabProvider.computeTokens,
1034+
{ includeLineNumbers: true }
1035+
);
1036+
1037+
if (currentFileContentR.isError()) {
1038+
tracer.trace(`Failed to construct tagged file: ${currentFileContentR.err}`);
1039+
return;
1040+
}
1041+
1042+
const { taggedCurrentFileR: { taggedCurrentFileContent }, areaAroundCodeToEdit } = currentFileContentR.val;
1043+
1044+
const newPromptPieces = new PromptPieces(
1045+
promptPieces.currentDocument,
1046+
promptPieces.editWindowLinesRange,
1047+
promptPieces.areaAroundEditWindowLinesRange,
1048+
promptPieces.activeDoc,
1049+
promptPieces.xtabHistory,
1050+
taggedCurrentFileContent,
1051+
areaAroundCodeToEdit,
1052+
promptPieces.langCtx,
1053+
XtabProvider.computeTokens,
1054+
{
1055+
...promptPieces.opts,
1056+
includePostScript: false,
1057+
}
1058+
);
1059+
1060+
const userMessage = getUserPrompt(newPromptPieces);
1061+
1062+
const messages = constructMessages({
1063+
systemMsg: systemMessage,
1064+
userMsg: userMessage
1065+
});
1066+
1067+
const modelName = this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsNextCursorPredictionModelName, this.expService);
1068+
if (modelName === undefined) {
1069+
tracer.trace('Model name for cursor prediction is not defined; skipping prediction');
1070+
return;
1071+
}
1072+
1073+
const url = this.configService.getConfig(ConfigKey.Internal.InlineEditsNextCursorPredictionUrl);
1074+
const secretKey = this.configService.getConfig(ConfigKey.Internal.InlineEditsNextCursorPredictionApiKey);
1075+
1076+
const endpoint = this.instaService.createInstance(ChatEndpoint, {
1077+
id: modelName,
1078+
name: 'nes.nextCursorPosition',
1079+
urlOrRequestMetadata: url ? url : { type: RequestType.ProxyChatCompletions },
1080+
model_picker_enabled: false,
1081+
is_chat_default: false,
1082+
is_chat_fallback: false,
1083+
version: '',
1084+
capabilities: {
1085+
type: 'chat',
1086+
family: '',
1087+
tokenizer: TokenizerType.CL100K,
1088+
limits: undefined,
1089+
supports: {
1090+
parallel_tool_calls: false,
1091+
tool_calls: false,
1092+
streaming: true,
1093+
vision: false,
1094+
prediction: false,
1095+
thinking: false
1096+
}
1097+
},
1098+
});
1099+
1100+
const response = await endpoint.makeChatRequest2(
1101+
{
1102+
messages,
1103+
debugName: 'nes.nextCursorPosition',
1104+
finishedCb: undefined,
1105+
location: ChatLocation.Other,
1106+
requestOptions: secretKey ? {
1107+
secretKey,
1108+
} : undefined,
1109+
},
1110+
CancellationToken.None
1111+
);
1112+
1113+
if (response.type !== ChatFetchResponseType.Success) {
1114+
return;
1115+
}
1116+
1117+
try {
1118+
const trimmed = response.value.trim();
1119+
const lineNumber = parseInt(trimmed, 10);
1120+
if (isNaN(lineNumber) || lineNumber < 0) {
1121+
throw new Error(`parsed line number is NaN or negative: ${trimmed}`);
1122+
}
1123+
if (lineNumber > promptPieces.currentDocument.lines.length) {
1124+
this.tracer.trace(`Predicted line number ${lineNumber} is out of bounds, document has ${promptPieces.currentDocument.lines.length} lines`);
1125+
return undefined;
1126+
}
1127+
1128+
return lineNumber;
1129+
} catch (err) {
1130+
tracer.trace(`Failed to parse predicted line number from response '${response.value}': ${err}`);
1131+
return undefined;
1132+
}
1133+
1134+
}
1135+
10021136
private determinePromptingStrategy(): xtabPromptOptions.PromptingStrategy | undefined {
10031137
const isXtabUnifiedModel = this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabUseUnifiedModel, this.expService);
10041138
const isCodexV21NesUnified = this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabCodexV21NesUnified, this.expService);

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -665,6 +665,10 @@ export namespace ConfigKey {
665665
export const InlineEditsHideInternalInterface = defineValidatedSetting<boolean>('chat.advanced.inlineEdits.hideInternalInterface', vBoolean(), false, INTERNAL_RESTRICTED);
666666
export const InlineEditsLogCancelledRequests = defineValidatedSetting<boolean>('chat.advanced.inlineEdits.logCancelledRequests', vBoolean(), false, INTERNAL_RESTRICTED);
667667
export const InlineEditsUnification = defineExpSetting<boolean>('chat.advanced.inlineEdits.unification', false, INTERNAL_RESTRICTED);
668+
export const InlineEditsNextCursorPredictionEnabled = defineExpSetting<boolean>('chat.advanced.inlineEdits.nextCursorPrediction.enabled', false, INTERNAL_RESTRICTED);
669+
export const InlineEditsNextCursorPredictionModelName = defineExpSetting<string | undefined>('chat.advanced.inlineEdits.nextCursorPrediction.modelName', undefined, INTERNAL_RESTRICTED);
670+
export const InlineEditsNextCursorPredictionUrl = defineValidatedSetting<string | undefined>('chat.advanced.inlineEdits.nextCursorPrediction.url', vString(), undefined, INTERNAL_RESTRICTED);
671+
export const InlineEditsNextCursorPredictionApiKey = defineValidatedSetting<string | undefined>('chat.advanced.inlineEdits.nextCursorPrediction.apiKey', vString(), undefined, INTERNAL_RESTRICTED);
668672
export const InlineEditsXtabProviderUrl = defineValidatedSetting<string | undefined>('chat.advanced.inlineEdits.xtabProvider.url', vString(), undefined, INTERNAL_RESTRICTED);
669673
export const InlineEditsXtabProviderApiKey = defineValidatedSetting<string | undefined>('chat.advanced.inlineEdits.xtabProvider.apiKey', vString(), undefined, INTERNAL_RESTRICTED);
670674
export const InlineEditsXtabProviderModelConfiguration = defineValidatedSetting<xtabPromptOptions.ModelConfiguration | undefined>('chat.advanced.inlineEdits.xtabProvider.modelConfiguration', xtabPromptOptions.MODEL_CONFIGURATION_VALIDATOR, undefined, INTERNAL_RESTRICTED);

‎src/platform/inlineEdits/common/dataTypes/xtabPromptOptions.ts‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ export type PromptOptions = {
4040
readonly recentlyViewedDocuments: RecentlyViewedDocumentsOptions;
4141
readonly languageContext: LanguageContextOptions;
4242
readonly diffHistory: DiffHistoryOptions;
43+
readonly includePostScript: boolean;
4344
}
4445

4546
/**
@@ -81,6 +82,7 @@ export const DEFAULT_OPTIONS: PromptOptions = {
8182
onlyForDocsInPrompt: false,
8283
useRelativePaths: false,
8384
},
85+
includePostScript: true,
8486
};
8587

8688
// TODO: consider a better per language setting/experiment approach

0 commit comments

Comments
 (0)