@lobehub/chat 1.42.6 → 1.43.1

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> => {

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

@@ -1,13 +1,20 @@
 import { Column, count, sql } from 'drizzle-orm';
-import { and, asc, desc, eq, inArray, like, not, or } from 'drizzle-orm/expressions';
+import { and, asc, desc, eq, gt, inArray, isNull, like, not, or } from 'drizzle-orm/expressions';
 
 import { appEnv } from '@/config/app';
+import { DEFAULT_INBOX_AVATAR } from '@/const/meta';
 import { INBOX_SESSION_ID } from '@/const/session';
 import { DEFAULT_AGENT_CONFIG } from '@/const/settings';
 import { LobeChatDatabase } from '@/database/type';
+import {
+  genEndDateWhere,
+  genRangeWhere,
+  genStartDateWhere,
+  genWhere,
+} from '@/database/utils/genWhere';
 import { idGenerator } from '@/database/utils/idGenerator';
 import { parseAgentConfig } from '@/server/globalConfig/parseDefaultAgent';
-import { ChatSessionList, LobeAgentSession } from '@/types/session';
+import { ChatSessionList, LobeAgentSession, SessionRankItem } from '@/types/session';
 import { merge } from '@/utils/merge';
 
 import {
@@ -19,6 +26,7 @@ import {
   agentsToSessions,
   sessionGroups,
   sessions,
+  topics,
 } from '../../schemas';
 
 export class SessionModel {
@@ -84,17 +92,81 @@ export class SessionModel {
     return { ...result, agent: (result?.agentsToSessions?.[0] as any)?.agent } as any;
   };
 
-  count = async (): Promise<number> => {
+  count = async (params?: {
+    endDate?: string;
+    range?: [string, string];
+    startDate?: string;
+  }): Promise<number> => {
     const result = await this.db
       .select({
         count: count(sessions.id),
       })
       .from(sessions)
-      .where(eq(sessions.userId, this.userId));
+      .where(
+        genWhere([
+          eq(sessions.userId, this.userId),
+          params?.range
+            ? genRangeWhere(params.range, sessions.createdAt, (date) => date.toDate())
+            : undefined,
+          params?.endDate
+            ? genEndDateWhere(params.endDate, sessions.createdAt, (date) => date.toDate())
+            : undefined,
+          params?.startDate
+            ? genStartDateWhere(params.startDate, sessions.createdAt, (date) => date.toDate())
+            : undefined,
+        ]),
+      );
 
     return result[0].count;
   };
 
+  _rank = async (limit: number = 10): Promise<SessionRankItem[]> => {
+    return this.db
+      .select({
+        avatar: agents.avatar,
+        backgroundColor: agents.backgroundColor,
+        count: count(topics.id).as('count'),
+        id: sessions.id,
+        title: agents.title,
+      })
+      .from(sessions)
+      .where(and(eq(sessions.userId, this.userId)))
+      .leftJoin(topics, eq(sessions.id, topics.sessionId))
+      .leftJoin(agentsToSessions, eq(sessions.id, agentsToSessions.sessionId))
+      .leftJoin(agents, eq(agentsToSessions.agentId, agents.id))
+      .groupBy(sessions.id, agentsToSessions.agentId, agents.id)
+      .having(({ count }) => gt(count, 0))
+      .orderBy(desc(sql`count`))
+      .limit(limit);
+  };
+
+  // TODO: 未来将 Inbox id 入库后可以直接使用 _rank 方法
+  rank = async (limit: number = 10): Promise<SessionRankItem[]> => {
+    const inboxResult = await this.db
+      .select({
+        count: count(topics.id).as('count'),
+      })
+      .from(topics)
+      .where(and(eq(topics.userId, this.userId), isNull(topics.sessionId)));
+
+    const inboxCount = inboxResult[0].count;
+
+    if (!inboxCount || inboxCount === 0) return this._rank(limit);
+
+    const result = await this._rank(limit ? limit - 1 : undefined);
+
+    return [
+      {
+        avatar: DEFAULT_INBOX_AVATAR,
+        backgroundColor: null,
+        count: inboxCount,
+        id: INBOX_SESSION_ID,
+        title: 'inbox.title',
+      },
+      ...result,
+    ].sort((a, b) => b.count - a.count);
+  };
+
   hasMoreThanN = async (n: number): Promise<boolean> => {
     const result = await this.db
       .select({ id: sessions.id })

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

@@ -1,8 +1,15 @@
 import { Column, count, sql } from 'drizzle-orm';
-import { and, desc, eq, exists, inArray, isNull, like, or } from 'drizzle-orm/expressions';
+import { and, desc, eq, exists, gt, inArray, isNull, like, or } 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 { TopicRankItem } from '@/types/topic';
 
 import { NewMessage, TopicItem, messages, topics } from '../../schemas';
 
@@ -92,17 +99,51 @@ export class TopicModel {
     });
   };
 
-  count = async (): Promise<number> => {
+  count = async (params?: {
+    endDate?: string;
+    range?: [string, string];
+    startDate?: string;
+  }): Promise<number> => {
     const result = await this.db
       .select({
         count: count(topics.id),
       })
       .from(topics)
-      .where(eq(topics.userId, this.userId));
+      .where(
+        genWhere([
+          eq(topics.userId, this.userId),
+          params?.range
+            ? genRangeWhere(params.range, topics.createdAt, (date) => date.toDate())
+            : undefined,
+          params?.endDate
+            ? genEndDateWhere(params.endDate, topics.createdAt, (date) => date.toDate())
+            : undefined,
+          params?.startDate
+            ? genStartDateWhere(params.startDate, topics.createdAt, (date) => date.toDate())
+            : undefined,
+        ]),
+      );
 
     return result[0].count;
   };
 
+  rank = async (limit: number = 10): Promise<TopicRankItem[]> => {
+    return this.db
+      .select({
+        count: count(messages.id).as('count'),
+        id: topics.id,
+        sessionId: topics.sessionId,
+        title: topics.title,
+      })
+      .from(topics)
+      .where(and(eq(topics.userId, this.userId)))
+      .leftJoin(messages, eq(topics.id, messages.topicId))
+      .groupBy(topics.id)
+      .orderBy(desc(sql`count`))
+      .having(({ count }) => gt(count, 0))
+      .limit(limit);
+  };
+
   // **************** Create *************** //
 
   create = async (

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

@@ -1,4 +1,5 @@
 import { TRPCError } from '@trpc/server';
+import dayjs from 'dayjs';
 import { eq } from 'drizzle-orm/expressions';
 import { DeepPartial } from 'utility-types';
 
@@ -6,6 +7,7 @@ import { LobeChatDatabase } from '@/database/type';
 import { UserGuide, UserPreference } from '@/types/user';
 import { UserKeyVaults, UserSettings } from '@/types/user/settings';
 import { merge } from '@/utils/merge';
+import { today } from '@/utils/time';
 
 import { NewUser, UserItem, UserSettingsItem, userSettings, users } from '../../schemas';
 import { SessionModel } from './session';
@@ -30,6 +32,26 @@ export class UserModel {
     this.db = db;
   }
 
+  getUserRegistrationDuration = async (): Promise<{
+    createdAt: string;
+    duration: number;
+    updatedAt: string;
+  }> => {
+    const user = await this.db.query.users.findFirst({ where: eq(users.id, this.userId) });
+    if (!user)
+      return {
+        createdAt: today().format('YYYY-MM-DD'),
+        duration: 1,
+        updatedAt: today().format('YYYY-MM-DD'),
+      };
+
+    return {
+      createdAt: dayjs(user.createdAt).format('YYYY-MM-DD'),
+      duration: dayjs().diff(dayjs(user.createdAt), 'day') + 1,
+      updatedAt: today().format('YYYY-MM-DD'),
+    };
+  };
+
   getUserState = async (decryptor: DecryptUserKeyVaults) => {
     const result = await this.db
       .select({