Skip to content

Commit 818a914

Browse files
arvinxxarvinxx
committed
refactor to improve input performance
1 parent 059a514 commit 818a914

File tree

3 files changed

+76
-131
lines changed

3 files changed

+76
-131
lines changed

‎src/server/routers/lambda/aiChat.ts

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -24,30 +24,31 @@ export const aiChatRouter = router({
2424
.input(AiSendMessageServerSchema)
2525
.mutation(async ({ input, ctx }) => {
2626
let messageId: string;
27-
if (input.onlyAddUserMessage) {
28-
// TODO: 是否直接返回 refresh 后的 message list
29-
const item = await ctx.messageModel.create(input.newMessage as any);
30-
messageId = item.id;
31-
}
27+
let topicId = input.topicId!;
28+
29+
let isCreatNewTopic = false;
3230

33-
let topicId = input.newMessage.topicId;
3431
if (input.newTopic) {
3532
const topicItem = await ctx.topicModel.create({
3633
messages: input.newTopic.topicMessageIds,
3734
sessionId: input.newTopic.sessionId,
3835
title: input.newTopic.title,
3936
});
4037
topicId = topicItem.id;
38+
isCreatNewTopic = true;
4139
}
4240

43-
const messageItem = await ctx.messageModel.create({ ...(input.newMessage as any), topicId });
41+
const messageItem = await ctx.messageModel.create({
42+
...(input.newMessage as any),
43+
topicId,
44+
});
4445
messageId = messageItem.id;
4546

4647
const { messages, topics } = await ctx.aiChatService.getMessagesAndTopics({
4748
sessionId: input.newMessage.sessionId,
4849
topicId,
4950
});
5051

51-
return { messageId: messageId, messages, topicId, topics };
52+
return { isCreatNewTopic, messageId: messageId, messages, topicId, topics };
5253
}),
5354
});

‎src/store/chat/slices/aiChat/actions/generateAIChatV2.ts

Lines changed: 51 additions & 123 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,14 @@ import { t } from 'i18next';
44
import { StateCreator } from 'zustand/vanilla';
55

66
import { aiChatService } from '@/services/aiChat';
7-
import { agentChatConfigSelectors } from '@/store/agent/selectors';
8-
import { getAgentStoreState } from '@/store/agent/store';
7+
import { getAgentStoreState } from '@/store/agent';
98
import { ChatStore } from '@/store/chat/store';
9+
import { messageMapKey } from '@/store/chat/utils/messageMapKey';
10+
import { getSessionStoreState } from '@/store/session';
1011
import { CreateMessageParams, SendMessageParams } from '@/types/message';
1112
import { cleanObject } from '@/utils/object';
1213

13-
import { chatSelectors } from '../../../selectors';
14+
import { chatSelectors, topicSelectors } from '../../../selectors';
1415

1516
export interface AIGenerateV2Action {
1617
/**
@@ -25,8 +26,8 @@ export const generateAIChatV2: StateCreator<
2526
[],
2627
AIGenerateV2Action
2728
> = (set, get) => ({
28-
sendMessageInServer: async ({ message, files, onlyAddUserMessage }) => {
29-
const { activeTopicId, activeId, activeThreadId } = get();
29+
sendMessageInServer: async ({ message, files, onlyAddUserMessage, isWelcomeQuestion }) => {
30+
const { activeTopicId, activeId, activeThreadId, internal_coreProcessMessage } = get();
3031
if (!activeId) return;
3132

3233
const fileIdList = files?.map((f) => f.id);
@@ -36,7 +37,11 @@ export const generateAIChatV2: StateCreator<
3637
// if message is empty or no files, then stop
3738
if (!message && !hasFile) return;
3839

39-
// set({ isCreatingMessage: true }, false, n('creatingMessage/start'));
40+
if (onlyAddUserMessage) {
41+
await get().addUserMessage({ message, fileList: fileIdList });
42+
43+
return;
44+
}
4045

4146
const newMessage: CreateMessageParams = {
4247
content: message,
@@ -50,143 +55,66 @@ export const generateAIChatV2: StateCreator<
5055
};
5156
const messages = chatSelectors.activeBaseChats(get());
5257

53-
const agentConfig = agentChatConfigSelectors.currentChatConfig(getAgentStoreState());
58+
set({ isCreatingMessage: true }, false, 'creatingMessage/start');
5459

5560
const data = await aiChatService.sendMessageInServer({
5661
newMessage: cleanObject(newMessage) as any,
57-
newTopic:
58-
!activeTopicId && agentConfig.enableAutoCreateTopic
59-
? {
60-
topicMessageIds: messages.map((m) => m.id),
61-
sessionId: activeId,
62-
title: t('defaultTitle', { ns: 'topic' }),
63-
}
64-
: undefined,
62+
newTopic: !activeTopicId
63+
? {
64+
topicMessageIds: messages.map((m) => m.id),
65+
sessionId: activeId,
66+
title: t('defaultTitle', { ns: 'topic' }),
67+
}
68+
: undefined,
6569
onlyAddUserMessage,
6670
});
6771

68-
console.log('data:', data);
69-
70-
if (agentConfig.enableAutoCreateTopic) {
72+
if (!activeTopicId) {
7173
await get().switchTopic(data.topicId!, true);
7274
}
7375

7476
get().internal_dispatchTopic({ type: 'updateTopics', value: data.topics as any[] });
7577
get().internal_dispatchMessage({ type: 'updateMessages', value: data.messages });
7678

77-
return;
79+
// update assistant update to make it rerank
80+
getSessionStoreState().triggerSessionUpdate(get().activeId);
7881

79-
// let tempMessageId: string | undefined = undefined;
80-
// let newTopicId: string | undefined = undefined;
81-
//
82-
// // it should be the default topic, then
83-
// // if autoCreateTopic is enabled, check to whether we need to create a topic
84-
// if (!onlyAddUserMessage && !activeTopicId && agentConfig.enableAutoCreateTopic) {
85-
// // check activeTopic and then auto create topic
86-
// const chats = chatSelectors.activeBaseChats(get());
87-
//
88-
// // we will add two messages (user and assistant), so the finial length should +2
89-
// const featureLength = chats.length + 2;
90-
//
91-
// // if there is no activeTopicId and the feature length is greater than the threshold
92-
// // then create a new topic and active it
93-
// if (!activeTopicId && featureLength >= agentConfig.autoCreateTopicThreshold) {
94-
// // we need to create a temp message for optimistic update
95-
// tempMessageId = get().internal_createTmpMessage(newMessage);
96-
// get().internal_toggleMessageLoading(true, tempMessageId);
97-
//
98-
// const topicId = await get().createTopic();
99-
//
100-
// if (topicId) {
101-
// newTopicId = topicId;
102-
// newMessage.topicId = topicId;
103-
//
104-
// // we need to copy the messages to the new topic or the message will disappear
105-
// const mapKey = chatSelectors.currentChatKey(get());
106-
// const newMaps = {
107-
// ...get().messagesMap,
108-
// [messageMapKey(activeId, topicId)]: get().messagesMap[mapKey],
109-
// };
110-
// set({ messagesMap: newMaps }, false, n('moveMessagesToNewTopic'));
111-
//
112-
// // make the topic loading
113-
// get().internal_updateTopicLoading(topicId, true);
114-
// }
115-
// }
116-
// }
117-
// // update assistant update to make it rerank
118-
// useSessionStore.getState().triggerSessionUpdate(get().activeId);
119-
//
120-
// const id = await get().internal_createMessage(newMessage, {
121-
// tempMessageId,
122-
// skipRefresh: !onlyAddUserMessage && newMessage.fileList?.length === 0,
123-
// });
124-
//
125-
// if (!id) {
126-
// set({ isCreatingMessage: false }, false, n('creatingMessage/start'));
127-
// if (!!newTopicId) get().internal_updateTopicLoading(newTopicId, false);
128-
// return;
129-
// }
130-
//
131-
// if (tempMessageId) get().internal_toggleMessageLoading(false, tempMessageId);
132-
//
133-
// // switch to the new topic if create the new topic
134-
// if (!!newTopicId) {
135-
// await get().switchTopic(newTopicId, true);
136-
// await get().internal_fetchMessages();
137-
//
138-
// // delete previous messages
139-
// // remove the temp message map
140-
// const newMaps = { ...get().messagesMap, [messageMapKey(activeId, null)]: [] };
141-
// set({ messagesMap: newMaps }, false, 'internal_copyMessages');
142-
// }
143-
//
144-
// // if only add user message, then stop
145-
// if (onlyAddUserMessage) {
146-
// set({ isCreatingMessage: false }, false, 'creatingMessage/start');
147-
// return;
148-
// }
14982
//
15083
// // Get the current messages to generate AI response
151-
// const messages = chatSelectors.activeBaseChats(get());
152-
// const userFiles = chatSelectors.currentUserFiles(get()).map((f) => f.id);
153-
//
154-
// await internal_coreProcessMessage(messages, id, {
155-
// isWelcomeQuestion,
156-
// ragQuery: get().internal_shouldUseRAG() ? message : undefined,
157-
// threadId: activeThreadId,
158-
// });
159-
//
160-
// set({ isCreatingMessage: false }, false, n('creatingMessage/stop'));
161-
//
162-
// const summaryTitle = async () => {
163-
// // if autoCreateTopic is false, then stop
164-
// if (!agentConfig.enableAutoCreateTopic) return;
84+
const baseMessages = chatSelectors.activeBaseChats(get());
16585
//
166-
// // check activeTopic and then auto update topic title
167-
// if (newTopicId) {
168-
// const chats = chatSelectors.getBaseChatsByKey(messageMapKey(activeId, newTopicId))(get());
169-
// await get().summaryTopicTitle(newTopicId, chats);
170-
// return;
171-
// }
86+
await internal_coreProcessMessage(baseMessages, data.messageId, {
87+
isWelcomeQuestion,
88+
ragQuery: get().internal_shouldUseRAG() ? message : undefined,
89+
threadId: activeThreadId,
90+
});
17291
//
173-
// if (!activeTopicId) return;
174-
// const topic = topicSelectors.getTopicById(activeTopicId)(get());
92+
set({ isCreatingMessage: false }, false, 'creatingMessage/stop');
17593
//
176-
// if (topic && !topic.title) {
177-
// const chats = chatSelectors.getBaseChatsByKey(messageMapKey(activeId, topic.id))(get());
178-
// await get().summaryTopicTitle(topic.id, chats);
179-
// }
180-
// };
94+
const summaryTitle = async () => {
95+
// check activeTopic and then auto update topic title
96+
if (data.isCreatNewTopic) {
97+
await get().summaryTopicTitle(data.topicId, data.messages);
98+
return;
99+
}
100+
101+
if (!activeTopicId) return;
102+
103+
const topic = topicSelectors.getTopicById(activeTopicId)(get());
104+
105+
if (topic && !topic.title) {
106+
const chats = chatSelectors.getBaseChatsByKey(messageMapKey(activeId, topic.id))(get());
107+
await get().summaryTopicTitle(topic.id, chats);
108+
}
109+
};
181110
//
182111
// // if there is relative files, then add files to agent
183112
// // only available in server mode
184-
// const addFilesToAgent = async () => {
185-
// if (userFiles.length === 0 || !isServerMode) return;
186-
//
187-
// await useAgentStore.getState().addFilesToAgent(userFiles, false);
188-
// };
189-
//
190-
// await Promise.all([summaryTitle(), addFilesToAgent()]);
113+
const userFiles = chatSelectors.currentUserFiles(get()).map((f) => f.id);
114+
const addFilesToAgent = async () => {
115+
await getAgentStoreState().addFilesToAgent(userFiles, false);
116+
};
117+
118+
await Promise.all([summaryTitle(), addFilesToAgent()]);
191119
},
192120
});

‎src/store/chat/slices/message/action.ts

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ const SWR_USE_FETCH_MESSAGES = 'SWR_USE_FETCH_MESSAGES';
3939
export interface ChatMessageAction {
4040
// create
4141
addAIMessage: () => Promise<void>;
42+
addUserMessage: (params: { message: string; fileList?: string[] }) => Promise<void>;
4243
// delete
4344
/**
4445
* clear message on the active session
@@ -213,6 +214,21 @@ export const chatMessage: StateCreator<
213214

214215
updateInputMessage('');
215216
},
217+
addUserMessage: async ({ message, fileList }) => {
218+
const { internal_createMessage, updateInputMessage, activeTopicId, activeId } = get();
219+
if (!activeId) return;
220+
221+
await internal_createMessage({
222+
content: message,
223+
files: fileList,
224+
role: 'user',
225+
sessionId: activeId,
226+
// if there is activeTopicId，then add topicId to message
227+
topicId: activeTopicId,
228+
});
229+
230+
updateInputMessage('');
231+
},
216232
copyMessage: async (id, content) => {
217233
await copyToClipboard(content);
218234

0 commit comments

Comments
 (0)