feat: add onProgress to ChatGPTAPIBrowser.sendMessage

shynsky · Jan 12, 2023 · e0fd5f4 · e0fd5f4
1 parent 525524b
commit e0fd5f4
Show file tree

Hide file tree

Showing 5 changed files with 142 additions and 14 deletions.
diff --git a/demos/demo-on-progress.ts 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
@@ -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
@@ -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
@@ -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,20 +401,23 @@ 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())
     if (minimize) {
       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
@@ -9,6 +9,12 @@ import stripMarkdown from 'strip-markdown'
 
 import * as types from './types'
 
+declare global {
+  function ChatGPTAPIBrowserOnProgress(
+    partialChatResponse: types.ChatResponse
+  ): Promise<void>
+}
+
 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,24 +149,32 @@ export async function browserPostEventStream(
 
     const responseP = new Promise<types.ChatResponse>(
       async (resolve, reject) => {
-        function onMessage(data: string) {
+        async function onMessage(data: string) {
           if (data === '[DONE]') {
             return resolve({
               response,
               conversationId,
               messageId
             })
           }
+
+          let convoResponseEvent: types.ConversationResponseEvent
           try {
-            const checkJson = JSON.parse(data)
-          } catch (error) {
-            console.log('warning: parse error.')
+            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)