refactor to improve input performance

Author: arvinxx

packages/types/src/aiChat.ts

@@ -1,6 +1,7 @@
 import { z } from 'zod';
 
-import { MessageRoleType } from './message';
+import { ChatMessage, MessageRoleType } from './message';
+import { ChatTopic } from './topic';
 
 export interface SendNewMessage {
   content: string;
@@ -20,7 +21,6 @@ export interface SendMessageServerParams {
     title?: string;
     topicMessageIds?: string[];
   };
-  onlyAddUserMessage?: boolean;
 }
 
 export const AiSendMessageServerSchema = z.object({
@@ -39,6 +39,13 @@ export const AiSendMessageServerSchema = z.object({
       topicMessageIds: z.array(z.string()).optional(),
     })
     .optional(),
-  onlyAddUserMessage: z.boolean().optional(),
   topicId: z.string().optional(),
 });
+
+export interface SendMessageServerResponse {
+  isCreatNewTopic: boolean;
+  messageId: string;
+  messages: ChatMessage[];
+  topicId: string;
+  topics?: ChatTopic[];
+}

packages/types/src/index.ts

@@ -13,6 +13,7 @@ export * from './llm';
 export * from './message';
 export * from './meta';
 export * from './serverConfig';
+export * from './topic';
 export * from './user';
 export * from './user/settings';
 // FIXME: I think we need a refactor for the "openai" types

src/server/routers/lambda/aiChat.ts

@@ -1,10 +1,11 @@
+import { AiSendMessageServerSchema, SendMessageServerResponse } from '@lobechat/types';
+
 import { MessageModel } from '@/database/models/message';
 import { TopicModel } from '@/database/models/topic';
 import { authedProcedure, router } from '@/libs/trpc/lambda';
 import { serverDatabase } from '@/libs/trpc/lambda/middleware';
 import { AiChatService } from '@/server/services/aiChat';
 import { FileService } from '@/server/services/file';
-import { AiSendMessageServerSchema } from '@/types/aiChat';
 
 const aiChatProcedure = authedProcedure.use(serverDatabase).use(async (opts) => {
   const { ctx } = opts;
@@ -24,10 +25,11 @@ export const aiChatRouter = router({
     .input(AiSendMessageServerSchema)
     .mutation(async ({ input, ctx }) => {
       let messageId: string;
-      let topicId = input.topicId!;
+      let topicId = input.newMessage.topicId!;
 
       let isCreatNewTopic = false;
 
+      // create topic if there should be a new topic
       if (input.newTopic) {
         const topicItem = await ctx.topicModel.create({
           messages: input.newTopic.topicMessageIds,
@@ -44,11 +46,19 @@ export const aiChatRouter = router({
       });
       messageId = messageItem.id;
 
+      // retrieve latest messages and topic with
       const { messages, topics } = await ctx.aiChatService.getMessagesAndTopics({
+        isCreatNewTopic,
         sessionId: input.newMessage.sessionId,
         topicId,
       });
 
-      return { isCreatNewTopic, messageId: messageId, messages, topicId, topics };
+      return {
+        isCreatNewTopic,
+        messageId: messageId,
+        messages,
+        topicId,
+        topics,
+      } as SendMessageServerResponse;
     }),
 });

src/server/services/aiChat/index.ts

@@ -19,6 +19,7 @@ export class AiChatService {
 
   async getMessagesAndTopics(params: {
     current?: number;
+    isCreatNewTopic?: boolean;
     pageSize?: number;
     sessionId?: string;
     topicId?: string;
@@ -27,7 +28,7 @@ export class AiChatService {
       this.messageModel.query(params, {
         postProcessUrl: (path) => this.fileService.getFullFileUrl(path),
       }),
-      this.topicModel.query({ sessionId: params.sessionId }),
+      params.isCreatNewTopic ? this.topicModel.query({ sessionId: params.sessionId }) : undefined,
     ]);
 
     return { messages, topics };

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

@@ -1,5 +1,6 @@
 /* eslint-disable sort-keys-fix/sort-keys-fix, typescript-sort-keys/interface */
 // Disable the auto sort key eslint rule to make the code more logic and readable
+import { ChatMessage, ChatTopic, CreateMessageParams, SendMessageParams } from '@lobechat/types';
 import { t } from 'i18next';
 import { StateCreator } from 'zustand/vanilla';
 
@@ -8,7 +9,6 @@ import { getAgentStoreState } from '@/store/agent';
 import { ChatStore } from '@/store/chat/store';
 import { messageMapKey } from '@/store/chat/utils/messageMapKey';
 import { getSessionStoreState } from '@/store/session';
-import { CreateMessageParams, SendMessageParams } from '@/types/message';
 import { cleanObject } from '@/utils/object';
 
 import { chatSelectors, topicSelectors } from '../../../selectors';
@@ -18,6 +18,12 @@ export interface AIGenerateV2Action {
    * Sends a new message to the AI chat system
    */
   sendMessageInServer: (params: SendMessageParams) => Promise<void>;
+  internal_refreshAiChat: (params: {
+    topics?: ChatTopic[];
+    messages: ChatMessage[];
+    sessionId: string;
+    topicId?: string;
+  }) => void;
 }
 
 export const generateAIChatV2: StateCreator<
@@ -55,6 +61,10 @@ export const generateAIChatV2: StateCreator<
     };
     const messages = chatSelectors.activeBaseChats(get());
 
+    // use optimistic update to avoid the slow waiting
+    const tempId = get().internal_createTmpMessage(newMessage);
+
+    get().internal_toggleMessageLoading(true, tempId);
     set({ isCreatingMessage: true }, false, 'creatingMessage/start');
 
     const data = await aiChatService.sendMessageInServer({
@@ -66,15 +76,21 @@ export const generateAIChatV2: StateCreator<
             title: t('defaultTitle', { ns: 'topic' }),
           }
         : undefined,
-      onlyAddUserMessage,
+    });
+
+    // refresh the total data
+    get().internal_refreshAiChat({
+      messages: data.messages,
+      topics: data.topics,
+      sessionId: activeId,
+      topicId: data.topicId,
     });
 
     if (!activeTopicId) {
       await get().switchTopic(data.topicId!, true);
     }
 
-    get().internal_dispatchTopic({ type: 'updateTopics', value: data.topics as any[] });
-    get().internal_dispatchMessage({ type: 'updateMessages', value: data.messages });
+    get().internal_toggleMessageLoading(false, tempId);
 
     //  update assistant update to make it rerank
     getSessionStoreState().triggerSessionUpdate(get().activeId);
@@ -117,4 +133,15 @@ export const generateAIChatV2: StateCreator<
 
     await Promise.all([summaryTitle(), addFilesToAgent()]);
   },
+
+  internal_refreshAiChat: ({ topics, messages, sessionId, topicId }) => {
+    set(
+      {
+        topicMaps: topics ? { ...get().topicMaps, [sessionId]: topics } : get().topicMaps,
+        messagesMap: { ...get().messagesMap, [messageMapKey(sessionId, topicId)]: messages },
+      },
+      false,
+      'refreshAiChat',
+    );
+  },
 });