Skip to content

Commit

Permalink
wss_url
Browse files Browse the repository at this point in the history
  • Loading branch information
@ishandutta2007
ishandutta2007 committed Feb 16, 2024
1 parent 77f43c5 commit 3872e80
Show file tree
Hide file tree
Showing 9 changed files with 1,288 additions and 325 deletions.
1,213 changes: 928 additions & 285 deletions package-lock.json
View file

Large diffs are not rendered by default.

10 changes: 8 additions & 2 deletions package.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,10 @@
"@geist-ui/icons": "^1.0.2",
"@primer/octicons-react": "^17.9.0",
"@ungap/structured-clone": "^1.2.0",
"buffer": "^6.0.3",
"dayjs": "^1.11.9",
"esbuild-copy-static-files": "^0.1.0",
"eventsource-parser": "^0.0.5",
"eventsource-parser": "^1.0.0",
"expiry-map": "^2.0.0",
"github-markdown-css": "^5.1.0",
"i18next-browser-languagedetector": "^7.0.1",
Expand All @@ -32,7 +33,8 @@
"react-toastify": "^9.1.3",
"rehype-highlight": "^6.0.0",
"swr": "^2.0.0",
"uuid": "^9.0.0"
"uuid": "^9.0.0",
"websocket-as-promised": "^2.0.1"
},
"devDependencies": {
"@types/fs-extra": "^9.0.13",
Expand Down Expand Up @@ -68,5 +70,9 @@
"npx prettier --write",
"npx eslint --fix"
]
},
"engines": {
"node": ">=16.15.0 <17.0.0",
"npm": ">=8.19.4 <9.0.0"
}
}
4 changes: 4 additions & 0 deletions src/background/fetch-sse.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,10 @@ export async function fetchSSE(
})
for await (const chunk of streamAsyncIterable(resp.body!)) {
const str = new TextDecoder().decode(chunk)
console.log('fetchSSE', str)
if (str.includes('wss_url')) {
onMessage(str)
}
parser.feed(str)
}
}
4 changes: 4 additions & 0 deletions src/background/index.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,10 @@ async function generateAnswers(
prompt: question,
signal: controller.signal,
onEvent(event) {
if (event.type === 'error') {
port.postMessage({ event: 'ERROR', message: event.message })
return
}
if (event.type === 'done') {
port.postMessage({ event: 'DONE' })
return
Expand Down
241 changes: 208 additions & 33 deletions src/background/providers/chatgpt.ts
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,13 @@
import { Buffer } from 'buffer'
import dayjs from 'dayjs'
import { createParser } from 'eventsource-parser'
import ExpiryMap from 'expiry-map'
import { v4 as uuidv4 } from 'uuid'
import { ADAY, APPSHORTNAME, HALFHOUR } from '../../utils/consts'
import Browser from 'webextension-polyfill'
import WebSocketAsPromised from 'websocket-as-promised'
import { ChatgptMode, getUserConfig } from '~config'
import { ADAY, APPSHORTNAME, HALFHOUR } from '~utils/consts'
import { parseSSEResponse } from '~utils/sse'
import { fetchSSE } from '../fetch-sse'
import { GenerateAnswerParams, Provider } from '../types'
dayjs().format()
Expand All @@ -17,6 +23,33 @@ async function request(token: string, method: string, path: string, data?: unkno
})
}

function removeCitations(text: string) {
return text.replaceAll(/\u3010\d+\u2020source\u3011/g, '')
}

const getConversationTitle = (bigtext: string) => {
let ret = bigtext.split('\n', 1)[0]
ret = ret.split('.', 1)[0]
ret = APPSHORTNAME + ':' + ret.split(':')[1].trim()
console.log('getConversationTitle:', ret)
return ret
}

const countWords = (text) => {
return text.trim().split(/\s+/).length
}

async function getChatgptwssIsOpenFlag() {
const { chatgptwssIsOpenFlag = false } = await Browser.storage.sync.get('chatgptwssIsOpenFlag')
return chatgptwssIsOpenFlag
}

async function setChatgptwssIsOpenFlag(isOpen: boolean) {
const { chatgptwssIsOpenFlag = false } = await Browser.storage.sync.get('chatgptwssIsOpenFlag')
Browser.storage.sync.set({ chatgptwssIsOpenFlag: isOpen })
return chatgptwssIsOpenFlag
}

async function request_new(
token: string,
method: string,
Expand Down Expand Up @@ -162,41 +195,17 @@ export class ChatGPTProvider implements Provider {
}
}

async generateAnswer(params: GenerateAnswerParams) {
let conversationId: string | undefined

const countWords = (text) => {
return text.trim().split(/\s+/).length
}

const getConversationTitle = (bigtext: string) => {
let ret = bigtext.split('\n', 1)[0]
ret = ret.split('.', 1)[0]
ret = APPSHORTNAME + ':' + ret.split(':')[1].trim()
console.log('getConversationTitle:', ret)
return ret
}

const renameConversationTitle = (convId: string) => {
const titl: string = getConversationTitle(params.prompt)
console.log('renameConversationTitle:', this.token, convId, titl)
setConversationProperty(this.token, convId, { title: titl })
}
const cleanup = () => {
if (conversationId) {
// setConversationProperty(this.token, conversationId, { is_visible: false })
}
}

async generateAnswerBySSE(params: GenerateAnswerParams, cleanup: () => void) {
console.debug('ChatGPTProvider:generateAnswerBySSE:', params)
const modelName = await this.getModelName()
console.debug('Using model:', modelName)

console.debug('ChatGPTProvider:modelName:', modelName)
await fetchSSE('https://chat.openai.com/backend-api/conversation', {
method: 'POST',
signal: params.signal,
headers: {
'Content-Type': 'application/json',
Authorization: `Bearer ${this.token}`,
'Openai-Sentinel-Arkose-Token': params.arkoseToken,
},
body: JSON.stringify({
action: 'next',
Expand All @@ -216,7 +225,13 @@ export class ChatGPTProvider implements Provider {
arkose_token: params.arkoseToken,
}),
onMessage(message: string) {
console.debug('sse message', message)
console.debug('ChatGPTProvider:generateAnswerBySSE:message', message)
let finaltext = ''
if (message.includes('wss_url')) {
params.onEvent({ type: 'error', message: message })
cleanup()
return
}
if (message === '[DONE]') {
params.onEvent({ type: 'done' })
cleanup()
Expand All @@ -231,16 +246,16 @@ export class ChatGPTProvider implements Provider {
}
const text = data.message?.content?.parts?.[0]
if (text) {
finaltext += text
if (countWords(text) == 1 && data.message?.author?.role == 'assistant') {
if (params.prompt.indexOf('search query:') !== -1) {
renameConversationTitle(data.conversation_id)
this.renameConversationTitle(data.conversation_id)
}
}
conversationId = data.conversation_id
params.onEvent({
type: 'answer',
data: {
text,
finaltext,
messageId: data.message.id,
parentMessageId: data.parent_message_id,
conversationId: data.conversation_id,
Expand All @@ -249,6 +264,166 @@ export class ChatGPTProvider implements Provider {
}
},
})
}

async setupWSS(params: GenerateAnswerParams, regResp: any) {
console.log('ChatGPTProvider:setupWSS:regResp', regResp)
let jj
await parseSSEResponse(regResp, (message) => {
console.log('ChatGPTProvider:setupWSS:parseSSEResponse:message', message)
jj = JSON.parse(message)
})
console.log('ChatGPTProvider:jj', jj)
if (jj) {
const wsAddress = jj['wss_url']
const wsp: WebSocketAsPromised = new WebSocketAsPromised(wsAddress, {
createWebSocket: (url) => {
const ws = new WebSocket(wsAddress, [
'Sec-Websocket-Protocol',
'json.reliable.webpubsub.azure.v1',
])
ws.binaryType = 'arraybuffer'
return ws
},
})
console.log('ChatGPTProvider:setupWebsocket:wsp', wsp)

const openListener = async () => {
console.log('ChatGPTProvider:setupWSSopenListener::wsp.onOpen')
await setChatgptwssIsOpenFlag(true)
}

let next_check_seqid = Math.round(Math.random() * 50)
const messageListener = (message: any) => {
// console.log('ChatGPTProvider:setupWebsocket:wsp.onMessage:', message)
const jjws = JSON.parse(message)
console.log('ChatGPTProvider:setupWSS:messageListener:jjws:', jjws)
const rawMessage = jjws['data'] ? jjws['data']['body'] : ''
console.log('ChatGPTProvider:setupWSS:wsp.onMessage:rawMessage:', rawMessage)
const b64decodedMessage = Buffer.from(rawMessage, 'base64')
const finalMessageStr = b64decodedMessage.toString()
console.log('ChatGPTProvider:setupWebsocket:wsp.onMessage:finalMessage:', finalMessageStr)

const parser = createParser((parent_message) => {
console.log('ChatGPTProvider:setupWSS:createParser:parent_message', parent_message) //event=`{data:'{}',event:undefine,id=undefined,type='event'}`
let data
try {
if ((parent_message['data' as keyof typeof parent_message] as string) === '[DONE]') {
console.log('ChatGPTProvider:setupWSS:createParser:returning DONE to frontend2')
params.onEvent({ type: 'done' })
wsp.close()
return
} else if (parent_message['data' as keyof typeof parent_message]) {
data = JSON.parse(parent_message['data' as keyof typeof parent_message])
console.log('ChatGPTProvider:setupWSS:createParser:data', data)
}
} catch (err) {
console.log('ChatGPTProvider:setupWSS:createParser:Error', err)
params.onEvent({ type: 'error', message: (err as any)?.message })
wsp.close()
return
}
const content = data?.message?.content as ResponseContent | undefined
if (!content) {
console.log('ChatGPTProvider:returning DONE to frontend3')
params.onEvent({ type: 'done' })
wsp.close()
return
}
let text: string
if (content.content_type === 'text') {
text = content.parts[0]
text = removeCitations(text)
} else if (content.content_type === 'code') {
text = '_' + content.text + '_'
} else {
console.log('ChatGPTProvider:returning DONE to frontend4')
params.onEvent({ type: 'done' })
wsp.close()
return
}
if (text) {
if (countWords(text) == 1 && data.message?.author?.role == 'assistant') {
if (params.prompt.indexOf('search query:') !== -1) {
this.renameConversationTitle(data.conversation_id)
}
}
params.onEvent({
type: 'answer',
data: {
text,
messageId: data.message.id,
parentMessageId: data.parent_message_id,
conversationId: data.conversation_id,
},
})
}
})
parser.feed(finalMessageStr)

const sequenceId = jjws['sequenceId']
console.log('ChatGPTProvider:doSendMessage:sequenceId:', sequenceId)
if (sequenceId === next_check_seqid) {
const t = {
type: 'sequenceAck',
sequenceId: next_check_seqid,
}
wsp.send(JSON.stringify(t))
next_check_seqid += Math.round(Math.random() * 50)
}
}
wsp.removeAllListeners()
wsp.close()
wsp.onOpen.addListener(openListener)
wsp.onMessage.addListener(messageListener)
wsp.onClose.removeListener(messageListener)
wsp.open().catch(async (e) => {
console.log('ChatGPTProvider:doSendMessage:showError:Error caught while opening ws', e)
wsp.removeAllListeners()
wsp.close()
await setChatgptwssIsOpenFlag(false)
params.onEvent({ type: 'error', message: (e as any)?.message })
})
}
}

async registerWSS(params: GenerateAnswerParams) {
const resp = await fetch('https://chat.openai.com/backend-api/register-websocket', {
method: 'POST',
signal: params.signal,
headers: {
'Content-Type': 'application/json',
Authorization: `Bearer ${this.token}`,
},
body: void 0,
})
return resp
}

async renameConversationTitle(convId: string) {
const titl: string = getConversationTitle(params.prompt)
console.log('renameConversationTitle:', this.token, convId, titl)
setConversationProperty(this.token, convId, { title: titl })
}

async generateAnswer(params: GenerateAnswerParams) {
console.log('ChatGPTProvider:generateAnswer', params.arkoseToken)
// let conversationId: string | undefined
const config = await getUserConfig()
const cleanup = () => {
// if (conversationId) {
// setConversationProperty(this.token, conversationId, { is_visible: false })
// }
}

if (config.chatgptMode == ChatgptMode.SSE) {
this.generateAnswerBySSE(params, cleanup)
} else {
const regResp = await this.registerWSS(params)
await this.setupWSS(params, regResp) // Since params change WSS have to be setup up every time
this.generateAnswerBySSE(params, cleanup)
}

return { cleanup }
}
}
8 changes: 7 additions & 1 deletion src/config.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,12 +39,18 @@ export interface SitePrompt {
prompt: string
}

export enum ChatgptMode {
SSE = 'Server-Sent-Events',
WSS = 'Web-Sockets',
}

const userConfigWithDefaultValue = {
triggerMode: TriggerMode.Always,
theme: Theme.Auto,
activeLanguage: "hi-IN",
activeLanguage: 'hi-IN',
prompt: Prompt,
promptOverrides: [] as SitePrompt[],
chatgptMode: ChatgptMode.SSE,
}

export type UserConfig = typeof userConfigWithDefaultValue
Expand Down
Loading

0 comments on commit 3872e80

Please sign in to comment.