From e0fd5f4652cb1dd2a6ac77f0780fee32c751115c Mon Sep 17 00:00:00 2001 From: Travis Fischer Date: Thu, 12 Jan 2023 03:55:20 -0600 Subject: [PATCH] feat: add onProgress to ChatGPTAPIBrowser.sendMessage --- demos/demo-on-progress.ts | 47 ++++++++++++++++++++++++++++++++++++++ readme.md | 14 +++++++++++- src/chatgpt-api-browser.ts | 43 +++++++++++++++++++++++++++++++--- src/openai-auth.ts | 14 ++++++++---- src/utils.ts | 40 ++++++++++++++++++++++++++------ 5 files changed, 143 insertions(+), 15 deletions(-) create mode 100644 demos/demo-on-progress.ts diff --git a/demos/demo-on-progress.ts b/demos/demo-on-progress.ts new file mode 100644 index 0000000..a8a259e --- /dev/null +++ b/demos/demo-on-progress.ts @@ -0,0 +1,47 @@ +import dotenv from 'dotenv-safe' +import { oraPromise } from 'ora' + +import { ChatGPTAPIBrowser } from '../src' + +dotenv.config() + +/** + * Demo CLI for testing the `onProgress` handler. + * + * ``` + * npx tsx demos/demo-on-progress.ts + * ``` + */ +async function main() { + const email = process.env.OPENAI_EMAIL + const password = process.env.OPENAI_PASSWORD + + const api = new ChatGPTAPIBrowser({ + email, + password, + debug: false, + minimize: true + }) + await api.initSession() + + const prompt = + 'Write a python version of bubble sort. Do not include example usage.' + + console.log(prompt) + + const res = await api.sendMessage(prompt, { + onProgress: (partialResponse) => { + console.log('p') + console.log('progress', partialResponse?.response) + } + }) + console.log(res.response) + + // close the browser at the end + await api.closeSession() +} + +main().catch((err) => { + console.error(err) + process.exit(1) +}) diff --git a/readme.md b/readme.md index 8093d69..8ae50dd 100644 --- a/readme.md +++ b/readme.md @@ -197,7 +197,19 @@ A [basic demo](./demos/demo.ts) is included for testing purposes: npx tsx demos/demo.ts ``` -A [conversation demo](./demos/demo-conversation.ts) is also included: +A [google auth demo](./demos/demo-google-auth.ts): + +```bash +npx tsx demos/demo-google-auth.ts +``` + +A [demo showing on progress handler](./demos/demo-on-progress.ts): + +```bash +npx tsx demos/demo-on-progress.ts +``` + +A [conversation demo](./demos/demo-conversation.ts): ```bash npx tsx demos/demo-conversation.ts diff --git a/src/chatgpt-api-browser.ts b/src/chatgpt-api-browser.ts index 885170e..64ea367 100644 --- a/src/chatgpt-api-browser.ts +++ b/src/chatgpt-api-browser.ts @@ -33,6 +33,10 @@ export class ChatGPTAPIBrowser extends AChatGPTAPI { protected _page: Page protected _proxyServer: string protected _isRefreshing: boolean + protected _messageOnProgressHandlers: Record< + string, + (partialResponse: types.ChatResponse) => void + > /** * Creates a new client for automating the ChatGPT webapp. @@ -97,6 +101,7 @@ export class ChatGPTAPIBrowser extends AChatGPTAPI { this._executablePath = executablePath this._proxyServer = proxyServer this._isRefreshing = false + this._messageOnProgressHandlers = {} if (!this._email) { const error = new types.ChatGPTError('ChatGPT invalid email') @@ -196,6 +201,24 @@ export class ChatGPTAPIBrowser extends AChatGPTAPI { }) } + // TODO: will this exist after page reload and navigation? + await this._page.exposeFunction( + 'ChatGPTAPIBrowserOnProgress', + (partialResponse: types.ChatResponse) => { + if ((partialResponse as any)?.origMessageId) { + const onProgress = + this._messageOnProgressHandlers[ + (partialResponse as any).origMessageId + ] + + if (onProgress) { + onProgress(partialResponse) + return + } + } + } + ) + // dismiss welcome modal (and other modals) do { const modalSelector = '[data-headlessui-state="open"]' @@ -482,9 +505,8 @@ export class ChatGPTAPIBrowser extends AChatGPTAPI { parentMessageId = uuidv4(), messageId = uuidv4(), action = 'next', - timeoutMs - // TODO - // onProgress + timeoutMs, + onProgress } = opts const url = `https://chat.openai.com/backend-api/conversation` @@ -508,6 +530,16 @@ export class ChatGPTAPIBrowser extends AChatGPTAPI { body.conversation_id = conversationId } + if (onProgress) { + this._messageOnProgressHandlers[messageId] = onProgress + } + + const cleanup = () => { + if (this._messageOnProgressHandlers[messageId]) { + delete this._messageOnProgressHandlers[messageId] + } + } + let result: types.ChatResponse | types.ChatError let numTries = 0 let is401 = false @@ -528,6 +560,7 @@ export class ChatGPTAPIBrowser extends AChatGPTAPI { if (!(await this.getIsAuthenticated())) { const error = new types.ChatGPTError('Not signed in') error.statusCode = 401 + cleanup() throw error } } @@ -551,6 +584,7 @@ export class ChatGPTAPIBrowser extends AChatGPTAPI { const error = new types.ChatGPTError(err.toString()) error.statusCode = err.response?.statusCode error.statusText = err.response?.statusText + cleanup() throw error } @@ -570,6 +604,7 @@ export class ChatGPTAPIBrowser extends AChatGPTAPI { is401 = true if (numTries >= 2) { + cleanup() throw error } else { continue @@ -590,10 +625,12 @@ export class ChatGPTAPIBrowser extends AChatGPTAPI { result.response = markdownToText(result.response) } + cleanup() return result } } while (!result) + cleanup() // console.log('<<< EVALUATE', result) // const lastMessage = await this.getLastMessage() diff --git a/src/openai-auth.ts b/src/openai-auth.ts index 9983713..dd75776 100644 --- a/src/openai-auth.ts +++ b/src/openai-auth.ts @@ -272,6 +272,7 @@ export async function getBrowser( nopechaKey?: string proxyServer?: string minimize?: boolean + debug?: boolean timeoutMs?: number } = {} ) { @@ -281,6 +282,7 @@ export async function getBrowser( executablePath = defaultChromeExecutablePath(), proxyServer = process.env.PROXY_SERVER, minimize = false, + debug = false, timeoutMs = DEFAULT_TIMEOUT_MS, ...launchOptions } = opts @@ -387,8 +389,9 @@ export async function getBrowser( } await initializeNopechaExtension(browser, { - minimize, nopechaKey, + minimize, + debug, timeoutMs }) @@ -398,12 +401,13 @@ export async function getBrowser( export async function initializeNopechaExtension( browser: Browser, opts: { - minimize?: boolean nopechaKey?: string + minimize?: boolean + debug?: boolean timeoutMs?: number } ) { - const { minimize = false, nopechaKey } = opts + const { minimize = false, debug = false, nopechaKey } = opts if (hasNopechaExtension) { const page = (await browser.pages())[0] || (await browser.newPage()) @@ -411,7 +415,9 @@ export async function initializeNopechaExtension( await minimizePage(page) } - console.log('initializing nopecha extension with key', nopechaKey, '...') + if (debug) { + console.log('initializing nopecha extension with key', nopechaKey, '...') + } // TODO: setting the nopecha extension key is really, really error prone... for (let i = 0; i < 5; ++i) { diff --git a/src/utils.ts b/src/utils.ts index 180d87b..5f4cfb3 100644 --- a/src/utils.ts +++ b/src/utils.ts @@ -9,6 +9,12 @@ import stripMarkdown from 'strip-markdown' import * as types from './types' +declare global { + function ChatGPTAPIBrowserOnProgress( + partialChatResponse: types.ChatResponse + ): Promise +} + export function markdownToText(markdown?: string): string { return remark() .use(stripMarkdown) @@ -103,6 +109,7 @@ export async function browserPostEventStream( const BOM = [239, 187, 191] let conversationId: string = body?.conversation_id + const origMessageId = body?.messages?.[0]?.id let messageId: string = body?.messages?.[0]?.id let response = '' @@ -142,7 +149,7 @@ export async function browserPostEventStream( const responseP = new Promise( async (resolve, reject) => { - function onMessage(data: string) { + async function onMessage(data: string) { if (data === '[DONE]') { return resolve({ response, @@ -150,16 +157,24 @@ export async function browserPostEventStream( messageId }) } - try { - const checkJson = JSON.parse(data) - } catch (error) { - console.log('warning: parse error.') + let convoResponseEvent: types.ConversationResponseEvent + try { + convoResponseEvent = JSON.parse(data) + } catch (err) { + console.warn( + 'warning: chatgpt even stream parse error', + err.toString(), + data + ) return } + + if (!convoResponseEvent) { + return + } + try { - const convoResponseEvent: types.ConversationResponseEvent = - JSON.parse(data) if (convoResponseEvent.conversation_id) { conversationId = convoResponseEvent.conversation_id } @@ -172,6 +187,17 @@ export async function browserPostEventStream( convoResponseEvent.message?.content?.parts?.[0] if (partialResponse) { response = partialResponse + + if (window.ChatGPTAPIBrowserOnProgress) { + const partialChatResponse = { + origMessageId, + response, + conversationId, + messageId + } + + await window.ChatGPTAPIBrowserOnProgress(partialChatResponse) + } } } catch (err) { console.warn('fetchSSE onMessage unexpected error', err)