From b0049bca743a9152584e7ab65d693f2580d3b493 Mon Sep 17 00:00:00 2001 From: Bram Adams <3282661+bramses@users.noreply.github.com> Date: Thu, 23 Mar 2023 12:08:24 -0400 Subject: [PATCH] Add a stop button (more testing required) Fixes #20 --- README.md | 7 +- main.ts | 21 ++++- stream.ts | 277 +++++++++++++++++++++++++++++++----------------------- 3 files changed, 181 insertions(+), 124 deletions(-) diff --git a/README.md b/README.md index 101c795..83fb4da 100644 --- 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 index 3b9dca9..e957e68 100644 --- 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 index 44121d7..bd538de 100644 --- 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
\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
\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); + } + }); + }; +}