Skip to content

Commit c74c68a

Browse files
committed
refactor(MessageHandler -> useServerStream): convert all relating files to TS and correct typings based on this change: properly refactor MessageHandler to a custom hook, where it's passed a submission object to instantiate the stream. This is the bare minimum groundwork for potentially having multiple streams running, which would be a big project to modularize a lot of the global state into maps/multiple streams, particular useful for having multiple views in place
1 parent 8b4d3c2 commit c74c68a

File tree

5 files changed

+82
-86
lines changed

5 files changed

+82
-86
lines changed

client/src/hooks/useMessageHandler.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import { v4 } from 'uuid';
22
import { useRecoilState, useRecoilValue, useSetRecoilState } from 'recoil';
33
import { parseConvo, getResponseSender } from 'librechat-data-provider';
4-
import type { TMessage } from 'librechat-data-provider';
4+
import type { TMessage, TSubmission } from 'librechat-data-provider';
55
import store from '~/store';
66

77
type TAskProps = {
@@ -98,7 +98,7 @@ const useMessageHandler = () => {
9898
error: false,
9999
};
100100

101-
const submission = {
101+
const submission: TSubmission = {
102102
conversation: {
103103
...currentConversation,
104104
conversationId,

client/src/routes/MessageHandler.jsx renamed to client/src/hooks/useServerStream.ts

Lines changed: 37 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,29 @@
11
import { useEffect } from 'react';
2-
import { useRecoilValue, useResetRecoilState, useSetRecoilState } from 'recoil';
3-
import { SSE, createPayload } from 'librechat-data-provider';
2+
import { useResetRecoilState, useSetRecoilState } from 'recoil';
3+
import { SSE, createPayload, tMessageSchema, tConversationSchema } from 'librechat-data-provider';
4+
import type { TPlugin, TMessage, TConversation, TSubmission } from 'librechat-data-provider';
45
import { useAuthContext } from '~/hooks/AuthContext';
56
import store from '~/store';
67

7-
export default function MessageHandler() {
8-
const submission = useRecoilValue(store.submission);
9-
const setIsSubmitting = useSetRecoilState(store.isSubmitting);
8+
type TResData = {
9+
plugin: TPlugin;
10+
final?: boolean;
11+
initial?: boolean;
12+
requestMessage: TMessage;
13+
responseMessage: TMessage;
14+
conversation: TConversation;
15+
};
16+
17+
export default function useServerStream(submission: TSubmission | null) {
1018
const setMessages = useSetRecoilState(store.messages);
19+
const setIsSubmitting = useSetRecoilState(store.isSubmitting);
1120
const setConversation = useSetRecoilState(store.conversation);
1221
const resetLatestMessage = useResetRecoilState(store.latestMessage);
1322
const { token } = useAuthContext();
1423

1524
const { refreshConversations } = store.useConversations();
1625

17-
const messageHandler = (data, submission) => {
26+
const messageHandler = (data: string, submission: TSubmission) => {
1827
const {
1928
messages,
2029
message,
@@ -30,9 +39,9 @@ export default function MessageHandler() {
3039
{
3140
...initialResponse,
3241
text: data,
33-
parentMessageId: message?.overrideParentMessageId,
42+
parentMessageId: message?.overrideParentMessageId ?? null,
3443
messageId: message?.overrideParentMessageId + '_',
35-
plugin: plugin ? plugin : null,
44+
plugin: plugin ?? null,
3645
submitting: true,
3746
// unfinished: true
3847
},
@@ -46,15 +55,15 @@ export default function MessageHandler() {
4655
text: data,
4756
parentMessageId: message?.messageId,
4857
messageId: message?.messageId + '_',
49-
plugin: plugin ? plugin : null,
58+
plugin: plugin ?? null,
5059
submitting: true,
5160
// unfinished: true
5261
},
5362
]);
5463
}
5564
};
5665

57-
const cancelHandler = (data, submission) => {
66+
const cancelHandler = (data: TResData, submission: TSubmission) => {
5867
const { requestMessage, responseMessage, conversation } = data;
5968
const { messages, isRegenerate = false, isEdited = false } = submission;
6069

@@ -84,7 +93,7 @@ export default function MessageHandler() {
8493
}));
8594
};
8695

87-
const createdHandler = (data, submission) => {
96+
const createdHandler = (data: TResData, submission: TSubmission) => {
8897
const {
8998
messages,
9099
message,
@@ -98,7 +107,7 @@ export default function MessageHandler() {
98107
...(isEdited ? messages.slice(0, -1) : messages),
99108
{
100109
...initialResponse,
101-
parentMessageId: message?.overrideParentMessageId,
110+
parentMessageId: message?.overrideParentMessageId ?? null,
102111
messageId: message?.overrideParentMessageId + '_',
103112
submitting: true,
104113
},
@@ -117,14 +126,16 @@ export default function MessageHandler() {
117126
}
118127

119128
const { conversationId } = message;
120-
setConversation((prevState) => ({
121-
...prevState,
122-
conversationId,
123-
}));
129+
setConversation((prevState) =>
130+
tConversationSchema.parse({
131+
...prevState,
132+
conversationId,
133+
}),
134+
);
124135
resetLatestMessage();
125136
};
126137

127-
const finalHandler = (data, submission) => {
138+
const finalHandler = (data: TResData, submission: TSubmission) => {
128139
const { requestMessage, responseMessage, conversation } = data;
129140
const { messages, isRegenerate = false, isEdited = false } = submission;
130141

@@ -154,21 +165,21 @@ export default function MessageHandler() {
154165
}));
155166
};
156167

157-
const errorHandler = (data, submission) => {
168+
const errorHandler = (data: TResData, submission: TSubmission) => {
158169
const { messages, message } = submission;
159170

160171
console.log('Error:', data);
161-
const errorResponse = {
172+
const errorResponse = tMessageSchema.parse({
162173
...data,
163174
error: true,
164175
parentMessageId: message?.messageId,
165-
};
176+
});
166177
setIsSubmitting(false);
167178
setMessages([...messages, message, errorResponse]);
168179
return;
169180
};
170181

171-
const abortConversation = (conversationId) => {
182+
const abortConversation = (conversationId = '', submission: TSubmission) => {
172183
console.log(submission);
173184
const { endpoint } = submission?.conversation || {};
174185

@@ -212,7 +223,7 @@ export default function MessageHandler() {
212223
headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${token}` },
213224
});
214225

215-
events.onmessage = (e) => {
226+
events.onmessage = (e: MessageEvent) => {
216227
const data = JSON.parse(e.data);
217228

218229
if (data.final) {
@@ -227,8 +238,8 @@ export default function MessageHandler() {
227238
createdHandler(data, { ...submission, message });
228239
console.log('created', message);
229240
} else {
230-
let text = data.text || data.response;
231-
let { initial, plugin } = data;
241+
const text = data.text || data.response;
242+
const { initial, plugin } = data;
232243
if (initial) {
233244
console.log(data);
234245
}
@@ -242,9 +253,9 @@ export default function MessageHandler() {
242253
events.onopen = () => console.log('connection is opened');
243254

244255
events.oncancel = () =>
245-
abortConversation(message?.conversationId || submission?.conversationId);
256+
abortConversation(message?.conversationId ?? submission?.conversationId, submission);
246257

247-
events.onerror = function (e) {
258+
events.onerror = function (e: MessageEvent) {
248259
console.log('error in opening conn.');
249260
events.close();
250261

@@ -268,6 +279,4 @@ export default function MessageHandler() {
268279
};
269280
// eslint-disable-next-line react-hooks/exhaustive-deps
270281
}, [submission]);
271-
272-
return null;
273282
}

client/src/routes/Root.jsx renamed to client/src/routes/Root.tsx

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
/* eslint-disable react-hooks/exhaustive-deps */
22
import { useEffect, useState } from 'react';
3-
import { useSetRecoilState } from 'recoil';
3+
import { useRecoilValue, useSetRecoilState } from 'recoil';
44
import { Outlet } from 'react-router-dom';
55
import {
66
useGetEndpointsQuery,
@@ -9,8 +9,7 @@ import {
99
} from 'librechat-data-provider';
1010

1111
import { Nav, MobileNav } from '~/components/Nav';
12-
import { useAuthContext } from '~/hooks/AuthContext';
13-
import MessageHandler from './MessageHandler';
12+
import { useAuthContext, useServerStream } from '~/hooks';
1413
import store from '~/store';
1514

1615
export default function Root() {
@@ -19,6 +18,9 @@ export default function Root() {
1918
return savedNavVisible !== null ? JSON.parse(savedNavVisible) : false;
2019
});
2120

21+
const submission = useRecoilValue(store.submission);
22+
useServerStream(submission ?? null);
23+
2224
const setIsSearchEnabled = useSetRecoilState(store.isSearchEnabled);
2325
const setEndpointsConfig = useSetRecoilState(store.endpointsConfig);
2426
const setPresets = useSetRecoilState(store.presets);
@@ -71,7 +73,6 @@ export default function Root() {
7173
</div>
7274
</div>
7375
</div>
74-
<MessageHandler />
7576
</>
7677
);
7778
}

packages/data-provider/src/schemas.ts

Lines changed: 30 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -12,32 +12,6 @@ export enum EModelEndpoint {
1212

1313
export const eModelEndpointSchema = z.nativeEnum(EModelEndpoint);
1414

15-
export const tMessageSchema = z.object({
16-
messageId: z.string(),
17-
clientId: z.string().nullable().optional(),
18-
conversationId: z.string().nullable(),
19-
parentMessageId: z.string().nullable(),
20-
sender: z.string(),
21-
text: z.string(),
22-
isCreatedByUser: z.boolean(),
23-
error: z.boolean(),
24-
createdAt: z
25-
.string()
26-
.optional()
27-
.default(() => new Date().toISOString()),
28-
updatedAt: z
29-
.string()
30-
.optional()
31-
.default(() => new Date().toISOString()),
32-
current: z.boolean().optional(),
33-
unfinished: z.boolean().optional(),
34-
submitting: z.boolean().optional(),
35-
searchResult: z.boolean().optional(),
36-
finish_reason: z.string().optional(),
37-
});
38-
39-
export type TMessage = z.input<typeof tMessageSchema>;
40-
4115
export const tPluginAuthConfigSchema = z.object({
4216
authField: z.string(),
4317
label: z.string(),
@@ -76,6 +50,36 @@ export const tAgentOptionsSchema = z.object({
7650
temperature: z.number(),
7751
});
7852

53+
export const tMessageSchema = z.object({
54+
messageId: z.string(),
55+
clientId: z.string().nullable().optional(),
56+
conversationId: z.string().nullable(),
57+
parentMessageId: z.string().nullable(),
58+
responseMessageId: z.string().nullable().optional(),
59+
overrideParentMessageId: z.string().nullable().optional(),
60+
plugin: tPluginSchema.nullable().optional(),
61+
sender: z.string(),
62+
text: z.string(),
63+
generation: z.string().nullable().optional(),
64+
isCreatedByUser: z.boolean(),
65+
error: z.boolean(),
66+
createdAt: z
67+
.string()
68+
.optional()
69+
.default(() => new Date().toISOString()),
70+
updatedAt: z
71+
.string()
72+
.optional()
73+
.default(() => new Date().toISOString()),
74+
current: z.boolean().optional(),
75+
unfinished: z.boolean().optional(),
76+
submitting: z.boolean().optional(),
77+
searchResult: z.boolean().optional(),
78+
finish_reason: z.string().optional(),
79+
});
80+
81+
export type TMessage = z.input<typeof tMessageSchema>;
82+
7983
export const tConversationSchema = z.object({
8084
conversationId: z.string().nullable(),
8185
title: z.string(),

packages/data-provider/src/types.ts

Lines changed: 8 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import type { TMessage, EModelEndpoint, TConversation, TEndpointOption } from './schemas';
1+
import type { TPlugin, TMessage, TConversation, TEndpointOption } from './schemas';
22

33
export * from './schemas';
44

@@ -7,32 +7,14 @@ export type TMessages = TMessage[];
77
export type TMessagesAtom = TMessages | null;
88

99
export type TSubmission = {
10-
clientId?: string;
11-
context?: string;
12-
conversationId?: string;
13-
conversationSignature?: string;
14-
current: boolean;
15-
endpoint: EModelEndpoint | null;
16-
invocationId: number;
17-
isCreatedByUser: boolean;
18-
jailbreak: boolean;
19-
jailbreakConversationId?: string;
20-
messageId: string;
21-
overrideParentMessageId?: string | boolean;
22-
parentMessageId?: string;
23-
sender: string;
24-
systemMessage?: string;
25-
text: string;
26-
toneStyle?: string;
27-
model?: string;
28-
promptPrefix?: string;
29-
temperature?: number;
30-
top_p?: number;
31-
presence_penalty?: number;
32-
frequence_penalty?: number;
10+
plugin?: TPlugin;
11+
message: TMessage;
3312
isEdited?: boolean;
13+
messages: TMessage[];
14+
isRegenerate?: boolean;
15+
conversationId?: string;
16+
initialResponse: TMessage;
3417
conversation: TConversation;
35-
message: TMessage;
3618
endpointOption: TEndpointOption;
3719
};
3820

@@ -167,7 +149,7 @@ export type TResetPassword = {
167149
};
168150

169151
export type TStartupConfig = {
170-
appTitle: boolean;
152+
appTitle: string;
171153
googleLoginEnabled: boolean;
172154
openidLoginEnabled: boolean;
173155
githubLoginEnabled: boolean;

0 commit comments

Comments
 (0)