Add a stop button (more testing required)

Fixes #20
bramses · Mar 23, 2023 · b0049bc · b0049bc
1 parent 7e4bf6d
commit b0049bc
Show file tree

Hide file tree

Showing 3 changed files with 181 additions and 124 deletions.
diff --git a/README.md b/README.md
@@ -25,7 +25,8 @@ https://user-images.githubusercontent.com/3282661/223005882-6632c997-b9a6-445b-8
 - [*Infer title* from messages](https://github.com/bramses/chatgpt-md/discussions/11). Can be set to run automatically after >4 messages.
 - Stream at cursor position or at end of file. Can be set in settings.
 - (NEW!) Choose [heading level for role](https://github.com/bramses/chatgpt-md/pull/22) h1-h6. Can be set in settings.
-- (NEW!) Custom endpoints can be specified using the url parameter in your front matter. See FAQ for an example
+- (NEW!) Custom endpoints can be specified using the url parameter in your front matter. See FAQ for an example.
+- (NEW!) Stop a running stream with a command. See commands section below.
 
 ### Commands
 
@@ -46,6 +47,10 @@ Create a new chat file from a template specified in `Chat Template Folder`. Reme
 
 [Infer the title of the chat from the messages](https://github.com/bramses/chatgpt-md/discussions/11). Requires at least 2 messages. Can be set in settings to run automatically after >4 messages.
 
+#### Stop Streaming
+
+Stops the stream. Useful if you want to stop the stream if you don't like where ChatGPT is heading/going too long.
+
 #### Add Divider
 
 Add a ChatGPT MD Horizontal Rule and `role::user`. 

diff --git a/main.ts b/main.ts
@@ -14,7 +14,7 @@ import {
 	Platform,
 } from "obsidian";
 
-import { streamSSE } from "./stream";
+import { StreamManager } from "./stream";
 import { unfinishedCodeBlock } from "helpers";
 
 interface ChatGPT_MDSettings {
@@ -64,6 +64,7 @@ export default class ChatGPT_MD extends Plugin {
 	settings: ChatGPT_MDSettings;
 
 	async callOpenAIAPI(
+		streamManager: StreamManager,
 		editor: Editor,
 		messages: { role: string; content: string }[],
 		model = "gpt-3.5-turbo",
@@ -98,7 +99,7 @@ export default class ChatGPT_MD extends Plugin {
 					// user: user, // not yet supported
 				};
 
-				const response = await streamSSE(
+				const response = await streamManager.streamSSE(
 					editor,
 					this.settings.apiKey,
 					url,
@@ -322,6 +323,7 @@ export default class ChatGPT_MD extends Plugin {
 	}
 
 	async inferTitleFromMessages(messages: string[]) {
+		console.log("[ChtGPT MD] Inferring Title")
 		try {
 			if (messages.length < 2) {
 				new Notice(
@@ -368,7 +370,8 @@ export default class ChatGPT_MD extends Plugin {
 				.trim()
 				.replace(/[:/\\]/g, "");
 		} catch (err) {
-			throw new Error("Error inferring title from messages" + err);
+			new Notice("[ChatGPT MD] Error inferring title from messages");
+			throw new Error("[ChatGPT MD] Error inferring title from messages" + err);
 		}
 	}
 
@@ -428,6 +431,8 @@ export default class ChatGPT_MD extends Plugin {
 
 		await this.loadSettings();
 
+		const streamManager = new StreamManager();
+
 		// This adds an editor command that can perform some operation on the current editor instance
 		this.addCommand({
 			id: "call-chatgpt-api",
@@ -471,6 +476,7 @@ export default class ChatGPT_MD extends Plugin {
 				}
 
 				this.callOpenAIAPI(
+					streamManager,
 					editor,
 					messagesWithRoleAndMessage,
 					frontmatter.model,
@@ -590,6 +596,15 @@ export default class ChatGPT_MD extends Plugin {
 			},
 		});
 
+		this.addCommand({
+			id: "stop-streaming",
+			name: "Stop streaming",
+			icon: "stop",
+			editorCallback: (editor: Editor, view: MarkdownView) => {
+				streamManager.stopStreaming();
+			},
+		});
+
 		this.addCommand({
 			id: "infer-title",
 			name: "Infer title",

diff --git a/stream.ts b/stream.ts
@@ -17,137 +17,174 @@ export interface OpenAIStreamPayload {
 	stream: boolean;
 }
 
-// todo: await this function and return the text on stream completion
-export const streamSSE = async (
-	editor: Editor,
-	apiKey: string,
-	url: string,
-	options: OpenAIStreamPayload,
-	setAtCursor: boolean,
-	headingPrefix: string
-) => {
-	return new Promise((resolve, reject) => {
-		try {
-			console.log("[ChatGPT MD] streamSSE", options);
-
-			const source = new SSE(url, {
-				headers: {
-					"Content-Type": "application/json",
-					Authorization: `Bearer ${apiKey}`,
-				},
-				method: "POST",
-				payload: JSON.stringify(options),
-			});
-
-			let txt = "";
-			let initialCursorPosCh = editor.getCursor().ch;
-			let initialCursorPosLine = editor.getCursor().line;
-
-			source.addEventListener("open", (e: any) => {
-				console.log("[ChatGPT MD] SSE Opened");
-
-				const newLine = `\n\n<hr class="__chatgpt_plugin">\n\n${headingPrefix}role::assistant\n\n`;
-				editor.replaceRange(newLine, editor.getCursor());
-
-				// move cursor to end of line
-				const cursor = editor.getCursor();
-				const newCursor = {
-					line: cursor.line,
-					ch: cursor.ch + newLine.length,
-				};
-				editor.setCursor(newCursor);
-
-				initialCursorPosCh = newCursor.ch;
-				initialCursorPosLine = newCursor.line;
-			});
-
-			source.addEventListener("message", (e: any) => {
-				if (e.data != "[DONE]") {
-					const payload = JSON.parse(e.data);
-					const text = payload.choices[0].delta.content;
-
-					// if text undefined, then do nothing
-					if (!text) {
-						return;
-					}
-
+export class StreamManager {
+	sse: any | null = null;
+	manualClose = false;
+
+	constructor() {}
+
+	stopStreaming = () => {
+		if (this.sse) {
+			this.sse.close();
+			console.log("[ChatGPT MD] SSE manually closed");
+			this.manualClose = true;
+			this.sse = null;
+		}
+	};
+
+	streamSSE = async (
+		editor: Editor,
+		apiKey: string,
+		url: string,
+		options: OpenAIStreamPayload,
+		setAtCursor: boolean,
+		headingPrefix: string
+	) => {
+		return new Promise((resolve, reject) => {
+			try {
+				console.log("[ChatGPT MD] streamSSE", options);
+
+				const source = new SSE(url, {
+					headers: {
+						"Content-Type": "application/json",
+						Authorization: `Bearer ${apiKey}`,
+					},
+					method: "POST",
+					payload: JSON.stringify(options),
+				});
+
+				this.sse = source;
+
+				let txt = "";
+				let initialCursorPosCh = editor.getCursor().ch;
+				let initialCursorPosLine = editor.getCursor().line;
+
+				source.addEventListener("open", (e: any) => {
+					console.log("[ChatGPT MD] SSE Opened");
+
+					const newLine = `\n\n<hr class="__chatgpt_plugin">\n\n${headingPrefix}role::assistant\n\n`;
+					editor.replaceRange(newLine, editor.getCursor());
+
+					// move cursor to end of line
 					const cursor = editor.getCursor();
-					const convPos = editor.posToOffset(cursor);
-
-					// @ts-ignore
-					const cm6 = editor.cm;
-					const transaction = cm6.state.update({
-						changes: { from: convPos, to: convPos, insert: text },
-					});
-					cm6.dispatch(transaction);
-
-					txt += text;
-
 					const newCursor = {
 						line: cursor.line,
-						ch: cursor.ch + text.length,
+						ch: cursor.ch + newLine.length,
 					};
 					editor.setCursor(newCursor);
-				} else {
-					source.close();
-					console.log("[ChatGPT MD] SSE Closed");
 
-          if (unfinishedCodeBlock(txt)) {
-            txt += "\n```";
-          }
+					initialCursorPosCh = newCursor.ch;
+					initialCursorPosLine = newCursor.line;
+				});
+
+				source.addEventListener("message", (e: any) => {
+					if (e.data != "[DONE]") {
+						const payload = JSON.parse(e.data);
+						const text = payload.choices[0].delta.content;
+
+						// if text undefined, then do nothing
+						if (!text) {
+							return;
+						}
+
+						const cursor = editor.getCursor();
+						const convPos = editor.posToOffset(cursor);
+
+						// @ts-ignore
+						const cm6 = editor.cm;
+						const transaction = cm6.state.update({
+							changes: {
+								from: convPos,
+								to: convPos,
+								insert: text,
+							},
+						});
+						cm6.dispatch(transaction);
 
-					// replace the text from initialCursor to fix any formatting issues from streaming
-					const cursor = editor.getCursor();
-					editor.replaceRange(
-						txt,
-						{ line: initialCursorPosLine, ch: initialCursorPosCh },
-						cursor
-					);
+						txt += text;
 
+						const newCursor = {
+							line: cursor.line,
+							ch: cursor.ch + text.length,
+						};
+						editor.setCursor(newCursor);
+					} else {
+						source.close();
+						console.log("[ChatGPT MD] SSE Closed");
+
+						if (unfinishedCodeBlock(txt)) {
+							txt += "\n```";
+						}
+
+						// replace the text from initialCursor to fix any formatting issues from streaming
+						const cursor = editor.getCursor();
+						editor.replaceRange(
+							txt,
+							{
+								line: initialCursorPosLine,
+								ch: initialCursorPosCh,
+							},
+							cursor
+						);
 
+						// set cursor to end of replacement text
+						const newCursor = {
+							line: initialCursorPosLine,
+							ch: initialCursorPosCh + txt.length,
+						};
+						editor.setCursor(newCursor);
+
+						if (!setAtCursor) {
+							// remove the text after the cursor
+							editor.replaceRange("", newCursor, {
+								line: Infinity,
+								ch: Infinity,
+							});
+						} else {
+							new Notice(
+								"[ChatGPT MD] Text pasted at cursor may leave artifacts. Please remove them manually. ChatGPT MD cannot safely remove text when pasting at cursor."
+							);
+						}
+
+						resolve(txt);
+						// return txt;
+					}
+				});
 
-					// set cursor to end of replacement text
-					const newCursor = {
-						line: initialCursorPosLine,
-						ch: initialCursorPosCh + txt.length,
-					};
-					editor.setCursor(newCursor);
+				source.addEventListener("abort", (e: any) => {
+					console.log("[ChatGPT MD] SSE Closed Event");
 
-					if (!setAtCursor) {
-						// remove the text after the cursor
-						editor.replaceRange("", newCursor, {
-							line: Infinity,
-							ch: Infinity,
-						});
+					// if e was triggered by stopStreaming, then resolve
+					if (this.manualClose) {
+						resolve(txt);
 					} else {
-						new Notice(
-							"[ChatGPT MD] Text pasted at cursor may leave artifacts. Please remove them manually. ChatGPT MD cannot safely remove text when pasting at cursor."
-						);
+						reject(e);
 					}
+				});
 
-					resolve(txt);
-					// return txt;
-				}
-			});
-
-			source.addEventListener("error", (e: any) => {
-				try {
-					console.log("[ChatGPT MD] SSE Error: ", JSON.parse(e.data));
-					source.close();
-					console.log("[ChatGPT MD] SSE Closed");
-					reject(JSON.parse(e.data));
-				} catch (err) {
-					console.log("[ChatGPT MD] Unknown Error: ", e);
-					source.close();
-					console.log("[ChatGPT MD] SSE Closed");
-					reject(e);
-				}
-			});
-
-			source.stream();
-		} catch (err) {
-			console.log("SSE Error", err);
-			reject(err);
-		}
-	});
-};
+
+				source.addEventListener("error", (e: any) => {
+					try {
+						console.log(
+							"[ChatGPT MD] SSE Error: ",
+							JSON.parse(e.data)
+						);
+						source.close();
+						console.log("[ChatGPT MD] SSE Closed");
+						reject(JSON.parse(e.data));
+					} catch (err) {
+						console.log("[ChatGPT MD] Unknown Error: ", e);
+						source.close();
+						console.log("[ChatGPT MD] SSE Closed");
+						reject(e);
+					}
+				});
+
+				source.stream();
+			} catch (err) {
+				console.log("SSE Error", err);
+				reject(err);
+			}
+		});
+	};
+}