@lobehub/chat 1.42.6 → 1.43.0

package/src/database/server/models/__tests__/session.test.ts

@@ -626,8 +626,6 @@ describe('SessionModel', () => {
     });
   });
 
-  // 在原有的 describe('SessionModel') 中添加以下测试套件
-
   describe('createInbox', () => {
     it('should create inbox session if not exists', async () => {
       const inbox = await sessionModel.createInbox();
@@ -782,4 +780,189 @@ describe('SessionModel', () => {
       });
     });
   });
+
+  describe('rank', () => {
+    it('should return ranked sessions based on topic count', async () => {
+      // Create test data
+      await serverDB.transaction(async (trx) => {
+        // Create sessions
+        await trx.insert(sessions).values([
+          { id: '1', userId },
+          { id: '2', userId },
+          { id: '3', userId },
+        ]);
+
+        // Create agents
+        await trx.insert(agents).values([
+          { id: 'a1', userId, title: 'Agent 1', avatar: 'avatar1', backgroundColor: 'bg1' },
+          { id: 'a2', userId, title: 'Agent 2', avatar: 'avatar2', backgroundColor: 'bg2' },
+          { id: 'a3', userId, title: 'Agent 3', avatar: 'avatar3', backgroundColor: 'bg3' },
+        ]);
+
+        // Link agents to sessions
+        await trx.insert(agentsToSessions).values([
+          { sessionId: '1', agentId: 'a1' },
+          { sessionId: '2', agentId: 'a2' },
+          { sessionId: '3', agentId: 'a3' },
+        ]);
+
+        // Create topics (different counts for ranking)
+        await trx.insert(topics).values([
+          { id: 't1', sessionId: '1', userId },
+          { id: 't2', sessionId: '1', userId },
+          { id: 't3', sessionId: '1', userId }, // Session 1 has 3 topics
+          { id: 't4', sessionId: '2', userId },
+          { id: 't5', sessionId: '2', userId }, // Session 2 has 2 topics
+          { id: 't6', sessionId: '3', userId }, // Session 3 has 1 topic
+        ]);
+      });
+
+      // Get ranked sessions with default limit
+      const result = await sessionModel.rank();
+
+      // Verify results
+      expect(result).toHaveLength(3);
+      // Should be ordered by topic count (descending)
+      expect(result[0]).toMatchObject({
+        id: '1',
+        count: 3,
+        title: 'Agent 1',
+        avatar: 'avatar1',
+        backgroundColor: 'bg1',
+      });
+      expect(result[1]).toMatchObject({
+        id: '2',
+        count: 2,
+        title: 'Agent 2',
+        avatar: 'avatar2',
+        backgroundColor: 'bg2',
+      });
+      expect(result[2]).toMatchObject({
+        id: '3',
+        count: 1,
+        title: 'Agent 3',
+        avatar: 'avatar3',
+        backgroundColor: 'bg3',
+      });
+    });
+
+    it('should respect the limit parameter', async () => {
+      // Create test data
+      await serverDB.transaction(async (trx) => {
+        // Create sessions and related data
+        await trx.insert(sessions).values([
+          { id: '1', userId },
+          { id: '2', userId },
+          { id: '3', userId },
+        ]);
+
+        await trx.insert(agents).values([
+          { id: 'a1', userId, title: 'Agent 1' },
+          { id: 'a2', userId, title: 'Agent 2' },
+          { id: 'a3', userId, title: 'Agent 3' },
+        ]);
+
+        await trx.insert(agentsToSessions).values([
+          { sessionId: '1', agentId: 'a1' },
+          { sessionId: '2', agentId: 'a2' },
+          { sessionId: '3', agentId: 'a3' },
+        ]);
+
+        await trx.insert(topics).values([
+          { id: 't1', sessionId: '1', userId },
+          { id: 't2', sessionId: '1', userId },
+          { id: 't3', sessionId: '2', userId },
+          { id: 't4', sessionId: '3', userId },
+        ]);
+      });
+
+      // Get ranked sessions with limit of 2
+      const result = await sessionModel.rank(2);
+
+      // Verify results
+      expect(result).toHaveLength(2);
+      expect(result[0].id).toBe('1'); // Most topics (2)
+      expect(result[1].id).toBe('2'); // Second most topics (1)
+    });
+
+    it('should handle sessions with no topics', async () => {
+      // Create test data
+      await serverDB.transaction(async (trx) => {
+        await trx.insert(sessions).values([
+          { id: '1', userId },
+          { id: '2', userId },
+        ]);
+
+        await trx.insert(agents).values([
+          { id: 'a1', userId, title: 'Agent 1' },
+          { id: 'a2', userId, title: 'Agent 2' },
+        ]);
+
+        await trx.insert(agentsToSessions).values([
+          { sessionId: '1', agentId: 'a1' },
+          { sessionId: '2', agentId: 'a2' },
+        ]);
+
+        // No topics created
+      });
+
+      const result = await sessionModel.rank();
+
+      expect(result).toHaveLength(0);
+    });
+  });
+
+  describe('hasMoreThanN', () => {
+    it('should return true when session count is more than N', async () => {
+      // Create test data
+      await serverDB.insert(sessions).values([
+        { id: '1', userId },
+        { id: '2', userId },
+        { id: '3', userId },
+      ]);
+
+      const result = await sessionModel.hasMoreThanN(2);
+      expect(result).toBe(true);
+    });
+
+    it('should return false when session count is equal to N', async () => {
+      // Create test data
+      await serverDB.insert(sessions).values([
+        { id: '1', userId },
+        { id: '2', userId },
+      ]);
+
+      const result = await sessionModel.hasMoreThanN(2);
+      expect(result).toBe(false);
+    });
+
+    it('should return false when session count is less than N', async () => {
+      // Create test data
+      await serverDB.insert(sessions).values([{ id: '1', userId }]);
+
+      const result = await sessionModel.hasMoreThanN(2);
+      expect(result).toBe(false);
+    });
+
+    it('should only count sessions for the current user', async () => {
+      // Create sessions for current user and another user
+      await serverDB.transaction(async (trx) => {
+        await trx.insert(users).values([{ id: 'other-user' }]);
+        await trx.insert(sessions).values([
+          { id: '1', userId }, // Current user
+          { id: '2', userId: 'other-user' }, // Other user
+          { id: '3', userId: 'other-user' }, // Other user
+        ]);
+      });
+
+      const result = await sessionModel.hasMoreThanN(1);
+      // Should return false as current user only has 1 session
+      expect(result).toBe(false);
+    });
+
+    it('should return false when no sessions exist', async () => {
+      const result = await sessionModel.hasMoreThanN(0);
+      expect(result).toBe(false);
+    });
+  });
 });

package/src/database/server/models/__tests__/topic.test.ts

@@ -620,4 +620,140 @@ describe('TopicModel', () => {
       );
     });
   });
+
+  describe('rank', () => {
+    it('should return ranked topics based on message count', async () => {
+      // 创建测试数据
+      await serverDB.transaction(async (tx) => {
+        await tx.insert(topics).values([
+          { id: 'topic1', title: 'Topic 1', sessionId, userId },
+          { id: 'topic2', title: 'Topic 2', sessionId, userId },
+          { id: 'topic3', title: 'Topic 3', sessionId, userId },
+        ]);
+
+        // topic1 有 3 条消息
+        await tx.insert(messages).values([
+          { id: 'msg1', role: 'user', topicId: 'topic1', userId },
+          { id: 'msg2', role: 'assistant', topicId: 'topic1', userId },
+          { id: 'msg3', role: 'user', topicId: 'topic1', userId },
+        ]);
+
+        // topic2 有 2 条消息
+        await tx.insert(messages).values([
+          { id: 'msg4', role: 'user', topicId: 'topic2', userId },
+          { id: 'msg5', role: 'assistant', topicId: 'topic2', userId },
+        ]);
+
+        // topic3 有 1 条消息
+        await tx.insert(messages).values([{ id: 'msg6', role: 'user', topicId: 'topic3', userId }]);
+      });
+
+      // 调用 rank 方法
+      const result = await topicModel.rank(2);
+
+      // 断言返回结果符合预期
+      expect(result).toHaveLength(2);
+      expect(result[0]).toMatchObject({
+        id: 'topic1',
+        title: 'Topic 1',
+        count: 3,
+        sessionId,
+      });
+      expect(result[1]).toMatchObject({
+        id: 'topic2',
+        title: 'Topic 2',
+        count: 2,
+        sessionId,
+      });
+    });
+
+    it('should return empty array if no topics exist', async () => {
+      const result = await topicModel.rank();
+      expect(result).toHaveLength(0);
+    });
+
+    it('should respect the limit parameter', async () => {
+      // 创建测试数据
+      await serverDB.transaction(async (tx) => {
+        await tx.insert(topics).values([
+          { id: 'topic1', title: 'Topic 1', sessionId, userId },
+          { id: 'topic2', title: 'Topic 2', sessionId, userId },
+        ]);
+
+        await tx.insert(messages).values([
+          { id: 'msg1', role: 'user', topicId: 'topic1', userId },
+          { id: 'msg2', role: 'user', topicId: 'topic2', userId },
+        ]);
+      });
+
+      // 使用限制为 1 调用 rank 方法
+      const result = await topicModel.rank(1);
+
+      // 断言只返回一个结果
+      expect(result).toHaveLength(1);
+    });
+  });
+
+  describe('count with date filters', () => {
+    beforeEach(async () => {
+      // 创建测试数据
+      await serverDB.insert(topics).values([
+        {
+          id: 'topic1',
+          userId,
+          createdAt: new Date('2023-01-01'),
+        },
+        {
+          id: 'topic2',
+          userId,
+          createdAt: new Date('2023-02-01'),
+        },
+        {
+          id: 'topic3',
+          userId,
+          createdAt: new Date('2023-03-01'),
+        },
+      ]);
+    });
+
+    it('should count topics with start date filter', async () => {
+      const result = await topicModel.count({
+        startDate: '2023-02-01',
+      });
+
+      expect(result).toBe(2); // should count topics from Feb 1st onwards
+    });
+
+    it('should count topics with end date filter', async () => {
+      const result = await topicModel.count({
+        endDate: '2023-02-01',
+      });
+
+      expect(result).toBe(2); // should count topics up to Feb 1st
+    });
+
+    it('should count topics within date range', async () => {
+      const result = await topicModel.count({
+        range: ['2023-01-15', '2023-02-15'],
+      });
+
+      expect(result).toBe(1); // should only count topic2
+    });
+
+    it('should return 0 if no topics match date filters', async () => {
+      const result = await topicModel.count({
+        range: ['2024-01-01', '2024-12-31'],
+      });
+
+      expect(result).toBe(0);
+    });
+
+    it('should handle invalid date filters gracefully', async () => {
+      const result = await topicModel.count({
+        startDate: 'invalid-date',
+      });
+
+      expect(result).toBe(3); // should return all topics if date is invalid
+    });
+  });
 });

package/src/database/server/models/__tests__/user.test.ts

@@ -1,3 +1,4 @@
+import dayjs from 'dayjs';
 import { eq } from 'drizzle-orm/expressions';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
@@ -5,7 +6,6 @@ import { INBOX_SESSION_ID } from '@/const/session';
 import { getTestDBInstance } from '@/database/server/core/dbForTest';
 import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt';
 import { UserGuide, UserPreference } from '@/types/user';
-import { UserSettings } from '@/types/user/settings';
 
 import { UserSettingsItem, userSettings, users } from '../../../schemas';
 import { SessionModel } from '../session';
@@ -268,4 +268,143 @@ describe('UserModel', () => {
       expect(result).toEqual({});
     });
   });
+
+  describe('getUserRegistrationDuration', () => {
+    it('should return default values when user not found', async () => {
+      const duration = await userModel.getUserRegistrationDuration();
+      const today = dayjs().format('YYYY-MM-DD');
+
+      expect(duration).toEqual({
+        createdAt: today,
+        duration: 1,
+        updatedAt: today,
+      });
+    });
+
+    it('should calculate correct duration for existing user', async () => {
+      // Mock the current date
+      const now = new Date('2024-01-15');
+      vi.setSystemTime(now);
+
+      const createdAt = new Date('2024-01-10'); // 5 days ago
+      await serverDB.insert(users).values({
+        id: userId,
+        createdAt,
+      });
+
+      const duration = await userModel.getUserRegistrationDuration();
+
+      expect(duration).toEqual({
+        createdAt: '2024-01-10',
+        duration: 6, // 5 days difference + 1
+        updatedAt: '2024-01-15',
+      });
+
+      vi.useRealTimers();
+    });
+  });
+
+  // 补充一些边界情况的测试
+  describe('edge cases', () => {
+    describe('updatePreference', () => {
+      it('should handle undefined preference', async () => {
+        await serverDB.insert(users).values({ id: userId });
+
+        const newPreference: Partial<UserPreference> = {
+          guide: { topic: true },
+        };
+
+        await userModel.updatePreference(newPreference);
+
+        const updatedUser = await serverDB.query.users.findFirst({
+          where: eq(users.id, userId),
+        });
+
+        expect(updatedUser?.preference).toMatchObject(newPreference);
+      });
+
+      it('should do nothing if user not found', async () => {
+        const nonExistentUserModel = new UserModel(serverDB, 'non-existent-id');
+        const result = await nonExistentUserModel.updatePreference({ guide: { topic: true } });
+        expect(result).toBeUndefined();
+      });
+    });
+
+    describe('updateGuide', () => {
+      it('should handle undefined guide', async () => {
+        await serverDB.insert(users).values({
+          id: userId,
+          preference: {} as UserPreference,
+        });
+
+        const newGuide: Partial<UserGuide> = {
+          topic: true,
+        };
+
+        await userModel.updateGuide(newGuide);
+
+        const updatedUser = await serverDB.query.users.findFirst({
+          where: eq(users.id, userId),
+        });
+        expect(updatedUser?.preference).toEqual({ guide: newGuide });
+      });
+
+      it('should do nothing if user not found', async () => {
+        const nonExistentUserModel = new UserModel(serverDB, 'non-existent-id');
+        const result = await nonExistentUserModel.updateGuide({ topic: true });
+        expect(result).toBeUndefined();
+      });
+    });
+
+    describe('createUser', () => {
+      it('should not create duplicate user with same id', async () => {
+        const params = {
+          id: userId,
+          username: 'existinguser',
+          email: 'existing@example.com',
+        };
+
+        // First creation
+        await UserModel.createUser(serverDB, params);
+
+        // Attempt to create with same ID but different details
+        await UserModel.createUser(serverDB, {
+          ...params,
+          username: 'newuser',
+          email: 'new@example.com',
+        });
+
+        const user = await UserModel.findById(serverDB, userId);
+        expect(user?.username).toBe('existinguser');
+        expect(user?.email).toBe('existing@example.com');
+      });
+    });
+
+    describe('getUserState', () => {
+      it('should handle empty settings', async () => {
+        await serverDB.insert(users).values({
+          id: userId,
+          preference: {} as UserPreference,
+          isOnboarded: true,
+        });
+
+        const state = await userModel.getUserState(KeyVaultsGateKeeper.getUserKeyVaults);
+
+        expect(state).toMatchObject({
+          userId,
+          isOnboarded: true,
+          preference: {},
+          settings: {
+            defaultAgent: {},
+            general: {},
+            keyVaults: {},
+            languageModel: {},
+            systemAgent: {},
+            tool: {},
+            tts: {},
+          },
+        });
+      });
+    });
+  });
 });

package/src/database/server/models/message.ts

@@ -1,7 +1,15 @@
-import { count } from 'drizzle-orm';
-import { and, asc, desc, eq, gte, inArray, isNull, like, lt } from 'drizzle-orm/expressions';
+import type { HeatmapsProps } from '@lobehub/charts';
+import dayjs from 'dayjs';
+import { count, sql } from 'drizzle-orm';
+import { and, asc, desc, eq, gt, inArray, isNotNull, isNull, like } from 'drizzle-orm/expressions';
 
 import { LobeChatDatabase } from '@/database/type';
+import {
+  genEndDateWhere,
+  genRangeWhere,
+  genStartDateWhere,
+  genWhere,
+} from '@/database/utils/genWhere';
 import { idGenerator } from '@/database/utils/idGenerator';
 import {
   ChatFileItem,
@@ -9,8 +17,10 @@ import {
   ChatTTS,
   ChatToolPayload,
   CreateMessageParams,
+  ModelRankItem,
 } from '@/types/message';
 import { merge } from '@/utils/merge';
+import { today } from '@/utils/time';
 
 import {
   MessageItem,
@@ -265,37 +275,122 @@ export class MessageModel {
     });
   };
 
-  count = async (): Promise<number> => {
+  count = async (params?: {
+    endDate?: string;
+    range?: [string, string];
+    startDate?: string;
+  }): Promise<number> => {
     const result = await this.db
       .select({
         count: count(messages.id),
       })
       .from(messages)
-      .where(eq(messages.userId, this.userId));
+      .where(
+        genWhere([
+          eq(messages.userId, this.userId),
+          params?.range
+            ? genRangeWhere(params.range, messages.createdAt, (date) => date.toDate())
+            : undefined,
+          params?.endDate
+            ? genEndDateWhere(params.endDate, messages.createdAt, (date) => date.toDate())
+            : undefined,
+          params?.startDate
+            ? genStartDateWhere(params.startDate, messages.createdAt, (date) => date.toDate())
+            : undefined,
+        ]),
+      );
 
     return result[0].count;
   };
 
-  countToday = async (): Promise<number> => {
-    const today = new Date();
-    today.setHours(0, 0, 0, 0);
-    const tomorrow = new Date(today);
-    tomorrow.setDate(tomorrow.getDate() + 1);
+  countWords = async (params?: {
+    endDate?: string;
+    range?: [string, string];
+    startDate?: string;
+  }): Promise<number> => {
+    const result = await this.db
+      .select({
+        count: sql<string>`sum(length(${messages.content}))`.as('total_length'),
+      })
+      .from(messages)
+      .where(
+        genWhere([
+          eq(messages.userId, this.userId),
+          params?.range
+            ? genRangeWhere(params.range, messages.createdAt, (date) => date.toDate())
+            : undefined,
+          params?.endDate
+            ? genEndDateWhere(params.endDate, messages.createdAt, (date) => date.toDate())
+            : undefined,
+          params?.startDate
+            ? genStartDateWhere(params.startDate, messages.createdAt, (date) => date.toDate())
+            : undefined,
+        ]),
+      );
+
+    return Number(result[0].count);
+  };
+
+  rankModels = async (limit: number = 10): Promise<ModelRankItem[]> => {
+    return this.db
+      .select({
+        count: count(messages.id).as('count'),
+        id: messages.model,
+      })
+      .from(messages)
+      .where(and(eq(messages.userId, this.userId), isNotNull(messages.model)))
+      .having(({ count }) => gt(count, 0))
+      .groupBy(messages.model)
+      .orderBy(desc(sql`count`), asc(messages.model))
+      .limit(limit);
+  };
+
+  getHeatmaps = async (): Promise<HeatmapsProps['data']> => {
+    const startDate = today().subtract(1, 'year').startOf('day');
+    const endDate = today().endOf('day');
 
     const result = await this.db
       .select({
         count: count(messages.id),
+        date: sql`DATE(${messages.createdAt})`.as('heatmaps_date'),
       })
       .from(messages)
       .where(
-        and(
+        genWhere([
           eq(messages.userId, this.userId),
-          gte(messages.createdAt, today),
-          lt(messages.createdAt, tomorrow),
-        ),
+          genRangeWhere(
+            [startDate.format('YYYY-MM-DD'), endDate.add(1, 'day').format('YYYY-MM-DD')],
+            messages.createdAt,
+            (date) => date.toDate(),
+          ),
+        ]),
+      )
+      .groupBy(sql`heatmaps_date`)
+      .orderBy(desc(sql`heatmaps_date`));
+
+    const heatmapData: HeatmapsProps['data'] = [];
+    let currentDate = startDate;
+
+    while (currentDate.isBefore(endDate) || currentDate.isSame(endDate, 'day')) {
+      const formattedDate = currentDate.format('YYYY-MM-DD');
+      const matchingResult = result.find(
+        (r) => r?.date && dayjs(r.date as string).format('YYYY-MM-DD') === formattedDate,
       );
 
-    return result[0].count;
+      const count = matchingResult ? matchingResult.count : 0;
+      const levelCount = count > 0 ? Math.floor(count / 5) : 0;
+      const level = levelCount > 4 ? 4 : levelCount;
+
+      heatmapData.push({
+        count,
+        date: formattedDate,
+        level,
+      });
+
+      currentDate = currentDate.add(1, 'day');
+    }
+
+    return heatmapData;
   };
 
   hasMoreThanN = async (n: number): Promise<boolean> => {