diff --git a/src/extension/xtab/common/promptCrafting.ts b/src/extension/xtab/common/promptCrafting.ts index 2848595b0..64a5f43eb 100644 --- a/src/extension/xtab/common/promptCrafting.ts +++ b/src/extension/xtab/common/promptCrafting.ts @@ -16,6 +16,7 @@ import { illegalArgument } from '../../../util/vs/base/common/errors'; import { Schemas } from '../../../util/vs/base/common/network'; import { OffsetRange } from '../../../util/vs/editor/common/core/ranges/offsetRange'; import { StringText } from '../../../util/vs/editor/common/core/text/abstractText'; +import { CurrentDocument } from './xtabCurrentDocument'; export namespace PromptTags { export const CURSOR = "<|cursor|>"; @@ -147,6 +148,9 @@ export const xtab275SystemPrompt = `Predict the next code edit based on user con export class PromptPieces { constructor( + public readonly currentDocument: CurrentDocument, + public readonly editWindowLinesRange: OffsetRange, + public readonly areaAroundEditWindowLinesRange: OffsetRange, public readonly activeDoc: StatelessNextEditDocument, public readonly xtabHistory: readonly IXtabHistoryEntry[], public readonly currentFileContent: string, @@ -172,7 +176,7 @@ export function getUserPrompt(promptPieces: PromptPieces): string { const currentFilePath = toUniquePath(activeDoc.id, activeDoc.workspaceRoot?.path); - const postScript = getPostScript(opts.promptingStrategy, currentFilePath); + const postScript = promptPieces.opts.includePostScript ? getPostScript(opts.promptingStrategy, currentFilePath) : ''; const mainPrompt = `${PromptTags.RECENT_FILES.start} ${recentlyViewedCodeSnippets} diff --git a/src/extension/xtab/node/xtabProvider.ts b/src/extension/xtab/node/xtabProvider.ts index ec2963666..b2acc41c4 100644 --- a/src/extension/xtab/node/xtabProvider.ts +++ b/src/extension/xtab/node/xtabProvider.ts @@ -3,6 +3,7 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ +import { RequestType } from '@vscode/copilot-api'; import { Raw } from '@vscode/prompt-tsx'; import { ChatCompletionContentPartKind } from '@vscode/prompt-tsx/dist/base/output/rawTypes'; import { FetchStreamSource } from '../../../platform/chat/common/chatMLFetcher'; @@ -34,6 +35,7 @@ import { IWorkspaceService } from '../../../platform/workspace/common/workspaceS import { raceFilter } from '../../../util/common/async'; import * as errors from '../../../util/common/errors'; import { Result } from '../../../util/common/result'; +import { TokenizerType } from '../../../util/common/tokenizer'; import { createTracer, ITracer } from '../../../util/common/tracing'; import { AsyncIterableObject, DeferredPromise, raceTimeout, timeout } from '../../../util/vs/base/common/async'; import { CancellationToken } from '../../../util/vs/base/common/cancellation'; @@ -287,6 +289,9 @@ export class XtabProvider implements IStatelessNextEditProvider { } const promptPieces = new PromptPieces( + currentDocument, + editWindowLinesRange, + areaAroundEditWindowLinesRange, activeDocument, request.xtabEditHistory, taggedCurrentFileContent, @@ -332,6 +337,7 @@ export class XtabProvider implements IStatelessNextEditProvider { cursorOriginalLinesOffset, cursorLineOffset, editWindowLinesRange, + promptPieces, prediction, { shouldRemoveCursorTagFromResponse, @@ -351,7 +357,7 @@ export class XtabProvider implements IStatelessNextEditProvider { currentDocument: CurrentDocument, editWindowLinesRange: OffsetRange, areaAroundEditWindowLinesRange: OffsetRange, - promptOptions: ModelConfig, + promptOptions: xtabPromptOptions.PromptOptions, computeTokens: (s: string) => number, opts: { includeLineNumbers: boolean; @@ -533,6 +539,7 @@ export class XtabProvider implements IStatelessNextEditProvider { cursorOriginalLinesOffset: number, cursorLineOffset: number, // cursor offset within the line it's in; 1-based editWindowLineRange: OffsetRange, + promptPieces: PromptPieces, prediction: Prediction | undefined, opts: { promptingStrategy: xtabPromptOptions.PromptingStrategy | undefined; @@ -679,7 +686,7 @@ export class XtabProvider implements IStatelessNextEditProvider { const trimmedLines = firstLine.value.trim(); if (trimmedLines === ResponseTags.NO_CHANGE.start) { - this.pushNoSuggestionsOrRetry(request, editWindow, pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, opts.retryState); + await this.pushNoSuggestionsOrRetry(request, editWindow, promptPieces, pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, opts.retryState); return; } @@ -807,7 +814,7 @@ export class XtabProvider implements IStatelessNextEditProvider { if (hadEdits) { pushEdit(Result.error(new NoNextEditReason.NoSuggestions(request.documentBeforeEdits, editWindow))); } else { - this.pushNoSuggestionsOrRetry(request, editWindow, pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, opts.retryState); + await this.pushNoSuggestionsOrRetry(request, editWindow, promptPieces, pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, opts.retryState); } } catch (err) { @@ -818,9 +825,10 @@ export class XtabProvider implements IStatelessNextEditProvider { })(); } - private pushNoSuggestionsOrRetry( + private async pushNoSuggestionsOrRetry( request: StatelessNextEditRequest, editWindow: OffsetRange, + promptPieces: PromptPieces, pushEdit: PushEdit, delaySession: DelaySession, logContext: InlineEditRequestLogContext, @@ -836,6 +844,17 @@ export class XtabProvider implements IStatelessNextEditProvider { return; } + // FIXME@ulugbekna: think out how it works with retrying logic + if (this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsNextCursorPredictionEnabled, this.expService)) { + // FIXME@ulugbekna: possibly convert from 1-based to 0-based + const nextCursorLine = await this.predictNextCursorPosition(promptPieces); + if (nextCursorLine) { + this.tracer.trace(`Predicted next cursor line: ${nextCursorLine}`); + this.doGetNextEditWithSelection(request, new Range(nextCursorLine, 1, nextCursorLine, 1), pushEdit, delaySession, logContext, cancellationToken, telemetryBuilder, RetryState.NotRetrying); + return; + } + } + pushEdit(Result.error(new NoNextEditReason.NoSuggestions(request.documentBeforeEdits, editWindow))); return; } @@ -977,7 +996,8 @@ export class XtabProvider implements IStatelessNextEditProvider { maxTokens: this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabDiffMaxTokens, this.expService), onlyForDocsInPrompt: this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabDiffOnlyForDocsInPrompt, this.expService), useRelativePaths: this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabDiffUseRelativePaths, this.expService), - } + }, + includePostScript: true, }; const overridingModelConfig = this.configService.getConfig(ConfigKey.Internal.InlineEditsXtabProviderModelConfiguration); @@ -996,6 +1016,120 @@ export class XtabProvider implements IStatelessNextEditProvider { return sourcedModelConfig; } + private async predictNextCursorPosition(promptPieces: PromptPieces) { + + const tracer = this.tracer.sub('predictNextCursorPosition'); + + 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.'; + + const currentFileContentR = this.constructTaggedFile( + promptPieces.currentDocument, + promptPieces.editWindowLinesRange, + promptPieces.areaAroundEditWindowLinesRange, + promptPieces.opts, + XtabProvider.computeTokens, + { includeLineNumbers: true } + ); + + if (currentFileContentR.isError()) { + tracer.trace(`Failed to construct tagged file: ${currentFileContentR.err}`); + return; + } + + const { taggedCurrentFileR: { taggedCurrentFileContent }, areaAroundCodeToEdit } = currentFileContentR.val; + + const newPromptPieces = new PromptPieces( + promptPieces.currentDocument, + promptPieces.editWindowLinesRange, + promptPieces.areaAroundEditWindowLinesRange, + promptPieces.activeDoc, + promptPieces.xtabHistory, + taggedCurrentFileContent, + areaAroundCodeToEdit, + promptPieces.langCtx, + XtabProvider.computeTokens, + { + ...promptPieces.opts, + includePostScript: false, + } + ); + + const userMessage = getUserPrompt(newPromptPieces); + + const messages = constructMessages({ + systemMsg: systemMessage, + userMsg: userMessage + }); + + const modelName = this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsNextCursorPredictionModelName, this.expService); + if (modelName === undefined) { + tracer.trace('Model name for cursor prediction is not defined; skipping prediction'); + return; + } + + const url = this.configService.getConfig(ConfigKey.Internal.InlineEditsNextCursorPredictionUrl); + const secretKey = this.configService.getConfig(ConfigKey.Internal.InlineEditsNextCursorPredictionApiKey); + + const endpoint = this.instaService.createInstance(ChatEndpoint, { + id: modelName, + name: 'nes.nextCursorPosition', + urlOrRequestMetadata: url ? url : { type: RequestType.ProxyChatCompletions }, + model_picker_enabled: false, + is_chat_default: false, + is_chat_fallback: false, + version: '', + capabilities: { + type: 'chat', + family: '', + tokenizer: TokenizerType.CL100K, + limits: undefined, + supports: { + parallel_tool_calls: false, + tool_calls: false, + streaming: true, + vision: false, + prediction: false, + thinking: false + } + }, + }); + + const response = await endpoint.makeChatRequest2( + { + messages, + debugName: 'nes.nextCursorPosition', + finishedCb: undefined, + location: ChatLocation.Other, + requestOptions: secretKey ? { + secretKey, + } : undefined, + }, + CancellationToken.None + ); + + if (response.type !== ChatFetchResponseType.Success) { + return; + } + + try { + const trimmed = response.value.trim(); + const lineNumber = parseInt(trimmed, 10); + if (isNaN(lineNumber) || lineNumber < 0) { + throw new Error(`parsed line number is NaN or negative: ${trimmed}`); + } + if (lineNumber > promptPieces.currentDocument.lines.length) { + this.tracer.trace(`Predicted line number ${lineNumber} is out of bounds, document has ${promptPieces.currentDocument.lines.length} lines`); + return undefined; + } + + return lineNumber; + } catch (err) { + tracer.trace(`Failed to parse predicted line number from response '${response.value}': ${err}`); + return undefined; + } + + } + private determinePromptingStrategy(): xtabPromptOptions.PromptingStrategy | undefined { const isXtabUnifiedModel = this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabUseUnifiedModel, this.expService); const isCodexV21NesUnified = this.configService.getExperimentBasedConfig(ConfigKey.Internal.InlineEditsXtabCodexV21NesUnified, this.expService); diff --git a/src/platform/configuration/common/configurationService.ts b/src/platform/configuration/common/configurationService.ts index 9a9d165bb..a3984d840 100644 --- a/src/platform/configuration/common/configurationService.ts +++ b/src/platform/configuration/common/configurationService.ts @@ -664,6 +664,10 @@ export namespace ConfigKey { export const InlineEditsHideInternalInterface = defineValidatedSetting('chat.advanced.inlineEdits.hideInternalInterface', vBoolean(), false, INTERNAL_RESTRICTED); export const InlineEditsLogCancelledRequests = defineValidatedSetting('chat.advanced.inlineEdits.logCancelledRequests', vBoolean(), false, INTERNAL_RESTRICTED); export const InlineEditsUnification = defineExpSetting('chat.advanced.inlineEdits.unification', false, INTERNAL_RESTRICTED); + export const InlineEditsNextCursorPredictionEnabled = defineExpSetting('chat.advanced.inlineEdits.nextCursorPrediction.enabled', false, INTERNAL_RESTRICTED); + export const InlineEditsNextCursorPredictionModelName = defineExpSetting('chat.advanced.inlineEdits.nextCursorPrediction.modelName', undefined, INTERNAL_RESTRICTED); + export const InlineEditsNextCursorPredictionUrl = defineValidatedSetting('chat.advanced.inlineEdits.nextCursorPrediction.url', vString(), undefined, INTERNAL_RESTRICTED); + export const InlineEditsNextCursorPredictionApiKey = defineValidatedSetting('chat.advanced.inlineEdits.nextCursorPrediction.apiKey', vString(), undefined, INTERNAL_RESTRICTED); export const InlineEditsXtabProviderUrl = defineValidatedSetting('chat.advanced.inlineEdits.xtabProvider.url', vString(), undefined, INTERNAL_RESTRICTED); export const InlineEditsXtabProviderApiKey = defineValidatedSetting('chat.advanced.inlineEdits.xtabProvider.apiKey', vString(), undefined, INTERNAL_RESTRICTED); export const InlineEditsXtabProviderModelConfiguration = defineValidatedSetting('chat.advanced.inlineEdits.xtabProvider.modelConfiguration', xtabPromptOptions.MODEL_CONFIGURATION_VALIDATOR, undefined, INTERNAL_RESTRICTED); diff --git a/src/platform/inlineEdits/common/dataTypes/xtabPromptOptions.ts b/src/platform/inlineEdits/common/dataTypes/xtabPromptOptions.ts index ef46d4c3f..afce733aa 100644 --- a/src/platform/inlineEdits/common/dataTypes/xtabPromptOptions.ts +++ b/src/platform/inlineEdits/common/dataTypes/xtabPromptOptions.ts @@ -40,6 +40,7 @@ export type PromptOptions = { readonly recentlyViewedDocuments: RecentlyViewedDocumentsOptions; readonly languageContext: LanguageContextOptions; readonly diffHistory: DiffHistoryOptions; + readonly includePostScript: boolean; } /** @@ -81,6 +82,7 @@ export const DEFAULT_OPTIONS: PromptOptions = { onlyForDocsInPrompt: false, useRelativePaths: false, }, + includePostScript: true, }; // TODO: consider a better per language setting/experiment approach