@lobehub/chat 1.42.5 → 1.43.0

package/src/database/server/models/__tests__/message.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';
 
@@ -1042,41 +1043,6 @@ describe('MessageModel', () => {
     });
   });
 
-  describe('countToday', () => {
-    it('should return the count of messages created today', async () => {
-      // 创建测试数据
-      await serverDB.insert(messages).values([
-        {
-          id: '1',
-          userId,
-          role: 'user',
-          content: 'message 1',
-          createdAt: new Date(),
-        },
-        {
-          id: '2',
-          userId,
-          role: 'user',
-          content: 'message 2',
-          createdAt: new Date(),
-        },
-        {
-          id: '3',
-          userId,
-          role: 'user',
-          content: 'message 3',
-          createdAt: new Date('2023-01-01'),
-        },
-      ]);
-
-      // 调用 countToday 方法
-      const result = await messageModel.countToday();
-
-      // 断言结果
-      expect(result).toBe(2);
-    });
-  });
-
   describe('findMessageQueriesById', () => {
     it('should return undefined for non-existent message query', async () => {
       const result = await messageModel.findMessageQueriesById('non-existent-id');
@@ -1211,4 +1177,349 @@ describe('MessageModel', () => {
       expect(id2).toMatch(/^msg_/);
     });
   });
+
+  describe('countWords', () => {
+    it('should count total words of messages belonging to the user', async () => {
+      // 创建测试数据
+      await serverDB.insert(messages).values([
+        { id: '1', userId, role: 'user', content: 'hello world' },
+        { id: '2', userId, role: 'user', content: 'test message' },
+        { id: '3', userId: '456', role: 'user', content: 'other user message' },
+      ]);
+
+      // 调用 countWords 方法
+      const result = await messageModel.countWords();
+
+      // 断言结果 - 'hello world' + 'test message' = 23 characters
+      expect(result).toEqual(23);
+    });
+
+    it('should count words within date range', async () => {
+      // 创建测试数据
+      await serverDB.insert(messages).values([
+        {
+          id: '1',
+          userId,
+          role: 'user',
+          content: 'old message',
+          createdAt: new Date('2023-01-01'),
+        },
+        {
+          id: '2',
+          userId,
+          role: 'user',
+          content: 'new message',
+          createdAt: new Date('2023-06-01'),
+        },
+      ]);
+
+      // 调用 countWords 方法，设置日期范围
+      const result = await messageModel.countWords({
+        range: ['2023-05-01', '2023-07-01'],
+      });
+
+      // 断言结果 - 只计算 'new message' = 11 characters
+      expect(result).toEqual(11);
+    });
+
+    it('should handle empty content', async () => {
+      // 创建测试数据
+      await serverDB.insert(messages).values([
+        { id: '1', userId, role: 'user', content: '' },
+        { id: '2', userId, role: 'user', content: null },
+      ]);
+
+      // 调用 countWords 方法
+      const result = await messageModel.countWords();
+
+      // 断言结果
+      expect(result).toEqual(0);
+    });
+  });
+
+  describe('getHeatmaps', () => {
+    it('should return heatmap data for the last year', async () => {
+      const today = dayjs();
+
+      // 创建测试数据
+      await serverDB.insert(messages).values([
+        {
+          id: '1',
+          userId,
+          role: 'user',
+          content: 'message 1',
+          createdAt: today.subtract(2, 'day').toDate(),
+        },
+        {
+          id: '2',
+          userId,
+          role: 'user',
+          content: 'message 2',
+          createdAt: today.subtract(2, 'day').toDate(),
+        },
+        {
+          id: '3',
+          userId,
+          role: 'user',
+          content: 'message 3',
+          createdAt: today.subtract(1, 'day').toDate(),
+        },
+      ]);
+
+      // 调用 getHeatmaps 方法
+      const result = await messageModel.getHeatmaps();
+
+      // 断言结果
+      expect(result.length).toBeGreaterThan(366);
+      expect(result.length).toBeLessThan(368);
+
+      // 检查两天前的数据
+      const twoDaysAgo = result.find(
+        (item) => item.date === today.subtract(2, 'day').format('YYYY-MM-DD'),
+      );
+      expect(twoDaysAgo?.count).toBe(2);
+      expect(twoDaysAgo?.level).toBe(0);
+
+      // 检查一天前的数据
+      const oneDayAgo = result.find(
+        (item) => item.date === today.subtract(1, 'day').format('YYYY-MM-DD'),
+      );
+      expect(oneDayAgo?.count).toBe(1);
+      expect(oneDayAgo?.level).toBe(0);
+
+      // 检查今天的数据
+      const todayData = result.find((item) => item.date === today.format('YYYY-MM-DD'));
+      expect(todayData?.count).toBe(0);
+      expect(todayData?.level).toBe(0);
+    });
+
+    it('should calculate correct levels based on message count', async () => {
+      const today = dayjs();
+
+      // 创建测试数据 - 不同数量的消息以测试不同的等级
+      await serverDB.insert(messages).values([
+        // 1 message - level 0
+        {
+          id: '1',
+          userId,
+          role: 'user',
+          content: 'message 1',
+          createdAt: today.subtract(4, 'day').toDate(),
+        },
+        // 6 messages - level 1
+        ...Array(6)
+          .fill(0)
+          .map((_, i) => ({
+            id: `2-${i}`,
+            userId,
+            role: 'user',
+            content: `message 2-${i}`,
+            createdAt: today.subtract(3, 'day').toDate(),
+          })),
+        // 11 messages - level 2
+        ...Array(11)
+          .fill(0)
+          .map((_, i) => ({
+            id: `3-${i}`,
+            userId,
+            role: 'user',
+            content: `message 3-${i}`,
+            createdAt: today.subtract(2, 'day').toDate(),
+          })),
+        // 16 messages - level 3
+        ...Array(16)
+          .fill(0)
+          .map((_, i) => ({
+            id: `4-${i}`,
+            userId,
+            role: 'user',
+            content: `message 4-${i}`,
+            createdAt: today.subtract(1, 'day').toDate(),
+          })),
+        // 21 messages - level 4
+        ...Array(21)
+          .fill(0)
+          .map((_, i) => ({
+            id: `5-${i}`,
+            userId,
+            role: 'user',
+            content: `message 5-${i}`,
+            createdAt: today.toDate(),
+          })),
+      ]);
+
+      // 调用 getHeatmaps 方法
+      const result = await messageModel.getHeatmaps();
+
+      // 检查不同天数的等级
+      const fourDaysAgo = result.find(
+        (item) => item.date === today.subtract(4, 'day').format('YYYY-MM-DD'),
+      );
+      expect(fourDaysAgo?.count).toBe(1);
+      expect(fourDaysAgo?.level).toBe(0);
+
+      const threeDaysAgo = result.find(
+        (item) => item.date === today.subtract(3, 'day').format('YYYY-MM-DD'),
+      );
+      expect(threeDaysAgo?.count).toBe(6);
+      expect(threeDaysAgo?.level).toBe(1);
+
+      const twoDaysAgo = result.find(
+        (item) => item.date === today.subtract(2, 'day').format('YYYY-MM-DD'),
+      );
+      expect(twoDaysAgo?.count).toBe(11);
+      expect(twoDaysAgo?.level).toBe(2);
+
+      const oneDayAgo = result.find(
+        (item) => item.date === today.subtract(1, 'day').format('YYYY-MM-DD'),
+      );
+      expect(oneDayAgo?.count).toBe(16);
+      expect(oneDayAgo?.level).toBe(3);
+
+      const todayData = result.find((item) => item.date === today.format('YYYY-MM-DD'));
+      expect(todayData?.count).toBe(21);
+      expect(todayData?.level).toBe(4);
+    });
+
+    it('should handle empty data', async () => {
+      // 不创建任何消息数据
+
+      // 调用 getHeatmaps 方法
+      const result = await messageModel.getHeatmaps();
+
+      // 断言结果
+      expect(result.length).toBeGreaterThan(366);
+      expect(result.length).toBeLessThan(368);
+
+      // 检查所有数据的 count 和 level 是否为 0
+      result.forEach((item) => {
+        expect(item.count).toBe(0);
+        expect(item.level).toBe(0);
+      });
+    });
+  });
+
+  describe('rankModels', () => {
+    it('should rank models by usage count', async () => {
+      // 创建测试数据
+      await serverDB.insert(messages).values([
+        { id: '1', userId, role: 'assistant', content: 'message 1', model: 'gpt-3.5' },
+        { id: '2', userId, role: 'assistant', content: 'message 2', model: 'gpt-3.5' },
+        { id: '3', userId, role: 'assistant', content: 'message 3', model: 'gpt-4' },
+        { id: '4', userId: '456', role: 'assistant', content: 'message 4', model: 'gpt-3.5' }, // 其他用户的消息
+      ]);
+
+      // 调用 rankModels 方法
+      const result = await messageModel.rankModels();
+
+      // 断言结果
+      expect(result).toHaveLength(2);
+      expect(result[0]).toEqual({ id: 'gpt-3.5', count: 2 }); // 当前用户使用 gpt-3.5 两次
+      expect(result[1]).toEqual({ id: 'gpt-4', count: 1 }); // 当前用户使用 gpt-4 一次
+    });
+
+    it('should only count messages with model field', async () => {
+      // 创建测试数据，包括没有 model 字段的消息
+      await serverDB.insert(messages).values([
+        { id: '1', userId, role: 'assistant', content: 'message 1', model: 'gpt-3.5' },
+        { id: '2', userId, role: 'assistant', content: 'message 2', model: null },
+        { id: '3', userId, role: 'user', content: 'message 3' }, // 用户消息通常没有 model
+      ]);
+
+      // 调用 rankModels 方法
+      const result = await messageModel.rankModels();
+
+      // 断言结果
+      expect(result).toHaveLength(1);
+      expect(result[0]).toEqual({ id: 'gpt-3.5', count: 1 });
+    });
+
+    it('should return empty array when no models are used', async () => {
+      // 创建测试数据，所有消息都没有 model
+      await serverDB.insert(messages).values([
+        { id: '1', userId, role: 'user', content: 'message 1' },
+        { id: '2', userId, role: 'assistant', content: 'message 2' },
+      ]);
+
+      // 调用 rankModels 方法
+      const result = await messageModel.rankModels();
+
+      // 断言结果
+      expect(result).toHaveLength(0);
+    });
+
+    it('should order models by count in descending order', async () => {
+      // 创建测试数据，使用不同次数的模型
+      await serverDB.insert(messages).values([
+        { id: '1', userId, role: 'assistant', content: 'message 1', model: 'gpt-4' },
+        { id: '2', userId, role: 'assistant', content: 'message 2', model: 'gpt-3.5' },
+        { id: '3', userId, role: 'assistant', content: 'message 3', model: 'gpt-3.5' },
+        { id: '4', userId, role: 'assistant', content: 'message 4', model: 'claude' },
+        { id: '5', userId, role: 'assistant', content: 'message 5', model: 'gpt-3.5' },
+      ]);
+
+      // 调用 rankModels 方法
+      const result = await messageModel.rankModels();
+
+      // 断言结果
+      expect(result).toHaveLength(3);
+      expect(result[0]).toEqual({ id: 'gpt-3.5', count: 3 }); // 最多使用
+      expect(result[1]).toEqual({ id: 'claude', count: 1 });
+      expect(result[2]).toEqual({ id: 'gpt-4', count: 1 });
+    });
+  });
+
+  describe('hasMoreThanN', () => {
+    it('should return true when message count is greater than N', async () => {
+      // 创建测试数据
+      await serverDB.insert(messages).values([
+        { id: '1', userId, role: 'user', content: 'message 1' },
+        { id: '2', userId, role: 'user', content: 'message 2' },
+        { id: '3', userId, role: 'user', content: 'message 3' },
+      ]);
+
+      // 测试不同的 N 值
+      const result1 = await messageModel.hasMoreThanN(2); // 3 > 2
+      const result2 = await messageModel.hasMoreThanN(3); // 3 ≯ 3
+      const result3 = await messageModel.hasMoreThanN(4); // 3 ≯ 4
+
+      expect(result1).toBe(true);
+      expect(result2).toBe(false);
+      expect(result3).toBe(false);
+    });
+
+    it('should only count messages belonging to the user', async () => {
+      // 创建测试数据，包括其他用户的消息
+      await serverDB.insert(messages).values([
+        { id: '1', userId, role: 'user', content: 'message 1' },
+        { id: '2', userId, role: 'user', content: 'message 2' },
+        { id: '3', userId: '456', role: 'user', content: 'message 3' }, // 其他用户的消息
+      ]);
+
+      const result = await messageModel.hasMoreThanN(2);
+
+      expect(result).toBe(false); // 当前用户只有 2 条消息，不大于 2
+    });
+
+    it('should return false when no messages exist', async () => {
+      const result = await messageModel.hasMoreThanN(0);
+      expect(result).toBe(false);
+    });
+
+    it('should handle edge cases', async () => {
+      // 创建一条消息
+      await serverDB
+        .insert(messages)
+        .values([{ id: '1', userId, role: 'user', content: 'message 1' }]);
+
+      // 测试边界情况
+      const result1 = await messageModel.hasMoreThanN(0); // 1 > 0
+      const result2 = await messageModel.hasMoreThanN(1); // 1 ≯ 1
+      const result3 = await messageModel.hasMoreThanN(-1); // 1 > -1
+
+      expect(result1).toBe(true);
+      expect(result2).toBe(false);
+      expect(result3).toBe(true);
+    });
+  });
 });

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);
+    });
+  });
 });