@lobehub/chat 1.35.9 → 1.35.11

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

@@ -13,15 +13,9 @@ import { UserModel } from '../user';
 
 let serverDB = await getTestDBInstance();
 
-vi.mock('@/database/server/core/db', async () => ({
-  get serverDB() {
-    return serverDB;
-  },
-}));
-
 const userId = 'user-db';
 const userEmail = 'user@example.com';
-const userModel = new UserModel();
+const userModel = new UserModel(serverDB, userId);
 
 beforeEach(async () => {
   await serverDB.delete(users);
@@ -44,14 +38,14 @@ describe('UserModel', () => {
         email: 'test@example.com',
       };
 
-      await UserModel.createUser(params);
+      await UserModel.createUser(serverDB, params);
 
       const user = await serverDB.query.users.findFirst({ where: eq(users.id, userId) });
       expect(user).not.toBeNull();
       expect(user?.username).toBe('testuser');
       expect(user?.email).toBe('test@example.com');
 
-      const sessionModel = new SessionModel(userId);
+      const sessionModel = new SessionModel(serverDB, userId);
       const inbox = await sessionModel.findByIdOrSlug(INBOX_SESSION_ID);
       expect(inbox).not.toBeNull();
     });
@@ -61,7 +55,7 @@ describe('UserModel', () => {
     it('should delete a user', async () => {
       await serverDB.insert(users).values({ id: userId });
 
-      await UserModel.deleteUser(userId);
+      await UserModel.deleteUser(serverDB, userId);
 
       const user = await serverDB.query.users.findFirst({ where: eq(users.id, userId) });
       expect(user).toBeUndefined();
@@ -72,7 +66,7 @@ describe('UserModel', () => {
     it('should find a user by ID', async () => {
       await serverDB.insert(users).values({ id: userId, username: 'testuser' });
 
-      const user = await UserModel.findById(userId);
+      const user = await UserModel.findById(serverDB, userId);
 
       expect(user).not.toBeNull();
       expect(user?.id).toBe(userId);
@@ -84,7 +78,7 @@ describe('UserModel', () => {
     it('should find a user by email', async () => {
       await serverDB.insert(users).values({ id: userId, email: userEmail });
 
-      const user = await UserModel.findByEmail(userEmail);
+      const user = await UserModel.findByEmail(serverDB, userEmail);
 
       expect(user).not.toBeNull();
       expect(user?.id).toBe(userId);
@@ -107,7 +101,7 @@ describe('UserModel', () => {
         keyVaults: encryptedKeyVaults,
       });
 
-      const state = await userModel.getUserState(userId);
+      const state = await userModel.getUserState();
 
       expect(state.userId).toBe(userId);
       expect(state.preference).toEqual(preference);
@@ -115,7 +109,9 @@ describe('UserModel', () => {
     });
 
     it('should throw an error if user not found', async () => {
-      await expect(userModel.getUserState('invalid-user-id')).rejects.toThrow('user not found');
+      const userModel = new UserModel(serverDB, 'invalid-user-id');
+
+      await expect(userModel.getUserState()).rejects.toThrow('user not found');
     });
   });
 
@@ -123,7 +119,7 @@ describe('UserModel', () => {
     it('should update user fields', async () => {
       await serverDB.insert(users).values({ id: userId, username: 'oldname' });
 
-      await userModel.updateUser(userId, { username: 'newname' });
+      await userModel.updateUser({ username: 'newname' });
 
       const updatedUser = await serverDB.query.users.findFirst({
         where: eq(users.id, userId),
@@ -137,7 +133,7 @@ describe('UserModel', () => {
       await serverDB.insert(users).values({ id: userId });
       await serverDB.insert(userSettings).values({ id: userId });
 
-      await userModel.deleteSetting(userId);
+      await userModel.deleteSetting();
 
       const settings = await serverDB.query.userSettings.findFirst({
         where: eq(users.id, userId),
@@ -155,7 +151,7 @@ describe('UserModel', () => {
       } as UserSettings;
       await serverDB.insert(users).values({ id: userId });
 
-      await userModel.updateSetting(userId, settings);
+      await userModel.updateSetting(settings);
 
       const updatedSettings = await serverDB.query.userSettings.findFirst({
         where: eq(users.id, userId),
@@ -178,7 +174,7 @@ describe('UserModel', () => {
       const newSettings = {
         general: { fontSize: 16, language: 'zh-CN', themeMode: 'dark' },
       } as UserSettings;
-      await userModel.updateSetting(userId, newSettings);
+      await userModel.updateSetting(newSettings);
 
       const updatedSettings = await serverDB.query.userSettings.findFirst({
         where: eq(users.id, userId),
@@ -195,7 +191,7 @@ describe('UserModel', () => {
       const newPreference: Partial<UserPreference> = {
         guide: { topic: true, moveSettingsToAvatar: true },
       };
-      await userModel.updatePreference(userId, newPreference);
+      await userModel.updatePreference(newPreference);
 
       const updatedUser = await serverDB.query.users.findFirst({ where: eq(users.id, userId) });
       expect(updatedUser?.preference).toEqual({ ...preference, ...newPreference });
@@ -212,10 +208,49 @@ describe('UserModel', () => {
         moveSettingsToAvatar: true,
         uploadFileInKnowledgeBase: true,
       };
-      await userModel.updateGuide(userId, newGuide);
+      await userModel.updateGuide(newGuide);
 
       const updatedUser = await serverDB.query.users.findFirst({ where: eq(users.id, userId) });
       expect(updatedUser?.preference).toEqual({ ...preference, guide: newGuide });
     });
   });
+
+  describe('getUserApiKeys', () => {
+    it('should get and decrypt user API keys', async () => {
+      const keyVaults = { openai: { apiKey: 'test-key' } };
+      const gateKeeper = await KeyVaultsGateKeeper.initWithEnvKey();
+      const encryptedKeyVaults = await gateKeeper.encrypt(JSON.stringify(keyVaults));
+
+      const userId = 'user-api-id';
+
+      await serverDB.insert(users).values({ id: userId });
+      await serverDB.insert(userSettings).values({
+        id: userId,
+        keyVaults: encryptedKeyVaults,
+      });
+
+      const result = await UserModel.getUserApiKeys(serverDB, userId);
+      expect(result).toEqual(keyVaults);
+    });
+
+    it('should throw error when user not found', async () => {
+      await expect(UserModel.getUserApiKeys(serverDB, 'non-existent-id')).rejects.toThrow(
+        'user not found',
+      );
+    });
+
+    it('should handle decrypt failure and return empty object', async () => {
+      const userId = 'user-api-test-id';
+      // 模拟解密失败的情况
+      const invalidEncryptedData = 'invalid:-encrypted-:data';
+      await serverDB.insert(users).values({ id: userId });
+      await serverDB.insert(userSettings).values({
+        id: userId,
+        keyVaults: invalidEncryptedData,
+      });
+
+      const result = await UserModel.getUserApiKeys(serverDB, userId);
+      expect(result).toEqual({});
+    });
+  });
 });

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

@@ -1,19 +1,21 @@
 import { eq } from 'drizzle-orm';
 import { and, desc } from 'drizzle-orm/expressions';
 
-import { serverDB } from '@/database/server';
+import { LobeChatDatabase } from '@/database/type';
 
 import { NewSessionGroup, SessionGroupItem, sessionGroups } from '../../schemas';
 
 export class TemplateModel {
   private userId: string;
+  private db: LobeChatDatabase;
 
-  constructor(userId: string) {
+  constructor(db: LobeChatDatabase, userId: string) {
     this.userId = userId;
+    this.db = db;
   }
 
   create = async (params: NewSessionGroup) => {
-    const [result] = await serverDB
+    const [result] = await this.db
       .insert(sessionGroups)
       .values({ ...params, userId: this.userId })
       .returning();
@@ -22,30 +24,30 @@ export class TemplateModel {
   };
 
   delete = async (id: string) => {
-    return serverDB
+    return this.db
       .delete(sessionGroups)
       .where(and(eq(sessionGroups.id, id), eq(sessionGroups.userId, this.userId)));
   };
 
   deleteAll = async () => {
-    return serverDB.delete(sessionGroups).where(eq(sessionGroups.userId, this.userId));
+    return this.db.delete(sessionGroups).where(eq(sessionGroups.userId, this.userId));
   };
 
   query = async () => {
-    return serverDB.query.sessionGroups.findMany({
+    return this.db.query.sessionGroups.findMany({
       orderBy: [desc(sessionGroups.updatedAt)],
       where: eq(sessionGroups.userId, this.userId),
     });
   };
 
   findById = async (id: string) => {
-    return serverDB.query.sessionGroups.findFirst({
+    return this.db.query.sessionGroups.findFirst({
       where: and(eq(sessionGroups.id, id), eq(sessionGroups.userId, this.userId)),
     });
   };
 
   async update(id: string, value: Partial<SessionGroupItem>) {
-    return serverDB
+    return this.db
       .update(sessionGroups)
       .set({ ...value, updatedAt: new Date() })
       .where(and(eq(sessionGroups.id, id), eq(sessionGroups.userId, this.userId)));

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

@@ -1,7 +1,8 @@
 import { inArray } from 'drizzle-orm';
 import { and, desc, eq } from 'drizzle-orm/expressions';
 
-import { serverDB } from '@/database/server';
+import { LobeChatDatabase } from '@/database/type';
+
 import {
   agents,
   agentsFiles,
@@ -13,12 +14,15 @@ import {
 
 export class AgentModel {
   private userId: string;
-  constructor(userId: string) {
+  private db: LobeChatDatabase;
+
+  constructor(db: LobeChatDatabase, userId: string) {
     this.userId = userId;
+    this.db = db;
   }
 
   async getAgentConfigById(id: string) {
-    const agent = await serverDB.query.agents.findFirst({ where: eq(agents.id, id) });
+    const agent = await this.db.query.agents.findFirst({ where: eq(agents.id, id) });
 
     const knowledge = await this.getAgentAssignedKnowledge(id);
 
@@ -26,14 +30,14 @@ export class AgentModel {
   }
 
   async getAgentAssignedKnowledge(id: string) {
-    const knowledgeBaseResult = await serverDB
+    const knowledgeBaseResult = await this.db
       .select({ enabled: agentsKnowledgeBases.enabled, knowledgeBases })
       .from(agentsKnowledgeBases)
       .where(eq(agentsKnowledgeBases.agentId, id))
       .orderBy(desc(agentsKnowledgeBases.createdAt))
       .leftJoin(knowledgeBases, eq(knowledgeBases.id, agentsKnowledgeBases.knowledgeBaseId));
 
-    const fileResult = await serverDB
+    const fileResult = await this.db
       .select({ enabled: agentsFiles.enabled, files })
       .from(agentsFiles)
       .where(eq(agentsFiles.agentId, id))
@@ -56,7 +60,7 @@ export class AgentModel {
    * Find agent by session id
    */
   async findBySessionId(sessionId: string) {
-    const item = await serverDB.query.agentsToSessions.findFirst({
+    const item = await this.db.query.agentsToSessions.findFirst({
       where: eq(agentsToSessions.sessionId, sessionId),
     });
     if (!item) return;
@@ -71,7 +75,7 @@ export class AgentModel {
     knowledgeBaseId: string,
     enabled: boolean = true,
   ) => {
-    return serverDB
+    return this.db
       .insert(agentsKnowledgeBases)
       .values({
         agentId,
@@ -83,7 +87,7 @@ export class AgentModel {
   };
 
   deleteAgentKnowledgeBase = async (agentId: string, knowledgeBaseId: string) => {
-    return serverDB
+    return this.db
       .delete(agentsKnowledgeBases)
       .where(
         and(
@@ -96,7 +100,7 @@ export class AgentModel {
   };
 
   toggleKnowledgeBase = async (agentId: string, knowledgeBaseId: string, enabled?: boolean) => {
-    return serverDB
+    return this.db
       .update(agentsKnowledgeBases)
       .set({ enabled })
       .where(
@@ -111,7 +115,7 @@ export class AgentModel {
 
   createAgentFiles = async (agentId: string, fileIds: string[], enabled: boolean = true) => {
     // Exclude the fileIds that already exist in agentsFiles, and then insert them
-    const existingFiles = await serverDB
+    const existingFiles = await this.db
       .select({ id: agentsFiles.fileId })
       .from(agentsFiles)
       .where(
@@ -128,7 +132,7 @@ export class AgentModel {
 
     if (needToInsertFileIds.length === 0) return;
 
-    return serverDB
+    return this.db
       .insert(agentsFiles)
       .values(
         needToInsertFileIds.map((fileId) => ({ agentId, enabled, fileId, userId: this.userId })),
@@ -137,7 +141,7 @@ export class AgentModel {
   };
 
   deleteAgentFile = async (agentId: string, fileId: string) => {
-    return serverDB
+    return this.db
       .delete(agentsFiles)
       .where(
         and(
@@ -150,7 +154,7 @@ export class AgentModel {
   };
 
   toggleFile = async (agentId: string, fileId: string, enabled?: boolean) => {
-    return serverDB
+    return this.db
       .update(agentsFiles)
       .set({ enabled })
       .where(

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

@@ -1,7 +1,7 @@
 import { eq, inArray, lt } from 'drizzle-orm';
 import { and } from 'drizzle-orm/expressions';
 
-import { serverDB } from '@/database/server';
+import { LobeChatDatabase } from '@/database/type';
 import {
   AsyncTaskError,
   AsyncTaskErrorType,
@@ -16,13 +16,15 @@ export const ASYNC_TASK_TIMEOUT = 298 * 1000;
 
 export class AsyncTaskModel {
   private userId: string;
+  private db: LobeChatDatabase;
 
-  constructor(userId: string) {
+  constructor(db: LobeChatDatabase, userId: string) {
     this.userId = userId;
+    this.db = db;
   }
 
   create = async (params: Pick<NewAsyncTaskItem, 'type' | 'status'>): Promise<string> => {
-    const data = await serverDB
+    const data = await this.db
       .insert(asyncTasks)
       .values({ ...params, userId: this.userId })
       .returning();
@@ -31,17 +33,17 @@ export class AsyncTaskModel {
   };
 
   delete = async (id: string) => {
-    return serverDB
+    return this.db
       .delete(asyncTasks)
       .where(and(eq(asyncTasks.id, id), eq(asyncTasks.userId, this.userId)));
   };
 
   findById = async (id: string) => {
-    return serverDB.query.asyncTasks.findFirst({ where: and(eq(asyncTasks.id, id)) });
+    return this.db.query.asyncTasks.findFirst({ where: and(eq(asyncTasks.id, id)) });
   };
 
   update(taskId: string, value: Partial<AsyncTaskSelectItem>) {
-    return serverDB
+    return this.db
       .update(asyncTasks)
       .set({ ...value, updatedAt: new Date() })
       .where(and(eq(asyncTasks.id, taskId)));
@@ -52,7 +54,7 @@ export class AsyncTaskModel {
 
     if (taskIds.length > 0) {
       await this.checkTimeoutTasks(taskIds);
-      chunkTasks = await serverDB.query.asyncTasks.findMany({
+      chunkTasks = await this.db.query.asyncTasks.findMany({
         where: and(inArray(asyncTasks.id, taskIds), eq(asyncTasks.type, type)),
       });
     }
@@ -64,7 +66,7 @@ export class AsyncTaskModel {
    * make the task status to be `error` if the task is not finished in 20 seconds
    */
   async checkTimeoutTasks(ids: string[]) {
-    const tasks = await serverDB
+    const tasks = await this.db
       .select({ id: asyncTasks.id })
       .from(asyncTasks)
       .where(
@@ -76,7 +78,7 @@ export class AsyncTaskModel {
       );
 
     if (tasks.length > 0) {
-      await serverDB
+      await this.db
         .update(asyncTasks)
         .set({
           error: new AsyncTaskError(

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

@@ -2,7 +2,7 @@ import { asc, cosineDistance, count, eq, inArray, sql } from 'drizzle-orm';
 import { and, desc, isNull } from 'drizzle-orm/expressions';
 import { chunk } from 'lodash-es';
 
-import { serverDB } from '@/database/server';
+import { LobeChatDatabase } from '@/database/type';
 import { ChunkMetadata, FileChunk } from '@/types/chunk';
 
 import {
@@ -18,12 +18,17 @@ import {
 export class ChunkModel {
   private userId: string;
 
-  constructor(userId: string) {
+  private db: LobeChatDatabase;
+
+  constructor(db: LobeChatDatabase, userId: string) {
     this.userId = userId;
+    this.db = db;
   }
 
   bulkCreate = async (params: NewChunkItem[], fileId: string) => {
-    return serverDB.transaction(async (trx) => {
+    return this.db.transaction(async (trx) => {
+      if (params.length === 0) return [];
+
       const result = await trx.insert(chunks).values(params).returning();
 
       const fileChunksData = result.map((chunk) => ({ chunkId: chunk.id, fileId }));
@@ -37,15 +42,15 @@ export class ChunkModel {
   };
 
   bulkCreateUnstructuredChunks = async (params: NewUnstructuredChunkItem[]) => {
-    return serverDB.insert(unstructuredChunks).values(params);
+    return this.db.insert(unstructuredChunks).values(params);
   };
 
   delete = async (id: string) => {
-    return serverDB.delete(chunks).where(and(eq(chunks.id, id), eq(chunks.userId, this.userId)));
+    return this.db.delete(chunks).where(and(eq(chunks.id, id), eq(chunks.userId, this.userId)));
   };
 
   deleteOrphanChunks = async () => {
-    const orphanedChunks = await serverDB
+    const orphanedChunks = await this.db
       .select({ chunkId: chunks.id })
       .from(chunks)
       .leftJoin(fileChunks, eq(chunks.id, fileChunks.chunkId))
@@ -56,7 +61,7 @@ export class ChunkModel {
 
     const list = chunk(ids, 500);
 
-    await serverDB.transaction(async (trx) => {
+    await this.db.transaction(async (trx) => {
       await Promise.all(
         list.map(async (chunkIds) => {
           await trx.delete(chunks).where(inArray(chunks.id, chunkIds));
@@ -66,13 +71,13 @@ export class ChunkModel {
   };
 
   findById = async (id: string) => {
-    return serverDB.query.chunks.findFirst({
+    return this.db.query.chunks.findFirst({
       where: and(eq(chunks.id, id)),
     });
   };
 
   async findByFileId(id: string, page = 0) {
-    const data = await serverDB
+    const data = await this.db
       .select({
         abstract: chunks.abstract,
         createdAt: chunks.createdAt,
@@ -98,7 +103,7 @@ export class ChunkModel {
   }
 
   async getChunksTextByFileId(id: string): Promise<{ id: string; text: string }[]> {
-    const data = await serverDB
+    const data = await this.db
       .select()
       .from(chunks)
       .innerJoin(fileChunks, eq(chunks.id, fileChunks.chunkId))
@@ -113,7 +118,7 @@ export class ChunkModel {
   async countByFileIds(ids: string[]) {
     if (ids.length === 0) return [];
 
-    return serverDB
+    return this.db
       .select({
         count: count(fileChunks.chunkId),
         id: fileChunks.fileId,
@@ -124,7 +129,7 @@ export class ChunkModel {
   }
 
   async countByFileId(ids: string) {
-    const data = await serverDB
+    const data = await this.db
       .select({
         count: count(fileChunks.chunkId),
         id: fileChunks.fileId,
@@ -146,7 +151,7 @@ export class ChunkModel {
   }) {
     const similarity = sql<number>`1 - (${cosineDistance(embeddings.embeddings, embedding)})`;
 
-    const data = await serverDB
+    const data = await this.db
       .select({
         fileId: fileChunks.fileId,
         fileName: files.name,
@@ -185,7 +190,7 @@ export class ChunkModel {
 
     if (!hasFiles) return [];
 
-    const result = await serverDB
+    const result = await this.db
       .select({
         fileId: files.id,
         fileName: files.name,

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

@@ -1,19 +1,21 @@
 import { count, eq } from 'drizzle-orm';
 import { and } from 'drizzle-orm/expressions';
 
-import { serverDB } from '@/database/server';
+import { LobeChatDatabase } from '@/database/type';
 
 import { NewEmbeddingsItem, embeddings } from '../../schemas';
 
 export class EmbeddingModel {
   private userId: string;
+  private db: LobeChatDatabase;
 
-  constructor(userId: string) {
+  constructor(db: LobeChatDatabase, userId: string) {
     this.userId = userId;
+    this.db = db;
   }
 
   create = async (value: Omit<NewEmbeddingsItem, 'userId'>) => {
-    const [item] = await serverDB
+    const [item] = await this.db
       .insert(embeddings)
       .values({ ...value, userId: this.userId })
       .returning();
@@ -22,7 +24,7 @@ export class EmbeddingModel {
   };
 
   bulkCreate = async (values: Omit<NewEmbeddingsItem, 'userId'>[]) => {
-    return serverDB
+    return this.db
       .insert(embeddings)
       .values(values.map((item) => ({ ...item, userId: this.userId })))
       .onConflictDoNothing({
@@ -31,25 +33,25 @@ export class EmbeddingModel {
   };
 
   delete = async (id: string) => {
-    return serverDB
+    return this.db
       .delete(embeddings)
       .where(and(eq(embeddings.id, id), eq(embeddings.userId, this.userId)));
   };
 
   query = async () => {
-    return serverDB.query.embeddings.findMany({
+    return this.db.query.embeddings.findMany({
       where: eq(embeddings.userId, this.userId),
     });
   };
 
   findById = async (id: string) => {
-    return serverDB.query.embeddings.findFirst({
+    return this.db.query.embeddings.findFirst({
       where: and(eq(embeddings.id, id), eq(embeddings.userId, this.userId)),
     });
   };
 
   countUsage = async () => {
-    const result = await serverDB
+    const result = await this.db
       .select({
         count: count(),
       })