Added context from the rest of the code + updated prompt to only replace selected

vscode_extension/wolverine/src/extension.ts

@@ -5,15 +5,16 @@ import {
 	TextEditor,
 	Range,
 	commands,
-	window
+	window,
+	TextDocument,
 } from 'vscode';
 import axios from 'axios';
 
-let openaikey = '';
 
 // See https://html.spec.whatwg.org/multipage/server-sent-events.html#server-sent-events
 // Responsiblity of caller to register event listener.
 const streamCompletion = async (prompt: string, onDataFunction: (chunk: any) => void): Promise<void> => {
+	const openaikey = await workspace.getConfiguration().get('wolverine.UNSAFE.OpenaiApiKeySetting') || '';
 	const messages: any[] = [
 		{ role: 'user', content: prompt }
 	];
@@ -22,7 +23,7 @@ const streamCompletion = async (prompt: string, onDataFunction: (chunk: any) =>
 		['Authorization']: `Bearer ${openaikey}`,
 	};
 
-	const configurationModel =  await workspace.getConfiguration().get('wolverine.model');
+	const configurationModel = await workspace.getConfiguration().get('wolverine.model');
 	const data = {
 		'model': configurationModel || 'gpt-3.5-turbo',
 		'messages': messages,
@@ -43,7 +44,6 @@ const streamCompletion = async (prompt: string, onDataFunction: (chunk: any) =>
 		});
 	});
 };
-
 const countCharacters = (text: string) => text.replace(/\n/g, '').length;
 const countNewLines = (text: string) => text.match(/\n/g)?.length || 0;
 const getNewCursorLocation = (textStream: string, currentLine: number, currentCharacter: number): { newCharacterLocation: number, newLineLocation: number } => {
@@ -52,13 +52,11 @@ const getNewCursorLocation = (textStream: string, currentLine: number, currentCh
 	const newLineLocation = numberOfNewLines + currentLine;
 	return { newCharacterLocation, newLineLocation };
 };
-
 const deleteRange = async (activeEditor: TextEditor, range: Range) => {
 	await activeEditor.edit(editBuilder => {
 		editBuilder.delete(range);
 	});
 };
-
 // Yacine threw most of the complexity into here.
 // Holds a buffer in a javascript array, registers a event listener on a server-sent events function, builds the buffer
 // Takes a position, and flushes the buffer on a preconfigured cron into the provided cursor position.
@@ -87,20 +85,26 @@ const useBufferToUpdateTextContentsWhileStreamingOpenAIResponse = async (activeE
 			currentLine = newLineLocation;
 		}
 		// TODO I should make this buffer flush configurable.
-		await sleep(30); 
+		await sleep(30);
 	}
 };
-
-const constructPrompt = async (text: string): Promise<string> => {
+const constructPrompt = async (text: string, filepath: string): Promise<string> => {
 	const defaultPrompt = `
+ADDITIONAL CONTEXT:
+filepath: ${filepath}
 INSTRUCTIONS: 
-The text under 'CODE' has come straight from my text editor. 
-When you respond, please only output valid code, and no raw english! If its english, put it under comments.
-We want your entire output to be valid codee.
-The code sent to you might have some instructions under comments. Follow those instructions.
+- The text under 'CODE' has come straight from my text editor. 
+- Your entire output MUST be valid code.
+- You may communicate back, but they MUST be in comments
+- The code sent to you might have some instructions under comments. Follow the instructions. Only attend to instructions contained between [[SELECTED]] & [[/SELECTED]]
+- The code between [[SELECTED]] [[/SELECTED]] is code that I need you to replace. So when you start responding, only respond with code that could replace it
+- MAKE SURE YOU ONLY WRITE CODE BETWEEN [[SELECTED]] AND [[/SELECTED]]. The rest of it doesnt need to be replace
+- The rest of the code is provided as context. Sometimes, I'll collapse my code for you, and hide the details under '...'. Follow the style of the rest of the provided code.
+- DO NOT WRITE THE ENTIRE CODE FILE. ONLY REPLACE WHAT IS NECESSARY IN SELECTED
+- Prefer functional programming principles.
 CODE:
 ${text}
-NEW CODE:
+NEW CODE TO REPLACE WHAT IS BETWEEN [[SELECTED]] and [/SELECTED]:
 `;
 	const configuredPrompt = await workspace.getConfiguration().get('wolverine.prompt');
 	if (configuredPrompt) {
@@ -108,12 +112,26 @@ NEW CODE:
 	}
 	return defaultPrompt;
 };
+// Gets the visible text, and adds ... between unseen text
+async function getVisibleText(document: TextDocument, visibleRanges: readonly Range[]): Promise<string> {
+	let visibleText = '';
+	let lastVisibleRangeEnd: Position | null = null;
 
+	visibleRanges.forEach((range, index) => {
+		if (lastVisibleRangeEnd && document.offsetAt(range.start) - document.offsetAt(lastVisibleRangeEnd) > 1) {
+			visibleText += '...';
+		}
+		visibleText += document.getText(range);
+		lastVisibleRangeEnd = range.end;
+	});
+
+	return visibleText;
+}
+
+// Main function to activate the extension
 export async function activate(context: ExtensionContext) {
 
 	let disposable = commands.registerCommand('wolverine.directedHeal', async () => {
-		// Forgive me father..
-		openaikey = await workspace.getConfiguration().get('wolverine.UNSAFE.OpenaiApiKeySetting') || '';
 		const activeEditor = window.activeTextEditor;
 		if (!activeEditor) {
 			window.showErrorMessage('No text editor is currently active.');
@@ -121,18 +139,33 @@ export async function activate(context: ExtensionContext) {
 		}
 
 		const selection = activeEditor.selection;
-		let range: Range;
-		if (!selection.isEmpty) {
-			range = new Range(selection.start, selection.end);
-		} else {
-			range = activeEditor.document.validateRange(new Range(0, 0, Number.MAX_VALUE, Number.MAX_VALUE));
+		if (selection.isEmpty) {
+			window.showErrorMessage('No text selected.');
+			return;
 		}
-		const text = activeEditor.document.getText(range);
-		const prompt = await constructPrompt(text);
+
+		const document = activeEditor.document;
+		const visibleText = await getVisibleText(document, activeEditor.visibleRanges);
+
+		// Get the selected text and replace it with marked selected text in the unfolded text
+		// This should then be used in the prompt; to instruct the LLM what to do
+		const selectedText = document.getText(selection);
+		const contextText = visibleText.replace(selectedText, '[[SELECTED]]' + selectedText + '[[/SELECTED]]');
+		const filePath = workspace.asRelativePath(document.uri);
+		const prompt = await constructPrompt(contextText, filePath);
+
+		// Determine the range to delete and replace
+		const range: Range = new Range(selection.start, selection.end);
 		await deleteRange(activeEditor, range);
+
+		// Update the text using OpenAI response
 		await useBufferToUpdateTextContentsWhileStreamingOpenAIResponse(activeEditor, range.start, prompt);
+
+		// Save the document
 		await activeEditor.document.save();
 	});
+
+	// Add the disposable to the context subscriptions
 	context.subscriptions.push(disposable);
 }