@lobehub/chat 1.121.0 → 1.122.0

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

@@ -0,0 +1,80 @@
+import { AiSendMessageServerSchema, SendMessageServerResponse } from '@lobechat/types';
+
+import { LOADING_FLAT } from '@/const/message';
+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';
+
+const aiChatProcedure = authedProcedure.use(serverDatabase).use(async (opts) => {
+  const { ctx } = opts;
+
+  return opts.next({
+    ctx: {
+      aiChatService: new AiChatService(ctx.serverDB, ctx.userId),
+      fileService: new FileService(ctx.serverDB, ctx.userId),
+      messageModel: new MessageModel(ctx.serverDB, ctx.userId),
+      topicModel: new TopicModel(ctx.serverDB, ctx.userId),
+    },
+  });
+});
+
+export const aiChatRouter = router({
+  sendMessageInServer: aiChatProcedure
+    .input(AiSendMessageServerSchema)
+    .mutation(async ({ input, ctx }) => {
+      let messageId: string;
+      let topicId = input.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,
+          sessionId: input.sessionId,
+          title: input.newTopic.title,
+        });
+        topicId = topicItem.id;
+        isCreatNewTopic = true;
+      }
+
+      // create user message
+      const userMessageItem = await ctx.messageModel.create({
+        content: input.newUserMessage.content,
+        files: input.newUserMessage.files,
+        role: 'user',
+        sessionId: input.sessionId!,
+        topicId,
+      });
+
+      messageId = userMessageItem.id;
+      // create assistant message
+      const assistantMessageItem = await ctx.messageModel.create({
+        content: LOADING_FLAT,
+        fromModel: input.newAssistantMessage.model,
+        parentId: messageId,
+        role: 'assistant',
+        sessionId: input.sessionId!,
+        topicId,
+      });
+
+      // retrieve latest messages and topic with
+      const { messages, topics } = await ctx.aiChatService.getMessagesAndTopics({
+        includeTopic: isCreatNewTopic,
+        sessionId: input.sessionId,
+        topicId,
+      });
+
+      return {
+        assistantMessageId: assistantMessageItem.id,
+        isCreatNewTopic,
+        messages,
+        topicId,
+        topics,
+        userMessageId: messageId,
+      } as SendMessageServerResponse;
+    }),
+});

package/src/server/routers/lambda/index.ts

@@ -4,6 +4,7 @@
 import { publicProcedure, router } from '@/libs/trpc/lambda';
 
 import { agentRouter } from './agent';
+import { aiChatRouter } from './aiChat';
 import { aiModelRouter } from './aiModel';
 import { aiProviderRouter } from './aiProvider';
 import { apiKeyRouter } from './apiKey';
@@ -30,6 +31,7 @@ import { userRouter } from './user';
 
 export const lambdaRouter = router({
   agent: agentRouter,
+  aiChat: aiChatRouter,
   aiModel: aiModelRouter,
   aiProvider: aiProviderRouter,
   apiKey: apiKeyRouter,

package/src/server/routers/lambda/message.ts

@@ -7,6 +7,7 @@ import { authedProcedure, publicProcedure, router } from '@/libs/trpc/lambda';
 import { serverDatabase } from '@/libs/trpc/lambda/middleware';
 import { FileService } from '@/server/services/file';
 import { ChatMessage } from '@/types/message';
+import { UpdateMessageRAGParamsSchema } from '@/types/message/rag';
 import { BatchTaskResult } from '@/types/service';
 
 type ChatMessageList = ChatMessage[];
@@ -174,6 +175,12 @@ export const messageRouter = router({
       return ctx.messageModel.updateMessagePlugin(input.id, input.value);
     }),
 
+  updateMessageRAG: messageProcedure
+    .input(UpdateMessageRAGParamsSchema)
+    .mutation(async ({ input, ctx }) => {
+      await ctx.messageModel.updateMessageRAG(input.id, input.value);
+    }),
+
   updatePluginError: messageProcedure
     .input(
       z.object({

package/src/server/services/aiChat/index.test.ts

@@ -0,0 +1,57 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import { MessageModel } from '@/database/models/message';
+import { TopicModel } from '@/database/models/topic';
+import { LobeChatDatabase } from '@/database/type';
+import { FileService } from '@/server/services/file';
+
+import { AiChatService } from '.';
+
+vi.mock('@/database/models/message');
+vi.mock('@/database/models/topic');
+vi.mock('@/server/services/file');
+
+describe('AiChatService', () => {
+  it('getMessagesAndTopics should fetch messages and topics concurrently', async () => {
+    const serverDB = {} as unknown as LobeChatDatabase;
+
+    const mockQueryMessages = vi.fn().mockResolvedValue([{ id: 'm1' }]);
+    const mockQueryTopics = vi.fn().mockResolvedValue([{ id: 't1' }]);
+
+    vi.mocked(MessageModel).mockImplementation(() => ({ query: mockQueryMessages }) as any);
+    vi.mocked(TopicModel).mockImplementation(() => ({ query: mockQueryTopics }) as any);
+    vi.mocked(FileService).mockImplementation(
+      () => ({ getFullFileUrl: vi.fn().mockResolvedValue('url') }) as any,
+    );
+
+    const service = new AiChatService(serverDB, 'u1');
+
+    const res = await service.getMessagesAndTopics({ includeTopic: true, sessionId: 's1' });
+
+    expect(mockQueryMessages).toHaveBeenCalledWith(
+      { includeTopic: true, sessionId: 's1' },
+      expect.objectContaining({ postProcessUrl: expect.any(Function) }),
+    );
+    expect(mockQueryTopics).toHaveBeenCalledWith({ sessionId: 's1' });
+    expect(res.messages).toEqual([{ id: 'm1' }]);
+    expect(res.topics).toEqual([{ id: 't1' }]);
+  });
+
+  it('getMessagesAndTopics should not query topics when includeTopic is false', async () => {
+    const serverDB = {} as unknown as LobeChatDatabase;
+
+    const mockQueryMessages = vi.fn().mockResolvedValue([]);
+    vi.mocked(MessageModel).mockImplementation(() => ({ query: mockQueryMessages }) as any);
+    vi.mocked(TopicModel).mockImplementation(() => ({ query: vi.fn() }) as any);
+    vi.mocked(FileService).mockImplementation(
+      () => ({ getFullFileUrl: vi.fn().mockResolvedValue('url') }) as any,
+    );
+
+    const service = new AiChatService(serverDB, 'u1');
+
+    const res = await service.getMessagesAndTopics({ includeTopic: false, topicId: 't1' });
+
+    expect(mockQueryMessages).toHaveBeenCalled();
+    expect(res.topics).toBeUndefined();
+  });
+});

package/src/server/services/aiChat/index.ts

@@ -0,0 +1,36 @@
+import { MessageModel } from '@/database/models/message';
+import { TopicModel } from '@/database/models/topic';
+import { LobeChatDatabase } from '@/database/type';
+import { FileService } from '@/server/services/file';
+
+export class AiChatService {
+  private userId: string;
+  private messageModel: MessageModel;
+  private fileService: FileService;
+  private topicModel: TopicModel;
+
+  constructor(serverDB: LobeChatDatabase, userId: string) {
+    this.userId = userId;
+
+    this.messageModel = new MessageModel(serverDB, userId);
+    this.topicModel = new TopicModel(serverDB, userId);
+    this.fileService = new FileService(serverDB, userId);
+  }
+
+  async getMessagesAndTopics(params: {
+    current?: number;
+    includeTopic?: boolean;
+    pageSize?: number;
+    sessionId?: string;
+    topicId?: string;
+  }) {
+    const [messages, topics] = await Promise.all([
+      this.messageModel.query(params, {
+        postProcessUrl: (path) => this.fileService.getFullFileUrl(path),
+      }),
+      params.includeTopic ? this.topicModel.query({ sessionId: params.sessionId }) : undefined,
+    ]);
+
+    return { messages, topics };
+  }
+}

package/src/services/aiChat.ts

@@ -0,0 +1,12 @@
+import { SendMessageServerParams } from '@lobechat/types';
+import { cleanObject } from '@lobechat/utils';
+
+import { lambdaClient } from '@/libs/trpc/client';
+
+class AiChatService {
+  sendMessageInServer = async (params: SendMessageServerParams) => {
+    return lambdaClient.aiChat.sendMessageInServer.mutate(cleanObject(params));
+  };
+}
+
+export const aiChatService = new AiChatService();

package/src/services/message/_deprecated.ts

@@ -143,4 +143,8 @@ export class ClientService implements IMessageService {
   async updateMessagePluginError() {
     throw new Error('Method not implemented.');
   }
+
+  async updateMessageRAG(): Promise<void> {
+    throw new Error('Method not implemented.');
+  }
 }

package/src/services/message/client.ts

@@ -96,6 +96,11 @@ export class ClientService extends BaseClientService implements IMessageService
     return this.messageModel.updateMessagePlugin(id, { error: value });
   };
 
+  updateMessageRAG: IMessageService['updateMessageRAG'] = async (id, value) => {
+    console.log(id, value);
+    throw new Error('not implemented');
+  };
+
   updateMessagePluginArguments: IMessageService['updateMessagePluginArguments'] = async (
     id,
     value,

package/src/services/message/server.ts

@@ -84,6 +84,10 @@ export class ServerService implements IMessageService {
     return lambdaClient.message.updatePluginError.mutate({ id, value: error as any });
   };
 
+  updateMessageRAG: IMessageService['updateMessageRAG'] = async (id, data) => {
+    return lambdaClient.message.updateMessageRAG.mutate({ id, value: data });
+  };
+
   removeMessage: IMessageService['removeMessage'] = async (id) => {
     return lambdaClient.message.removeMessage.mutate({ id });
   };

package/src/services/message/type.ts

@@ -11,6 +11,7 @@ import {
   ModelRankItem,
   UpdateMessageParams,
 } from '@/types/message';
+import { UpdateMessageRAGParams } from '@/types/message/rag';
 
 /* eslint-disable typescript-sort-keys/interface */
 
@@ -39,6 +40,7 @@ export interface IMessageService {
   updateMessageTranslate(id: string, translate: Partial<ChatTranslate> | false): Promise<any>;
   updateMessagePluginState(id: string, value: Record<string, any>): Promise<any>;
   updateMessagePluginError(id: string, value: ChatMessagePluginError | null): Promise<any>;
+  updateMessageRAG(id: string, value: UpdateMessageRAGParams): Promise<void>;
   updateMessagePluginArguments(id: string, value: string | Record<string, any>): Promise<any>;
   removeMessage(id: string): Promise<any>;
   removeMessages(ids: string[]): Promise<any>;

package/src/store/chat/slices/aiChat/actions/generateAIChat.ts

@@ -152,7 +152,13 @@ export const generateAIChat: StateCreator<
   },
 
   sendMessage: async ({ message, files, onlyAddUserMessage, isWelcomeQuestion }) => {
-    const { internal_coreProcessMessage, activeTopicId, activeId, activeThreadId } = get();
+    const {
+      internal_coreProcessMessage,
+      activeTopicId,
+      activeId,
+      activeThreadId,
+      sendMessageInServer,
+    } = get();
     if (!activeId) return;
 
     const fileIdList = files?.map((f) => f.id);
@@ -162,6 +168,10 @@ export const generateAIChat: StateCreator<
     // if message is empty or no files, then stop
     if (!message && !hasFile) return;
 
+    // router to server mode send message
+    if (isServerMode)
+      return sendMessageInServer({ message, files, onlyAddUserMessage, isWelcomeQuestion });
+
     set({ isCreatingMessage: true }, false, n('creatingMessage/start'));
 
     const newMessage: CreateMessageParams = {