nes: refactor prompt construction for extensibility (#1241)

* nes: refactor: get config from within computeEditWindowLinesRange

* nes: refactor: tagging

* nes: refactor: extract tagging current file into its own method

* nes: refactor: keep active document props in a single object

* nes: refactor: extract cursorLineOffset

* nes: refactor: more refactoring

* nes: extract language context computation

side effect that we now always await not when language context is enabled

* nes: refactor: move CurrentDocument to common/

* nes: refactor: abstract constructing messages

* nes: fix off by one

Author: ulugbekna

src/extension/xtab/common/promptCrafting.ts

@@ -20,35 +20,29 @@ import { StringText } from '../../../util/vs/editor/common/core/text/abstractTex
 export namespace PromptTags {
 	export const CURSOR = "<|cursor|>";
 
-	export const EDIT_WINDOW = {
-		start: "<|code_to_edit|>",
-		end: "<|/code_to_edit|>"
+	type Tag = {
+		start: string;
+		end: string;
 	};
 
-	export const AREA_AROUND = {
-		start: "<|area_around_code_to_edit|>",
-		end: "<|/area_around_code_to_edit|>"
-	};
+	function createTag(key: string): Tag {
+		return {
+			start: `<|${key}|>`,
+			end: `<|/${key}|>`
+		};
+	}
 
-	export const CURRENT_FILE = {
-		start: "<|current_file_content|>",
-		end: "<|/current_file_content|>"
-	};
+	export const EDIT_WINDOW = createTag("code_to_edit");
 
-	export const EDIT_HISTORY = {
-		start: "<|edit_diff_history|>",
-		end: "<|/edit_diff_history|>"
-	};
+	export const AREA_AROUND = createTag("area_around_code_to_edit");
 
-	export const RECENT_FILES = {
-		start: "<|recently_viewed_code_snippets|>",
-		end: "<|/recently_viewed_code_snippets|>"
-	};
+	export const CURRENT_FILE = createTag("current_file_content");
 
-	export const RECENT_FILE = {
-		start: "<|recently_viewed_code_snippet|>",
-		end: "<|/recently_viewed_code_snippet|>"
-	};
+	export const EDIT_HISTORY = createTag("edit_diff_history");
+
+	export const RECENT_FILES = createTag("recently_viewed_code_snippets");
+
+	export const RECENT_FILE = createTag("recently_viewed_code_snippet");
 }
 
 export const systemPromptTemplate = `Your role as an AI assistant is to help developers complete their code tasks by assisting in editing specific sections of code marked by the ${PromptTags.EDIT_WINDOW.start} and ${PromptTags.EDIT_WINDOW.end} tags, while adhering to Microsoft's content policies and avoiding the creation of content that violates copyrights.
@@ -671,40 +665,62 @@ function expandRangeToPageRange(
 	return { firstPageIdx, lastPageIdx, budgetLeft: tokenBudget };
 }
 
-/**
- * @remark exported for testing
- */
-export function createTaggedCurrentFileContentUsingPagedClipping(
-	currentDocLines: string[],
-	areaAroundCodeToEdit: string,
-	areaAroundEditWindowLinesRange: OffsetRange,
+export function clipPreservingRange(
+	docLines: string[],
+	rangeToPreserve: OffsetRange,
 	computeTokens: (s: string) => number,
 	pageSize: number,
-	opts: CurrentFileOptions
-): Result<{ taggedCurrentFileContent: string; nLines: number }, 'outOfBudget'> {
+	opts: CurrentFileOptions,
+): Result<OffsetRange, 'outOfBudget'> {
 
-	// subtract budget consumed by areaAroundCodeToEdit
-	const availableTokenBudget = opts.maxTokens - countTokensForLines(areaAroundCodeToEdit.split(/\r?\n/), computeTokens);
+	// subtract budget consumed by rangeToPreserve
+	const availableTokenBudget = opts.maxTokens - countTokensForLines(docLines.slice(rangeToPreserve.start, rangeToPreserve.endExclusive), computeTokens);
 	if (availableTokenBudget < 0) {
 		return Result.error('outOfBudget');
 	}
 
 	const { firstPageIdx, lastPageIdx } = expandRangeToPageRange(
-		currentDocLines,
-		areaAroundEditWindowLinesRange,
+		docLines,
+		rangeToPreserve,
 		pageSize,
 		availableTokenBudget,
 		computeTokens,
 		opts.prioritizeAboveCursor,
 	);
 
 	const linesOffsetStart = firstPageIdx * pageSize;
-	const linesOffsetEnd = lastPageIdx * pageSize + pageSize;
+	const linesOffsetEndExcl = lastPageIdx * pageSize + pageSize;
+
+	return Result.ok(new OffsetRange(linesOffsetStart, linesOffsetEndExcl));
+}
+
+export function createTaggedCurrentFileContentUsingPagedClipping(
+	currentDocLines: string[],
+	areaAroundCodeToEdit: string,
+	areaAroundEditWindowLinesRange: OffsetRange,
+	computeTokens: (s: string) => number,
+	pageSize: number,
+	opts: CurrentFileOptions
+): Result<{ taggedCurrentFileContent: string; nLines: number }, 'outOfBudget'> {
+
+	const r = clipPreservingRange(
+		currentDocLines,
+		areaAroundEditWindowLinesRange,
+		computeTokens,
+		pageSize,
+		opts
+	);
+
+	if (r.isError()) {
+		return Result.error('outOfBudget');
+	}
+
+	const clippedRange = r.val;
 
 	const taggedCurrentFileContent = [
-		...currentDocLines.slice(linesOffsetStart, areaAroundEditWindowLinesRange.start),
+		...currentDocLines.slice(clippedRange.start, areaAroundEditWindowLinesRange.start),
 		areaAroundCodeToEdit,
-		...currentDocLines.slice(areaAroundEditWindowLinesRange.endExclusive, linesOffsetEnd),
+		...currentDocLines.slice(areaAroundEditWindowLinesRange.endExclusive, clippedRange.endExclusive),
 	];
 
 	return Result.ok({ taggedCurrentFileContent: taggedCurrentFileContent.join('\n'), nLines: taggedCurrentFileContent.length });

src/extension/xtab/common/xtabCurrentDocument.ts

@@ -0,0 +1,31 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { Position } from '../../../util/vs/editor/common/core/position';
+import { StringText } from '../../../util/vs/editor/common/core/text/abstractText';
+import { PositionOffsetTransformer } from '../../../util/vs/editor/common/core/text/positionToOffsetImpl';
+
+export class CurrentDocument {
+
+	public readonly lines: string[];
+	public readonly cursorOffset: number;
+
+	public readonly transformer: PositionOffsetTransformer;
+
+	/**
+	 * The 0-based line number of the cursor.
+	 */
+	public readonly cursorLineOffset: number;
+
+	constructor(
+		public readonly content: StringText,
+		public readonly cursorPosition: Position,
+	) {
+		this.lines = content.getLines();
+		this.transformer = content.getTransformer();
+		this.cursorOffset = this.transformer.getOffset(cursorPosition);
+		this.cursorLineOffset = this.cursorPosition.lineNumber - 1;
+	}
+}