@lobehub/chat 1.36.46 → 1.37.1

package/src/services/session/pglite.test.ts

@@ -0,0 +1,411 @@
+import { eq, not } from 'drizzle-orm/expressions';
+import { Mock, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { INBOX_SESSION_ID } from '@/const/session';
+import { clientDB, initializeDB } from '@/database/client/db';
+import {
+  NewSession,
+  SessionItem,
+  agents,
+  agentsToSessions,
+  sessionGroups,
+  sessions,
+  users,
+} from '@/database/schemas';
+import { LobeAgentChatConfig, LobeAgentConfig } from '@/types/agent';
+import { LobeAgentSession, LobeSessionType, SessionGroups } from '@/types/session';
+
+import { ClientService } from './pglite';
+
+const userId = 'message-db';
+const sessionService = new ClientService(userId);
+
+const mockSessionId = 'mock-session-id';
+const mockAgentId = 'agent-id';
+
+// Mock data
+beforeEach(async () => {
+  await initializeDB();
+
+  // 在每个测试用例之前，清空表
+  await clientDB.transaction(async (trx) => {
+    await trx.insert(users).values([{ id: userId }, { id: '456' }]);
+    await trx.insert(sessions).values([{ id: mockSessionId, userId }]);
+    await trx.insert(agents).values([{ id: mockAgentId, userId }]);
+    await trx.insert(agentsToSessions).values([{ agentId: mockAgentId, sessionId: mockSessionId }]);
+    await trx.insert(sessionGroups).values([
+      { id: 'group-1', name: 'group-A', sort: 2, userId },
+      { id: 'group-2', name: 'group-B', sort: 1, userId },
+      { id: 'group-4', name: 'group-C', sort: 1, userId: '456' },
+    ]);
+  });
+});
+
+afterEach(async () => {
+  // 在每个测试用例之后，清空表
+  await clientDB.delete(users);
+});
+
+describe('SessionService', () => {
+  const mockSession = {
+    id: mockSessionId,
+    type: 'agent',
+    meta: { title: 'Mock Session' },
+  } as LobeAgentSession;
+
+  describe('createSession', () => {
+    it('should create a new session and return its id', async () => {
+      // Setup
+      const sessionType = LobeSessionType.Agent;
+      const defaultValue = { meta: { title: 'New Session' } } as Partial<LobeAgentSession>;
+
+      // Execute
+      const sessionId = await sessionService.createSession(sessionType, defaultValue);
+
+      // Assert
+      expect(sessionId).toMatch(/^ssn_/);
+    });
+  });
+
+  describe('removeSession', () => {
+    it('should remove a session by its id', async () => {
+      // Execute
+      await sessionService.removeSession(mockSessionId);
+
+      // Assert
+
+      const result = await clientDB.query.sessions.findFirst({
+        where: eq(sessions.id, mockSessionId),
+      });
+      // Assert
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe('removeAllSessions', () => {
+    it('should clear all sessions from the table', async () => {
+      // Setup
+      await clientDB
+        .insert(sessions)
+        .values([{ userId: userId }, { userId: userId }, { userId: userId }]);
+
+      // Execute
+      await sessionService.removeAllSessions();
+
+      // Assert
+      const result = await clientDB.query.sessions.findMany({
+        where: eq(sessionGroups.userId, userId),
+      });
+
+      expect(result.length).toBe(0);
+    });
+  });
+
+  describe('updateSession', () => {
+    it('should update the group of a session', async () => {
+      // Setup
+      const groupId = 'group-1';
+
+      // Execute
+      await sessionService.updateSession(mockSessionId, { group: groupId });
+
+      // Assert
+      const result = await clientDB.query.sessions.findFirst({
+        where: eq(sessions.id, mockSessionId),
+      });
+      expect(result).toMatchObject({ groupId });
+    });
+
+    it('should update the pinned status of a session', async () => {
+      // Setup
+      const pinned = true;
+
+      // Execute
+      await sessionService.updateSession(mockSessionId, { pinned });
+
+      // Assert
+      const result = await clientDB.query.sessions.findFirst({
+        where: eq(sessions.id, mockSessionId),
+      });
+
+      expect(result!.pinned).toBeTruthy();
+    });
+  });
+
+  describe('updateSessionConfig', () => {
+    it('should update the config of a session', async () => {
+      // Setup
+      const newConfig = { model: 'abc' } as LobeAgentConfig;
+
+      // Execute
+      await sessionService.updateSessionConfig(mockSessionId, newConfig);
+
+      // Assert
+      const result = await sessionService.getSessionConfig(mockSessionId);
+      expect(result).toMatchObject(newConfig);
+    });
+  });
+
+  describe('countSessions', () => {
+    it('should return false if no sessions exist', async () => {
+      await clientDB.delete(sessions);
+
+      // Execute
+      const result = await sessionService.countSessions();
+
+      // Assert
+      expect(result).toBe(0);
+    });
+
+    it('should return true if sessions exist', async () => {
+      // Setup
+      await clientDB.delete(sessions);
+      await clientDB.insert(sessions).values([{ userId }]);
+
+      // Execute
+      const result = await sessionService.countSessions();
+
+      // Assert
+      expect(result).toBe(1);
+    });
+  });
+
+  describe('searchSessions', () => {
+    it('should return sessions that match the keyword', async () => {
+      // Setup
+      await clientDB.insert(agents).values({ userId, id: 'agent-1', title: 'Session Name' });
+      await clientDB
+        .insert(agentsToSessions)
+        .values({ agentId: 'agent-1', sessionId: mockSessionId });
+
+      // Execute
+      const keyword = 'Name';
+      const result = await sessionService.searchSessions(keyword);
+
+      // Assert
+      // TODO: 后续需要把这个搜索的标题和描述都加上，现在这个 client 搜索会有问题
+      expect(result).toMatchObject([{ id: mockSessionId }]);
+    });
+  });
+
+  describe('cloneSession', () => {
+    it('should duplicate a session and return its id', async () => {
+      // Setup
+      const newTitle = 'Duplicated Session';
+      const session: NewSession = {
+        id: 'duplicated-session-id',
+        title: '123',
+        userId,
+      };
+      await clientDB.insert(sessions).values([session]);
+      await clientDB.insert(agents).values({ userId, id: 'agent-1' });
+      await clientDB
+        .insert(agentsToSessions)
+        .values({ agentId: 'agent-1', sessionId: 'duplicated-session-id' });
+
+      // Execute
+      const duplicatedSessionId = await sessionService.cloneSession(mockSessionId, newTitle);
+
+      // Assert
+
+      const result = await clientDB.query.sessions.findFirst({
+        where: eq(sessions.id, duplicatedSessionId!),
+      });
+      expect(result).toMatchObject({ title: 'Duplicated Session' });
+    });
+  });
+
+  describe('getGroupedSessions', () => {
+    it('should retrieve sessions with their group', async () => {
+      // Execute
+      const sessionsWithGroup = await sessionService.getGroupedSessions();
+
+      expect(sessionsWithGroup).toMatchObject({
+        sessionGroups: [
+          { id: 'group-2', name: 'group-B', sort: 1 },
+          { id: 'group-1', name: 'group-A', sort: 2 },
+        ],
+        sessions: [{ id: 'mock-session-id', type: 'agent' }],
+      });
+    });
+  });
+
+  describe('getSessionsByType', () => {
+    it('should get sessions by type "all"', async () => {
+      const sessions = await sessionService.getSessionsByType('all');
+      expect(sessions).toBeDefined();
+    });
+
+    it('should get sessions by type "agent"', async () => {
+      const sessions = await sessionService.getSessionsByType('agent');
+      expect(sessions).toBeDefined();
+    });
+
+    it('should get sessions by type "group"', async () => {
+      const sessions = await sessionService.getSessionsByType('group');
+      expect(sessions).toBeDefined();
+    });
+  });
+
+  describe('getSessionConfig', () => {
+    it.skip('should get default config for INBOX_SESSION_ID', async () => {
+      const config = await sessionService.getSessionConfig(INBOX_SESSION_ID);
+      expect(config).toBeDefined();
+    });
+
+    it('should throw error for non-existent session', async () => {
+      await expect(sessionService.getSessionConfig('non-existent')).rejects.toThrow(
+        'Session not found',
+      );
+    });
+  });
+
+  describe('updateSessionMeta', () => {
+    it('should not update meta for INBOX_SESSION_ID', async () => {
+      const result = await sessionService.updateSessionMeta(INBOX_SESSION_ID, {
+        title: 'New Title',
+      });
+      expect(result).toBeUndefined();
+    });
+
+    it('should update meta for normal session', async () => {
+      const meta = { title: 'Updated Title' };
+      await sessionService.updateSessionMeta(mockSessionId, meta);
+
+      const session = await clientDB.query.sessions.findFirst({
+        where: eq(sessions.id, mockSessionId),
+      });
+      expect(session).toBeDefined();
+    });
+  });
+
+  describe('updateSessionChatConfig', () => {
+    it('should update chat config', async () => {
+      const chatConfig = { temperature: 0.8 } as Partial<LobeAgentChatConfig>;
+      const result = await sessionService.updateSessionChatConfig(mockSessionId, chatConfig);
+      expect(result).toBeDefined();
+    });
+  });
+
+  describe('model getters', () => {
+    it('should return session model instance', () => {
+      // @ts-ignore - accessing private getter
+      const model = sessionService.sessionModel;
+      expect(model).toBeDefined();
+    });
+
+    it('should return session group model instance', () => {
+      // @ts-ignore - accessing private getter
+      const model = sessionService.sessionGroupModel;
+      expect(model).toBeDefined();
+    });
+  });
+
+  // SessionGroup related tests
+  describe('createSessionGroup', () => {
+    it('should create a new session group and return its id', async () => {
+      // Setup
+      const groupName = 'New Group';
+      const sort = 1;
+
+      // Execute
+      const groupId = await sessionService.createSessionGroup(groupName, sort);
+
+      // Assert
+      expect(groupId).toMatch(/^sg_/);
+
+      const result = await clientDB.query.sessionGroups.findFirst({
+        where: eq(sessionGroups.id, groupId),
+      });
+
+      expect(result).toMatchObject({ id: groupId, name: groupName, sort });
+    });
+  });
+
+  describe('removeSessionGroup', () => {
+    it('should remove a session group by its id', async () => {
+      const groupId = 'group-1';
+      // Execute
+      await sessionService.removeSessionGroup(groupId);
+
+      const result = await clientDB.query.sessionGroups.findFirst({
+        where: eq(sessionGroups.id, groupId),
+      });
+      // Assert
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe('clearSessionGroups', () => {
+    it('should clear all session groups', async () => {
+      // Execute
+      await sessionService.removeSessionGroups();
+
+      // Assert
+      const result = await clientDB.query.sessionGroups.findMany({
+        where: eq(sessionGroups.userId, userId),
+      });
+
+      expect(result.length).toBe(0);
+
+      const result2 = await clientDB.query.sessionGroups.findMany({
+        where: not(eq(sessionGroups.userId, userId)),
+      });
+
+      expect(result2.length).toBeGreaterThan(0);
+    });
+  });
+
+  describe('getSessionGroups', () => {
+    it('should retrieve all session groups', async () => {
+      // Execute
+      const result = await sessionService.getSessionGroups();
+
+      // Assert
+      const groups = [
+        { id: 'group-2', name: 'group-B', sort: 1 },
+        { id: 'group-1', name: 'group-A', sort: 2 },
+      ];
+      expect(result).toMatchObject(groups);
+    });
+  });
+
+  describe('updateSessionGroup', () => {
+    it('should update a session group', async () => {
+      // Setup
+      const groupId = 'group-1';
+      const data = { name: 'Updated Group', sort: 2 };
+
+      // Execute
+      await sessionService.updateSessionGroup(groupId, data);
+
+      // Assert
+      const result = await clientDB.query.sessionGroups.findFirst({
+        where: eq(sessionGroups.id, groupId),
+      });
+      expect(result).toMatchObject({ id: groupId, ...data });
+    });
+  });
+
+  describe('updateSessionGroupOrder', () => {
+    it('should update the order of session groups', async () => {
+      // Setup
+      const sortMap = [
+        { id: 'group-1', sort: 2 },
+        { id: 'group-2', sort: 1 },
+      ];
+
+      // Execute
+      await sessionService.updateSessionGroupOrder(sortMap);
+
+      // Assert
+      const data = await clientDB.query.sessionGroups.findMany({
+        where: eq(sessionGroups.userId, userId),
+      });
+      expect(data).toMatchObject([
+        { id: 'group-1', sort: 2 },
+        { id: 'group-2', sort: 1 },
+      ]);
+    });
+  });
+});

package/src/services/session/pglite.ts

@@ -0,0 +1,184 @@
+import { DeepPartial } from 'utility-types';
+
+import { INBOX_SESSION_ID } from '@/const/session';
+import { clientDB } from '@/database/client/db';
+import { AgentItem } from '@/database/schemas';
+import { SessionModel } from '@/database/server/models/session';
+import { SessionGroupModel } from '@/database/server/models/sessionGroup';
+import { BaseClientService } from '@/services/baseClientService';
+import { LobeAgentChatConfig, LobeAgentConfig } from '@/types/agent';
+import { MetaData } from '@/types/meta';
+import {
+  ChatSessionList,
+  LobeAgentSession,
+  LobeSessionType,
+  LobeSessions,
+  SessionGroupItem,
+  SessionGroups,
+  UpdateSessionParams,
+} from '@/types/session';
+
+import { ISessionService } from './type';
+
+export class ClientService extends BaseClientService implements ISessionService {
+  private get sessionModel(): SessionModel {
+    return new SessionModel(clientDB as any, this.userId);
+  }
+
+  private get sessionGroupModel(): SessionGroupModel {
+    return new SessionGroupModel(clientDB as any, this.userId);
+  }
+
+  async createSession(type: LobeSessionType, data: Partial<LobeAgentSession>): Promise<string> {
+    const { config, group, meta, ...session } = data;
+
+    const item = await this.sessionModel.create({
+      config: { ...config, ...meta } as any,
+      session: { ...session, groupId: group },
+      type,
+    });
+    if (!item) {
+      throw new Error('session create Error');
+    }
+    return item.id;
+  }
+
+  async batchCreateSessions(importSessions: LobeSessions) {
+    // @ts-ignore
+    return this.sessionModel.batchCreate(importSessions);
+  }
+
+  async cloneSession(id: string, newTitle: string): Promise<string | undefined> {
+    const res = await this.sessionModel.duplicate(id, newTitle);
+
+    if (res) return res?.id;
+  }
+
+  async getGroupedSessions(): Promise<ChatSessionList> {
+    return this.sessionModel.queryWithGroups();
+  }
+
+  async getSessionConfig(id: string): Promise<LobeAgentConfig> {
+    const res = await this.sessionModel.findByIdOrSlug(id);
+
+    if (!res) throw new Error('Session not found');
+
+    return res.agent as LobeAgentConfig;
+  }
+
+  /**
+   * 这个方法要对应移除的
+   */
+  async getSessionsByType(type: 'agent' | 'group' | 'all' = 'all'): Promise<LobeSessions> {
+    switch (type) {
+      // TODO: add a filter to get only agents or agents
+      case 'group': {
+        // @ts-ignore
+        return this.sessionModel.query();
+      }
+      case 'agent': {
+        // @ts-ignore
+        return this.sessionModel.query();
+      }
+
+      case 'all': {
+        // @ts-ignore
+        return this.sessionModel.query();
+      }
+    }
+  }
+
+  async countSessions() {
+    return this.sessionModel.count();
+  }
+
+  async searchSessions(keyword: string) {
+    return this.sessionModel.queryByKeyword(keyword);
+  }
+
+  async updateSession(id: string, value: Partial<UpdateSessionParams>) {
+    return this.sessionModel.update(id, {
+      ...value,
+      groupId: value.group === 'default' ? null : value.group,
+    });
+  }
+
+  async updateSessionConfig(
+    activeId: string,
+    config: DeepPartial<LobeAgentConfig>,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _?: AbortSignal,
+  ) {
+    const session = await this.sessionModel.findByIdOrSlug(activeId);
+    if (!session || !config) return;
+
+    return this.sessionModel.updateConfig(session.agent.id, config as AgentItem);
+  }
+
+  async updateSessionMeta(
+    activeId: string,
+    meta: Partial<MetaData>,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _?: AbortSignal,
+  ) {
+    // inbox 不允许修改 meta
+    if (activeId === INBOX_SESSION_ID) return;
+
+    return this.sessionModel.update(activeId, meta);
+  }
+
+  async updateSessionChatConfig(
+    activeId: string,
+    config: DeepPartial<LobeAgentChatConfig>,
+    // eslint-disable-next-line @typescript-eslint/no-unused-vars
+    _?: AbortSignal,
+  ) {
+    return this.updateSessionConfig(activeId, { chatConfig: config });
+  }
+
+  async removeSession(id: string) {
+    return this.sessionModel.delete(id);
+  }
+
+  async removeAllSessions() {
+    return this.sessionModel.deleteAll();
+  }
+
+  // ************************************** //
+  // ***********  SessionGroup  *********** //
+  // ************************************** //
+
+  async createSessionGroup(name: string, sort?: number) {
+    const item = await this.sessionGroupModel.create({ name, sort });
+    if (!item) {
+      throw new Error('session group create Error');
+    }
+
+    return item.id;
+  }
+
+  // eslint-disable-next-line @typescript-eslint/no-unused-vars
+  async batchCreateSessionGroups(_groups: SessionGroups) {
+    return { added: 0, ids: [], skips: [], success: true };
+  }
+
+  async removeSessionGroup(id: string) {
+    return await this.sessionGroupModel.delete(id);
+  }
+
+  async updateSessionGroup(id: string, data: Partial<SessionGroupItem>) {
+    return this.sessionGroupModel.update(id, data);
+  }
+
+  async updateSessionGroupOrder(sortMap: { id: string; sort: number }[]) {
+    return this.sessionGroupModel.updateOrder(sortMap);
+  }
+
+  async getSessionGroups(): Promise<SessionGroupItem[]> {
+    return this.sessionGroupModel.query();
+  }
+
+  async removeSessionGroups() {
+    return this.sessionGroupModel.deleteAll();
+  }
+}

package/src/services/session/type.ts

@@ -16,13 +16,21 @@ import {
 
 export interface ISessionService {
   createSession(type: LobeSessionType, defaultValue: Partial<LobeAgentSession>): Promise<string>;
+
+  /**
+   * 需要废弃
+   * @deprecated
+   */
   batchCreateSessions(importSessions: LobeSessions): Promise<any>;
   cloneSession(id: string, newTitle: string): Promise<string | undefined>;
 
   getGroupedSessions(): Promise<ChatSessionList>;
+
+  /**
+   * @deprecated
+   */
   getSessionsByType(type: 'agent' | 'group' | 'all'): Promise<LobeSessions>;
   countSessions(): Promise<number>;
-  hasSessions(): Promise<boolean>;
   searchSessions(keyword: string): Promise<LobeSessions>;
 
   updateSession(
@@ -53,6 +61,11 @@ export interface ISessionService {
   // ************************************** //
 
   createSessionGroup(name: string, sort?: number): Promise<string>;
+
+  /**
+   * 需要废弃
+   * @deprecated
+   */
   batchCreateSessionGroups(groups: SessionGroups): Promise<BatchTaskResult>;
 
   getSessionGroups(): Promise<SessionGroupItem[]>;

package/src/services/topic/client.test.ts

@@ -4,7 +4,7 @@ import { SessionModel } from '@/database/_deprecated/models/session';
 import { CreateTopicParams, TopicModel } from '@/database/_deprecated/models/topic';
 import { ChatTopic } from '@/types/topic';
 
-import { ClientService } from './client';
+import { ClientService } from './_deprecated';
 
 const topicService = new ClientService();
 // Mock the TopicModel

package/src/services/topic/index.ts

@@ -1,6 +1,9 @@
-
-import { ClientService } from './client';
+import { ClientService as DeprecatedService } from './_deprecated';
+import { ClientService } from './pglite';
 import { ServerService } from './server';
 
+const clientService =
+  process.env.NEXT_PUBLIC_CLIENT_DB === 'pglite' ? new ClientService() : new DeprecatedService();
+
 export const topicService =
-  process.env.NEXT_PUBLIC_SERVICE_MODE === 'server' ? new ServerService() : new ClientService();
+  process.env.NEXT_PUBLIC_SERVICE_MODE === 'server' ? new ServerService() : clientService;