Add ability to exit out, change the prompt to have instructions, some code cleaning, added readme. Enjoy!

vscode_extension/wolverine/README.md

@@ -1,10 +1,44 @@
 Hey chatGPT! This is some markdown, could you rewrite this to be a bit friendlier?
 
-# wolverine README
+
+
+# Wolverine: An AI-powered Code Healing Extension for VS Code
 It's 2020 + 3. Let's stop tabbing away from our text editor, shall we?
 
-#set up workspace settings
+Wolverine is a Visual Studio Code extension designed to assist developers in the code modification process. It utilizes OpenAI's GPT-3.5-turbo model to generate code suggestions and edits in real-time. By directly integrating with your text editor, Wolverine can improve your efficiency and help resolve coding issues.
 
-### Features
-- To auto-correct code, highlight the code and press ctrl+k ctrl+h, which will prompt GPT-3.5 and stream the fixed code back in.
-- If nothing is highlighted, the extension will fix the whole file.
+
+## Usage instructions 
+Right now, this isn't published, because there isn't a good mechanism for having the openAI key through a keychain. So to use; You're going to have to run a build.
+Furthermore; the code is in a bit of a sorry state. This was a really quick hack - which is planned to be improved upon.
+- `Pull the repository`
+- `cd vscode_extension/wolverine`
+- `npm install`
+- `npm run package && npm run vsce`. This will generate a vsix file
+- ctrl shift p => install extensions  
+- click the three little dots on top right of the extension window
+- choose the vsix file you just generated
+
+## Development instructions
+Do you want to mess with this and change things on your own?
+- Start a vscode instance in this current directory. Not the main repo; but `./vscode_extension/wolverine`
+- ctrl+shift+p => debug extension (f11 hotkey default) 
+- this will launch a new code editor; with a debugger attached to your main window
+
+## Key Features
+### Directed Code Healing
+Wolverine introduces a 'directedHeal' command. This command lets you select a block of code that you'd like to replace or improve. Simply select the code, execute the 'directedHeal' command, and the extension will replace the selected code with improved code generated by OpenAI's GPT-3.5-turbo model.
+
+### Contextual Assistance
+Wolverine understands the context of your code. It utilizes the surrounding code (visible in your text editor) to generate meaningful and coherent code replacements that align with your existing codebase.
+
+### Real-Time Streaming of AI Responses:
+The extension streams the AI-generated code in real-time, providing a seamless experience. It buffers the AI responses and updates the text editor at regular intervals, making the process efficient and smooth.
+
+### How to Use
+Select the piece of code you want to replace in your text editor.
+Run the 'wolverine.directedHeal' command.
+Watch as Wolverine replaces your selected code with AI-generated code.
+
+Note
+Please ensure that you have an active OpenAI key to use this extension, as it leverages the OpenAI API for generating the code.

vscode_extension/wolverine/package.json

@@ -36,7 +36,15 @@
     "commands": [
       {
         "command": "wolverine.directedHeal",
-        "title": "Directed Heal"
+        "title": "Directed Heal",
+        "key": "ctrl+k ctrl+h",
+        "when": "editorTextFocus"
+      },
+      {
+        "command": "wolverine.cancelOperation",
+        "title": "Cancel Current Buffer Updates",
+        "key": "escape",
+        "when": "wolverine.operationRunning"
       }
     ]
   },

vscode_extension/wolverine/src/extension.ts

@@ -7,13 +7,45 @@ import {
 	commands,
 	window,
 	TextDocument,
+	CancellationTokenSource,
 } from 'vscode';
 import axios from 'axios';
 
+// Flush interval is 30 milliseconds, because that's what the author found works well on the author's system.
+// Ryzen 5900x + arch linux (btw)
+const FLUSH_INTERVAL_MS = 30;
 
+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 } => {
+	const numberOfNewLines = countNewLines(textStream);
+	const newCharacterLocation = numberOfNewLines === 0 ? countCharacters(textStream) + currentCharacter : 0;
+	const newLineLocation = numberOfNewLines + currentLine;
+	return { newCharacterLocation, newLineLocation };
+};
+const deleteRange = async (activeEditor: TextEditor, range: Range) => {
+	await activeEditor.edit(editBuilder => {
+		editBuilder.delete(range);
+	});
+};
+// Gets the currently visible text, and if anything is folded, add ... 
+const getVisibleText = async (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;
+};
 // 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 streamCompletion = async (prompt: string, onDataFunction: (chunk: any) => void, cancelFunction: (source: any, reject: () => void) => void): Promise<void> => {
 	const openaikey = await workspace.getConfiguration().get('wolverine.UNSAFE.OpenaiApiKeySetting') || '';
 	const messages: any[] = [
 		{ role: 'user', content: prompt }
@@ -30,41 +62,33 @@ const streamCompletion = async (prompt: string, onDataFunction: (chunk: any) =>
 		'temperature': 0.9,
 		'stream': true,
 	};
-	return new Promise(async (resolve) => {
+	return new Promise(async (resolve, reject) => {
+		const cancelSource = axios.CancelToken.source();
 		const response = await axios({
 			method: 'post',
 			url: 'https://api.openai.com/v1/chat/completions',
 			headers: headers,
 			data: data,
 			responseType: 'stream',
+			cancelToken: cancelSource.token,
 		});
+		cancelFunction(cancelSource, resolve);
 		response.data.on('data', onDataFunction);
 		response.data.on('end', () => {
 			resolve();
 		});
 	});
 };
-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 } => {
-	const numberOfNewLines = countNewLines(textStream);
-	const newCharacterLocation = numberOfNewLines === 0 ? countCharacters(textStream) + currentCharacter : 0;
-	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.
-const useBufferToUpdateTextContentsWhileStreamingOpenAIResponse = async (activeEditor: TextEditor, position: Position, prompt: string) => {
+const useBufferToUpdateTextContentsWhileStreamingOpenAIResponse = async (activeEditor: TextEditor, position: Position, prompt: string, cancellationTokenSource: CancellationTokenSource): Promise<void> => {
 	let currentCharacter = position.character;
 	let currentLine = position.line;
 	let buffer: string[] = [];
 	let doneStreaming = false;
+	// Triggered on data response from openai
 	const onDataFunction = (word: any) => {
 		const newContent = word.toString().split('data: ')
 			.map((line: string) => line.trim())
@@ -72,7 +96,21 @@ const useBufferToUpdateTextContentsWhileStreamingOpenAIResponse = async (activeE
 			.map((line: string) => JSON.parse(line).choices[0].delta.content);
 		buffer = [...buffer, ...newContent];
 	};
-	streamCompletion(prompt, onDataFunction).then(() => doneStreaming = true);
+	// Triggered when the cancellationTokenSource is invoked
+	const cancelFunction = (source: any, resolve: any) => {
+		cancellationTokenSource.token.onCancellationRequested(() => {
+			doneStreaming = true;
+			source.cancel();
+			resolve();
+		});
+	};
+
+	streamCompletion(prompt, onDataFunction, cancelFunction).then(() => doneStreaming = true);
+
+	// While there is still data streaming or there is data in the buffer
+	// Shift data from buffer and add to the text editor
+	// Update currentCharacter and currentLine with the new cursor location
+	// Wait for the specified flush interval before processing the next data
 	while (!doneStreaming || buffer.length >= 0) {
 		const word: string | undefined = buffer.shift();
 		if (word) {
@@ -84,54 +122,41 @@ const useBufferToUpdateTextContentsWhileStreamingOpenAIResponse = async (activeE
 			currentCharacter = newCharacterLocation;
 			currentLine = newLineLocation;
 		}
-		// TODO I should make this buffer flush configurable.
-		await sleep(30);
+		await sleep(FLUSH_INTERVAL_MS);
 	}
+
+	return;
 };
-const constructPrompt = async (text: string, filepath: string): Promise<string> => {
-	const defaultPrompt = `
-ADDITIONAL CONTEXT:
-filepath: ${filepath}
-INSTRUCTIONS: 
+const directedHealPrompt = async (text: string, filepath: string): Promise<string> => {
+	const defaultInstructions = `
 - 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.
+- The code between [[SELECTED]] [[/SELECTED]] is the ONLY code that you should replace. So when you start responding, only respond with code that should replace it, and nothing else.
+- MAKE SURE YOU ONLY WRITE CODE TO REPLACE WHATS BETWEEN [[SELECTED]] AND [[/SELECTED]]. The rest of it doesnt need to be replaced
+- The rest of the code is ONLY 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.
+`;
+	const configurationInstructions = await workspace.getConfiguration().get('wolverine.prompt');
+	return `
+ADDITIONAL CONTEXT:
+filepath: ${filepath}
+INSTRUCTIONS:
+${configurationInstructions || defaultInstructions}
 CODE:
 ${text}
 NEW CODE TO REPLACE WHAT IS BETWEEN [[SELECTED]] and [/SELECTED]:
 `;
-	const configuredPrompt = await workspace.getConfiguration().get('wolverine.prompt');
-	if (configuredPrompt) {
-		return configuredPrompt + text;
-	}
-	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 cancellationTokenSource: any = undefined;
+	const directedHealDisposable = commands.registerCommand('wolverine.directedHeal', async () => {
+		commands.executeCommand('setContext', 'wolverine.operationRunning', true);
+		cancellationTokenSource = new CancellationTokenSource();
 
-	let disposable = commands.registerCommand('wolverine.directedHeal', async () => {
 		const activeEditor = window.activeTextEditor;
 		if (!activeEditor) {
 			window.showErrorMessage('No text editor is currently active.');
@@ -150,23 +175,29 @@ export async function activate(context: ExtensionContext) {
 		// 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 contextText = visibleText.replace(selectedText, '[[SELECTED]]\n' + selectedText + '\n[[/SELECTED]]');
 		const filePath = workspace.asRelativePath(document.uri);
-		const prompt = await constructPrompt(contextText, filePath);
+		const prompt = await directedHealPrompt(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 useBufferToUpdateTextContentsWhileStreamingOpenAIResponse(activeEditor, range.start, prompt, cancellationTokenSource);
+		commands.executeCommand('setContext', 'wolverine.operationRunning', false);
 		await activeEditor.document.save();
 	});
 
+	const cancelOpDisposable = commands.registerCommand('wolverine.cancelOperation', () => {
+		commands.executeCommand('setContext', 'wolverine.operationRunning', false);
+		if (cancellationTokenSource) {
+			cancellationTokenSource.cancel();
+			cancellationTokenSource = undefined;
+		}
+	});
 	// Add the disposable to the context subscriptions
-	context.subscriptions.push(disposable);
+	context.subscriptions.push(directedHealDisposable, cancelOpDisposable);
 }
 
 export function deactivate() { }