Enable flexible models to accommodate paid plans

Author: TiesdeKok

src/background/chatgpt.ts

@@ -1,54 +1,27 @@
-import ExpiryMap from 'expiry-map'
 import Browser from 'webextension-polyfill'
-import { sendMessage, sendMessageFeedback, setConversationProperty } from './chatgpt.js'
-
-const KEY_ACCESS_TOKEN = 'accessToken'
-
-const cache = new ExpiryMap(10 * 1000)
-
-async function getAccessToken(): Promise<string> {
-  if (cache.get(KEY_ACCESS_TOKEN)) {
-    return cache.get(KEY_ACCESS_TOKEN)
-  }
-  const resp = await fetch('https://chat.openai.com/api/auth/session')
-  if (resp.status === 403) {
-    throw new Error('CLOUDFLARE')
-  }
-  const data = await resp.json().catch(() => ({}))
-  if (!data.accessToken) {
-    throw new Error('UNAUTHORIZED')
-  }
-  cache.set(KEY_ACCESS_TOKEN, data.accessToken)
-  return data.accessToken
-}
+import { ChatGPTProvider, getChatGPTAccessToken, sendMessageFeedback } from './providers/chatgpt'
+import { Provider } from './types'
 
 async function generateAnswers(port: Browser.Runtime.Port, question: string) {
-  const accessToken = await getAccessToken()
 
-  let conversationId: string | undefined
-  const deleteConversation = () => {
-    if (conversationId) {
-      setConversationProperty(accessToken, conversationId, { is_visible: false })
-    }
-  }
+  let provider: Provider
+  const token = await getChatGPTAccessToken()
+  provider = new ChatGPTProvider(token)
 
   const controller = new AbortController()
   port.onDisconnect.addListener(() => {
     controller.abort()
-    deleteConversation()
+    cleanup?.()
   })
 
-  await sendMessage({
-    token: accessToken,
+  const { cleanup } = await provider.generateAnswer({
     prompt: question,
     signal: controller.signal,
     onEvent(event) {
       if (event.type === 'done') {
         port.postMessage({ event: 'DONE' })
-        deleteConversation()
         return
       }
-      conversationId = event.data.conversationId
       port.postMessage(event.data)
     },
   })
@@ -62,24 +35,15 @@ Browser.runtime.onConnect.addListener((port) => {
     } catch (err: any) {
       console.error(err)
       port.postMessage({ error: err.message })
-      cache.delete(KEY_ACCESS_TOKEN)
     }
   })
 })
 
 Browser.runtime.onMessage.addListener(async (message) => {
   if (message.type === 'FEEDBACK') {
-    const token = await getAccessToken()
+    const token = await getChatGPTAccessToken()
     await sendMessageFeedback(token, message.data)
-  } else if (message.type === 'OPEN_OPTIONS_PAGE') {
-    Browser.runtime.openOptionsPage()
   } else if (message.type === 'GET_ACCESS_TOKEN') {
-    return getAccessToken()
-  }
-})
-
-Browser.runtime.onInstalled.addListener((details) => {
-  if (details.reason === 'install') {
-    Browser.runtime.openOptionsPage()
+    return getChatGPTAccessToken()
   }
-})
+})

src/background/index.ts

@@ -0,0 +1,136 @@
+import { v4 as uuidv4 } from 'uuid'
+import ExpiryMap from 'expiry-map'
+import { fetchSSE } from '../fetch-sse'
+import { GenerateAnswerParams, Provider } from '../types'
+
+async function request(token: string, method: string, path: string, data?: unknown) {
+    return fetch(`https://chat.openai.com/backend-api${path}`, {
+      method,
+      headers: {
+        'Content-Type': 'application/json',
+        Authorization: `Bearer ${token}`,
+      },
+      body: data === undefined ? undefined : JSON.stringify(data),
+    })
+  }
+  
+
+export async function sendMessageFeedback(token: string, data: unknown) {
+  await request(token, 'POST', '/conversation/message_feedback', data)
+}
+
+export async function setConversationProperty(
+  token: string,
+  conversationId: string,
+  propertyObject: object,
+) {
+  await request(token, 'PATCH', `/conversation/${conversationId}`, propertyObject)
+}
+
+const KEY_ACCESS_TOKEN = 'accessToken'
+
+const cache = new ExpiryMap(10 * 1000)
+
+export async function getChatGPTAccessToken(): Promise<string> {
+  if (cache.get(KEY_ACCESS_TOKEN)) {
+    return cache.get(KEY_ACCESS_TOKEN)
+  }
+  const resp = await fetch('https://chat.openai.com/api/auth/session')
+  if (resp.status === 403) {
+    throw new Error('CLOUDFLARE')
+  }
+  const data = await resp.json().catch(() => ({}))
+  if (!data.accessToken) {
+    throw new Error('UNAUTHORIZED')
+  }
+  cache.set(KEY_ACCESS_TOKEN, data.accessToken)
+  return data.accessToken
+}
+
+export class ChatGPTProvider implements Provider {
+  constructor(private token: string) {
+    this.token = token
+  }
+
+  private async fetchModels(): Promise<
+    { slug: string; title: string; description: string; max_tokens: number }[]
+  > {
+    const resp = await request(this.token, 'GET', '/models').then((r) => r.json())
+    return resp.models
+  }
+
+  private async getModelName(): Promise<string> {
+    try {
+      const models = await this.fetchModels()
+      return models[0].slug
+    } catch (err) {
+      console.error(err)
+      return 'text-davinci-002-render'
+    }
+  }
+
+  async generateAnswer(params: GenerateAnswerParams) {
+    let conversationId: string | undefined
+
+    const cleanup = () => {
+      if (conversationId) {
+        setConversationProperty(this.token, conversationId, { is_visible: false })
+      }
+    }
+
+    const modelName = await this.getModelName()
+    console.log('Using model:', modelName)
+
+    await fetchSSE('https://chat.openai.com/backend-api/conversation', {
+      method: 'POST',
+      signal: params.signal,
+      headers: {
+        'Content-Type': 'application/json',
+        Authorization: `Bearer ${this.token}`,
+      },
+      body: JSON.stringify({
+        action: 'next',
+        messages: [
+          {
+            id: uuidv4(),
+            role: 'user',
+            content: {
+              content_type: 'text',
+              parts: [params.prompt],
+            },
+          },
+        ],
+        model: modelName,
+        parent_message_id: uuidv4(),
+      }),
+      onMessage(message: string) {
+        console.debug('sse message', message)
+        if (message === '[DONE]') {
+          params.onEvent({ type: 'done' })
+          cleanup()
+          return
+        }
+        let data
+        try {
+          data = JSON.parse(message)
+        } catch (err) {
+          console.error(err)
+          return
+        }
+        const text = data.message?.content?.parts?.[0]
+        if (text) {
+          conversationId = data.conversation_id
+          params.onEvent({
+            type: 'answer',
+            data: {
+              text,
+              messageId: data.message.id,
+              conversationId: data.conversation_id,
+            },
+          })
+        }
+      },
+    })
+    return { cleanup }
+  }
+}

src/background/providers/chatgpt.ts

@@ -0,0 +1,20 @@
+import { Answer } from '../messaging'
+
+export type Event =
+  | {
+      type: 'answer'
+      data: Answer
+    }
+  | {
+      type: 'done'
+    }
+
+export interface GenerateAnswerParams {
+  prompt: string
+  onEvent: (event: Event) => void
+  signal?: AbortSignal
+}
+
+export interface Provider {
+  generateAnswer(params: GenerateAnswerParams): Promise<{ cleanup?: () => void }>
+}

src/background/types.ts

@@ -1,7 +1,7 @@
 {
   "name": "ChatGPT - Jupyter - AI Assistant",
   "description": "ChatGPT-powered AI assistant for Jupyter Notebooks",
-  "version": "0.3.2",
+  "version": "0.3.5",
   "manifest_version": 3,
   "icons": {
     "16": "logo.png",

src/manifest.json

@@ -1,7 +1,7 @@
 {
   "name": "ChatGPT - Jupyter - AI Assistant",
   "description": "ChatGPT-powered AI assistant for Jupyter Notebooks",
-  "version": "0.3.2",
+  "version": "0.3.5",
   "manifest_version": 2,
   "icons": {
     "16": "logo.png",