Skip to content
This repository was archived by the owner on Sep 28, 2023. It is now read-only.

Commit f5193c9

Browse files
committed
Enable flexible models to accommodate paid plans
1 parent 71be9e8 commit f5193c9

File tree

6 files changed

+168
-130
lines changed

6 files changed

+168
-130
lines changed

src/background/chatgpt.ts

Lines changed: 0 additions & 82 deletions
This file was deleted.

src/background/index.ts

Lines changed: 10 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -1,54 +1,27 @@
1-
import ExpiryMap from 'expiry-map'
21
import Browser from 'webextension-polyfill'
3-
import { sendMessage, sendMessageFeedback, setConversationProperty } from './chatgpt.js'
4-
5-
const KEY_ACCESS_TOKEN = 'accessToken'
6-
7-
const cache = new ExpiryMap(10 * 1000)
8-
9-
async function getAccessToken(): Promise<string> {
10-
if (cache.get(KEY_ACCESS_TOKEN)) {
11-
return cache.get(KEY_ACCESS_TOKEN)
12-
}
13-
const resp = await fetch('https://chat.openai.com/api/auth/session')
14-
if (resp.status === 403) {
15-
throw new Error('CLOUDFLARE')
16-
}
17-
const data = await resp.json().catch(() => ({}))
18-
if (!data.accessToken) {
19-
throw new Error('UNAUTHORIZED')
20-
}
21-
cache.set(KEY_ACCESS_TOKEN, data.accessToken)
22-
return data.accessToken
23-
}
2+
import { ChatGPTProvider, getChatGPTAccessToken, sendMessageFeedback } from './providers/chatgpt'
3+
import { Provider } from './types'
244

255
async function generateAnswers(port: Browser.Runtime.Port, question: string) {
26-
const accessToken = await getAccessToken()
276

28-
let conversationId: string | undefined
29-
const deleteConversation = () => {
30-
if (conversationId) {
31-
setConversationProperty(accessToken, conversationId, { is_visible: false })
32-
}
33-
}
7+
let provider: Provider
8+
const token = await getChatGPTAccessToken()
9+
provider = new ChatGPTProvider(token)
3410

3511
const controller = new AbortController()
3612
port.onDisconnect.addListener(() => {
3713
controller.abort()
38-
deleteConversation()
14+
cleanup?.()
3915
})
4016

41-
await sendMessage({
42-
token: accessToken,
17+
const { cleanup } = await provider.generateAnswer({
4318
prompt: question,
4419
signal: controller.signal,
4520
onEvent(event) {
4621
if (event.type === 'done') {
4722
port.postMessage({ event: 'DONE' })
48-
deleteConversation()
4923
return
5024
}
51-
conversationId = event.data.conversationId
5225
port.postMessage(event.data)
5326
},
5427
})
@@ -62,24 +35,15 @@ Browser.runtime.onConnect.addListener((port) => {
6235
} catch (err: any) {
6336
console.error(err)
6437
port.postMessage({ error: err.message })
65-
cache.delete(KEY_ACCESS_TOKEN)
6638
}
6739
})
6840
})
6941

7042
Browser.runtime.onMessage.addListener(async (message) => {
7143
if (message.type === 'FEEDBACK') {
72-
const token = await getAccessToken()
44+
const token = await getChatGPTAccessToken()
7345
await sendMessageFeedback(token, message.data)
74-
} else if (message.type === 'OPEN_OPTIONS_PAGE') {
75-
Browser.runtime.openOptionsPage()
7646
} else if (message.type === 'GET_ACCESS_TOKEN') {
77-
return getAccessToken()
78-
}
79-
})
80-
81-
Browser.runtime.onInstalled.addListener((details) => {
82-
if (details.reason === 'install') {
83-
Browser.runtime.openOptionsPage()
47+
return getChatGPTAccessToken()
8448
}
85-
})
49+
})

src/background/providers/chatgpt.ts

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
import { v4 as uuidv4 } from 'uuid'
2+
import ExpiryMap from 'expiry-map'
3+
import { fetchSSE } from '../fetch-sse'
4+
import { GenerateAnswerParams, Provider } from '../types'
5+
6+
async function request(token: string, method: string, path: string, data?: unknown) {
7+
return fetch(`https://chat.openai.com/backend-api${path}`, {
8+
method,
9+
headers: {
10+
'Content-Type': 'application/json',
11+
Authorization: `Bearer ${token}`,
12+
},
13+
body: data === undefined ? undefined : JSON.stringify(data),
14+
})
15+
}
16+
17+
18+
export async function sendMessageFeedback(token: string, data: unknown) {
19+
await request(token, 'POST', '/conversation/message_feedback', data)
20+
}
21+
22+
export async function setConversationProperty(
23+
token: string,
24+
conversationId: string,
25+
propertyObject: object,
26+
) {
27+
await request(token, 'PATCH', `/conversation/${conversationId}`, propertyObject)
28+
}
29+
30+
const KEY_ACCESS_TOKEN = 'accessToken'
31+
32+
const cache = new ExpiryMap(10 * 1000)
33+
34+
export async function getChatGPTAccessToken(): Promise<string> {
35+
if (cache.get(KEY_ACCESS_TOKEN)) {
36+
return cache.get(KEY_ACCESS_TOKEN)
37+
}
38+
const resp = await fetch('https://chat.openai.com/api/auth/session')
39+
if (resp.status === 403) {
40+
throw new Error('CLOUDFLARE')
41+
}
42+
const data = await resp.json().catch(() => ({}))
43+
if (!data.accessToken) {
44+
throw new Error('UNAUTHORIZED')
45+
}
46+
cache.set(KEY_ACCESS_TOKEN, data.accessToken)
47+
return data.accessToken
48+
}
49+
50+
export class ChatGPTProvider implements Provider {
51+
constructor(private token: string) {
52+
this.token = token
53+
}
54+
55+
private async fetchModels(): Promise<
56+
{ slug: string; title: string; description: string; max_tokens: number }[]
57+
> {
58+
const resp = await request(this.token, 'GET', '/models').then((r) => r.json())
59+
return resp.models
60+
}
61+
62+
private async getModelName(): Promise<string> {
63+
try {
64+
const models = await this.fetchModels()
65+
return models[0].slug
66+
} catch (err) {
67+
console.error(err)
68+
return 'text-davinci-002-render'
69+
}
70+
}
71+
72+
async generateAnswer(params: GenerateAnswerParams) {
73+
let conversationId: string | undefined
74+
75+
const cleanup = () => {
76+
if (conversationId) {
77+
setConversationProperty(this.token, conversationId, { is_visible: false })
78+
}
79+
}
80+
81+
const modelName = await this.getModelName()
82+
console.log('Using model:', modelName)
83+
84+
await fetchSSE('https://chat.openai.com/backend-api/conversation', {
85+
method: 'POST',
86+
signal: params.signal,
87+
headers: {
88+
'Content-Type': 'application/json',
89+
Authorization: `Bearer ${this.token}`,
90+
},
91+
body: JSON.stringify({
92+
action: 'next',
93+
messages: [
94+
{
95+
id: uuidv4(),
96+
role: 'user',
97+
content: {
98+
content_type: 'text',
99+
parts: [params.prompt],
100+
},
101+
},
102+
],
103+
model: modelName,
104+
parent_message_id: uuidv4(),
105+
}),
106+
onMessage(message: string) {
107+
console.debug('sse message', message)
108+
if (message === '[DONE]') {
109+
params.onEvent({ type: 'done' })
110+
cleanup()
111+
return
112+
}
113+
let data
114+
try {
115+
data = JSON.parse(message)
116+
} catch (err) {
117+
console.error(err)
118+
return
119+
}
120+
const text = data.message?.content?.parts?.[0]
121+
if (text) {
122+
conversationId = data.conversation_id
123+
params.onEvent({
124+
type: 'answer',
125+
data: {
126+
text,
127+
messageId: data.message.id,
128+
conversationId: data.conversation_id,
129+
},
130+
})
131+
}
132+
},
133+
})
134+
return { cleanup }
135+
}
136+
}

src/background/types.ts

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
import { Answer } from '../messaging'
2+
3+
export type Event =
4+
| {
5+
type: 'answer'
6+
data: Answer
7+
}
8+
| {
9+
type: 'done'
10+
}
11+
12+
export interface GenerateAnswerParams {
13+
prompt: string
14+
onEvent: (event: Event) => void
15+
signal?: AbortSignal
16+
}
17+
18+
export interface Provider {
19+
generateAnswer(params: GenerateAnswerParams): Promise<{ cleanup?: () => void }>
20+
}

src/manifest.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
{
22
"name": "ChatGPT - Jupyter - AI Assistant",
33
"description": "ChatGPT-powered AI assistant for Jupyter Notebooks",
4-
"version": "0.3.2",
4+
"version": "0.3.5",
55
"manifest_version": 3,
66
"icons": {
77
"16": "logo.png",

src/manifest.v2.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
{
22
"name": "ChatGPT - Jupyter - AI Assistant",
33
"description": "ChatGPT-powered AI assistant for Jupyter Notebooks",
4-
"version": "0.3.2",
4+
"version": "0.3.5",
55
"manifest_version": 2,
66
"icons": {
77
"16": "logo.png",

0 commit comments

Comments
 (0)