-
Notifications
You must be signed in to change notification settings - Fork 304
/
usePostMessageStreaming.ts
72 lines (62 loc) · 1.83 KB
/
usePostMessageStreaming.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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
import { Auth } from 'aws-amplify';
import { PostMessageRequest } from '../@types/conversation';
import { create } from 'zustand';
import i18next from 'i18next';
const WS_ENDPOINT: string = import.meta.env.VITE_APP_WS_ENDPOINT;
const usePostMessageStreaming = create<{
post: (
input: PostMessageRequest,
dispatch: (completion: string) => void
) => Promise<void>;
}>(() => {
const post = async (
input: PostMessageRequest,
dispatch: (completion: string) => void
) => {
const token = (await Auth.currentSession()).getIdToken().getJwtToken();
return new Promise<void>((resolve, reject) => {
const ws = new WebSocket(WS_ENDPOINT);
let completion = '';
ws.onopen = () => {
ws.send(JSON.stringify({ ...input, token }));
};
ws.onmessage = (message) => {
try {
if (message.data === '') {
return;
}
const data = JSON.parse(message.data);
if (data.completion || data.completion === '') {
if (completion.endsWith('▍')) {
completion = completion.slice(0, -1);
}
completion += data.completion + (data.stop_reason ? '' : '▍');
dispatch(completion);
if (data.stop_reason) {
ws.close();
}
} else {
ws.close();
console.error(data);
throw new Error(i18next.t('error.predict.invalidResponse'));
}
} catch (e) {
console.error(e);
reject(i18next.t('error.predict.general'));
}
};
ws.onerror = (e) => {
ws.close();
console.error(e);
reject(i18next.t('error.predict.general'));
};
ws.onclose = () => {
resolve();
};
});
};
return {
post,
};
});
export default usePostMessageStreaming;