-
Notifications
You must be signed in to change notification settings - Fork 989
/
index.ts
46 lines (43 loc) · 1.59 KB
/
index.ts
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
import WebSocketAsPromised from 'websocket-as-promised'
import { GradioBot } from '../gradio'
import { ChatError, ErrorCode } from '~utils/errors'
export class LMSYSBot extends GradioBot {
constructor(model: string) {
super('wss:https://chat.lmsys.org/queue/join', model, [0.7, 1, 512], 'text')
}
private async initializeSession(
fnIndex: number,
sessionHash: string,
data: unknown[],
signal?: AbortSignal,
): Promise<void> {
const wsp = new WebSocketAsPromised(this.wsUrl, {
packMessage: (data) => JSON.stringify(data),
unpackMessage: (data) => JSON.parse(data as string),
})
signal?.addEventListener('abort', () => wsp.close())
return new Promise((resolve, reject) => {
wsp.onUnpackedMessage.addListener((event) => {
if (event.msg === 'send_hash') {
wsp.sendPacked({ fn_index: fnIndex, session_hash: sessionHash })
} else if (event.msg === 'send_data') {
wsp.sendPacked({ fn_index: fnIndex, data, event_data: null, session_hash: sessionHash })
} else if (event.msg === 'process_completed') {
resolve()
}
})
wsp.open().catch((err) => {
console.error('lmsys ws open error', err)
reject(new ChatError('Failed to establish websocket connection.', ErrorCode.LMSYS_WS_ERROR))
})
})
}
async createSession(signal?: AbortSignal) {
const sessionHash = await super.createSession(signal)
await Promise.all([
this.initializeSession(36, sessionHash, [], signal),
this.initializeSession(43, sessionHash, [{}], signal),
])
return sessionHash
}
}