🐛 fix(db): desktop local db can't vectorization (#8830)

src/app/(backend)/api/webhooks/casdoor/route.ts

@@ -1,6 +1,7 @@
 import { NextResponse } from 'next/server';
 
 import { authEnv } from '@/config/auth';
+import { serverDB } from '@/database/server';
 import { pino } from '@/libs/logger';
 import { NextAuthUserService } from '@/server/services/nextAuthUser';
 
@@ -18,7 +19,7 @@ export const POST = async (req: Request): Promise<NextResponse> => {
 
   const { action, object } = payload;
 
-  const nextAuthUserService = new NextAuthUserService();
+  const nextAuthUserService = new NextAuthUserService(serverDB);
   switch (action) {
     case 'update-user': {
       return nextAuthUserService.safeUpdateUser(

src/app/(backend)/api/webhooks/clerk/route.ts

@@ -2,6 +2,7 @@ import { NextResponse } from 'next/server';
 
 import { authEnv } from '@/config/auth';
 import { isServerMode } from '@/const/version';
+import { serverDB } from '@/database/server';
 import { pino } from '@/libs/logger';
 import { UserService } from '@/server/services/user';
 
@@ -25,7 +26,7 @@ export const POST = async (req: Request): Promise<NextResponse> => {
 
   pino.trace(`clerk webhook payload: ${{ data, type }}`);
 
-  const userService = new UserService();
+  const userService = new UserService(serverDB);
   switch (type) {
     case 'user.created': {
       pino.info('creating user due to clerk webhook');

src/app/(backend)/api/webhooks/logto/route.ts

@@ -1,6 +1,7 @@
 import { NextResponse } from 'next/server';
 
 import { authEnv } from '@/config/auth';
+import { serverDB } from '@/database/server';
 import { pino } from '@/libs/logger';
 import { NextAuthUserService } from '@/server/services/nextAuthUser';
 
@@ -20,7 +21,7 @@ export const POST = async (req: Request): Promise<NextResponse> => {
 
   pino.trace(`logto webhook payload: ${{ data, event }}`);
 
-  const nextAuthUserService = new NextAuthUserService();
+  const nextAuthUserService = new NextAuthUserService(serverDB);
   switch (event) {
     case 'User.Data.Updated': {
       return nextAuthUserService.safeUpdateUser(

src/app/(backend)/webapi/user/avatar/[id]/[image]/route.ts

@@ -1,3 +1,4 @@
+import { serverDB } from '@/database/server';
 import { UserService } from '@/server/services/user';
 
 export const runtime = 'nodejs';
@@ -31,7 +32,7 @@ export const GET = async (req: Request, segmentData: { params: Params }) => {
   try {
     const params = await segmentData.params;
     const type = getContentType(params.image);
-    const userService = new UserService();
+    const userService = new UserService(serverDB);
 
     const userAvatar = await userService.getUserAvatar(params.id, params.image);
     if (!userAvatar) {

src/database/server/models/ragEval/dataset.ts

@@ -1,18 +1,20 @@
 import { and, desc, eq } from 'drizzle-orm';
 
 import { NewEvalDatasetsItem, evalDatasets } from '@/database/schemas';
-import { serverDB } from '@/database/server';
+import { LobeChatDatabase } from '@/database/type';
 import { RAGEvalDataSetItem } from '@/types/eval';
 
 export class EvalDatasetModel {
   private userId: string;
+  private db: LobeChatDatabase;
 
-  constructor(userId: string) {
+  constructor(db: LobeChatDatabase, userId: string) {
+    this.db = db;
     this.userId = userId;
   }
 
   create = async (params: NewEvalDatasetsItem) => {
-    const [result] = await serverDB
+    const [result] = await this.db
       .insert(evalDatasets)
       .values({ ...params, userId: this.userId })
       .returning();
@@ -20,13 +22,13 @@ export class EvalDatasetModel {
   };
 
   delete = async (id: number) => {
-    return serverDB
+    return this.db
       .delete(evalDatasets)
       .where(and(eq(evalDatasets.id, id), eq(evalDatasets.userId, this.userId)));
   };
 
   query = async (knowledgeBaseId: string): Promise<RAGEvalDataSetItem[]> => {
-    return serverDB
+    return this.db
       .select({
         createdAt: evalDatasets.createdAt,
         description: evalDatasets.description,
@@ -45,13 +47,13 @@ export class EvalDatasetModel {
   };
 
   findById = async (id: number) => {
-    return serverDB.query.evalDatasets.findFirst({
+    return this.db.query.evalDatasets.findFirst({
       where: and(eq(evalDatasets.id, id), eq(evalDatasets.userId, this.userId)),
     });
   };
 
   update = async (id: number, value: Partial<NewEvalDatasetsItem>) => {
-    return serverDB
+    return this.db
       .update(evalDatasets)
       .set({ ...value, updatedAt: new Date() })
       .where(and(eq(evalDatasets.id, id), eq(evalDatasets.userId, this.userId)));

src/database/server/models/ragEval/datasetRecord.ts

@@ -1,18 +1,20 @@
 import { and, eq, inArray } from 'drizzle-orm';
 
 import { NewEvalDatasetRecordsItem, evalDatasetRecords, files } from '@/database/schemas';
-import { serverDB } from '@/database/server';
+import { LobeChatDatabase } from '@/database/type';
 import { EvalDatasetRecordRefFile } from '@/types/eval';
 
 export class EvalDatasetRecordModel {
   private userId: string;
+  private db: LobeChatDatabase;
 
-  constructor(userId: string) {
+  constructor(db: LobeChatDatabase, userId: string) {
+    this.db = db;
     this.userId = userId;
   }
 
   create = async (params: NewEvalDatasetRecordsItem) => {
-    const [result] = await serverDB
+    const [result] = await this.db
       .insert(evalDatasetRecords)
       .values({ ...params, userId: this.userId })
       .returning();
@@ -20,7 +22,7 @@ export class EvalDatasetRecordModel {
   };
 
   batchCreate = async (params: NewEvalDatasetRecordsItem[]) => {
-    const [result] = await serverDB
+    const [result] = await this.db
       .insert(evalDatasetRecords)
       .values(params.map((item) => ({ ...item, userId: this.userId })))
       .returning();
@@ -29,13 +31,13 @@ export class EvalDatasetRecordModel {
   };
 
   delete = async (id: number) => {
-    return serverDB
+    return this.db
       .delete(evalDatasetRecords)
       .where(and(eq(evalDatasetRecords.id, id), eq(evalDatasetRecords.userId, this.userId)));
   };
 
   query = async (datasetId: number) => {
-    const list = await serverDB.query.evalDatasetRecords.findMany({
+    const list = await this.db.query.evalDatasetRecords.findMany({
       where: and(
         eq(evalDatasetRecords.datasetId, datasetId),
         eq(evalDatasetRecords.userId, this.userId),
@@ -43,7 +45,7 @@ export class EvalDatasetRecordModel {
     });
     const fileList = list.flatMap((item) => item.referenceFiles).filter(Boolean) as string[];
 
-    const fileItems = await serverDB
+    const fileItems = await this.db
       .select({ fileType: files.fileType, id: files.id, name: files.name })
       .from(files)
       .where(and(inArray(files.id, fileList), eq(files.userId, this.userId)));
@@ -59,7 +61,7 @@ export class EvalDatasetRecordModel {
   };
 
   findByDatasetId = async (datasetId: number) => {
-    return serverDB.query.evalDatasetRecords.findMany({
+    return this.db.query.evalDatasetRecords.findMany({
       where: and(
         eq(evalDatasetRecords.datasetId, datasetId),
         eq(evalDatasetRecords.userId, this.userId),
@@ -68,13 +70,13 @@ export class EvalDatasetRecordModel {
   };
 
   findById = async (id: number) => {
-    return serverDB.query.evalDatasetRecords.findFirst({
+    return this.db.query.evalDatasetRecords.findFirst({
       where: and(eq(evalDatasetRecords.id, id), eq(evalDatasetRecords.userId, this.userId)),
     });
   };
 
   update = async (id: number, value: Partial<NewEvalDatasetRecordsItem>) => {
-    return serverDB
+    return this.db
       .update(evalDatasetRecords)
       .set(value)
       .where(and(eq(evalDatasetRecords.id, id), eq(evalDatasetRecords.userId, this.userId)));

src/database/server/models/ragEval/evaluation.ts

@@ -6,18 +6,20 @@ import {
   evalEvaluation,
   evaluationRecords,
 } from '@/database/schemas';
-import { serverDB } from '@/database/server';
+import { LobeChatDatabase } from '@/database/type';
 import { EvalEvaluationStatus, RAGEvalEvaluationItem } from '@/types/eval';
 
 export class EvalEvaluationModel {
   private userId: string;
+  private db: LobeChatDatabase;
 
-  constructor(userId: string) {
+  constructor(db: LobeChatDatabase, userId: string) {
+    this.db = db;
     this.userId = userId;
   }
 
   create = async (params: NewEvalEvaluationItem) => {
-    const [result] = await serverDB
+    const [result] = await this.db
       .insert(evalEvaluation)
       .values({ ...params, userId: this.userId })
       .returning();
@@ -25,13 +27,13 @@ export class EvalEvaluationModel {
   };
 
   delete = async (id: number) => {
-    return serverDB
+    return this.db
       .delete(evalEvaluation)
       .where(and(eq(evalEvaluation.id, id), eq(evalEvaluation.userId, this.userId)));
   };
 
   queryByKnowledgeBaseId = async (knowledgeBaseId: string) => {
-    const evaluations = await serverDB
+    const evaluations = await this.db
       .select({
         createdAt: evalEvaluation.createdAt,
         dataset: {
@@ -57,7 +59,7 @@ export class EvalEvaluationModel {
     // 然后查询每个评估的记录统计
     const evaluationIds = evaluations.map((evals) => evals.id);
 
-    const recordStats = await serverDB
+    const recordStats = await this.db
       .select({
         evaluationId: evaluationRecords.evaluationId,
         success: count(evaluationRecords.status).if(
@@ -82,13 +84,13 @@ export class EvalEvaluationModel {
   };
 
   findById = async (id: number) => {
-    return serverDB.query.evalEvaluation.findFirst({
+    return this.db.query.evalEvaluation.findFirst({
       where: and(eq(evalEvaluation.id, id), eq(evalEvaluation.userId, this.userId)),
     });
   };
 
   update = async (id: number, value: Partial<NewEvalEvaluationItem>) => {
-    return serverDB
+    return this.db
       .update(evalEvaluation)
       .set(value)
       .where(and(eq(evalEvaluation.id, id), eq(evalEvaluation.userId, this.userId)));

src/database/server/models/ragEval/evaluationRecord.ts

@@ -1,17 +1,19 @@
 import { and, eq } from 'drizzle-orm';
 
 import { NewEvaluationRecordsItem, evaluationRecords } from '@/database/schemas';
-import { serverDB } from '@/database/server';
+import { LobeChatDatabase } from '@/database/type';
 
 export class EvaluationRecordModel {
   private userId: string;
+  private db: LobeChatDatabase;
 
-  constructor(userId: string) {
+  constructor(db: LobeChatDatabase, userId: string) {
+    this.db = db;
     this.userId = userId;
   }
 
   create = async (params: NewEvaluationRecordsItem) => {
-    const [result] = await serverDB
+    const [result] = await this.db
       .insert(evaluationRecords)
       .values({ ...params, userId: this.userId })
       .returning();
@@ -19,20 +21,20 @@ export class EvaluationRecordModel {
   };
 
   batchCreate = async (params: NewEvaluationRecordsItem[]) => {
-    return serverDB
+    return this.db
       .insert(evaluationRecords)
       .values(params.map((item) => ({ ...item, userId: this.userId })))
       .returning();
   };
 
   delete = async (id: number) => {
-    return serverDB
+    return this.db
       .delete(evaluationRecords)
       .where(and(eq(evaluationRecords.id, id), eq(evaluationRecords.userId, this.userId)));
   };
 
   query = async (reportId: number) => {
-    return serverDB.query.evaluationRecords.findMany({
+    return this.db.query.evaluationRecords.findMany({
       where: and(
         eq(evaluationRecords.evaluationId, reportId),
         eq(evaluationRecords.userId, this.userId),
@@ -41,13 +43,13 @@ export class EvaluationRecordModel {
   };
 
   findById = async (id: number) => {
-    return serverDB.query.evaluationRecords.findFirst({
+    return this.db.query.evaluationRecords.findFirst({
       where: and(eq(evaluationRecords.id, id), eq(evaluationRecords.userId, this.userId)),
     });
   };
 
   findByEvaluationId = async (evaluationId: number) => {
-    return serverDB.query.evaluationRecords.findMany({
+    return this.db.query.evaluationRecords.findMany({
       where: and(
         eq(evaluationRecords.evaluationId, evaluationId),
         eq(evaluationRecords.userId, this.userId),
@@ -56,7 +58,7 @@ export class EvaluationRecordModel {
   };
 
   update = async (id: number, value: Partial<NewEvaluationRecordsItem>) => {
-    return serverDB
+    return this.db
       .update(evaluationRecords)
       .set(value)
       .where(and(eq(evaluationRecords.id, id), eq(evaluationRecords.userId, this.userId)));

src/server/routers/async/file.ts

@@ -32,7 +32,7 @@ const fileProcedure = asyncAuthedProcedure.use(async (opts) => {
     ctx: {
       asyncTaskModel: new AsyncTaskModel(ctx.serverDB, ctx.userId),
       chunkModel: new ChunkModel(ctx.serverDB, ctx.userId),
-      chunkService: new ChunkService(ctx.userId),
+      chunkService: new ChunkService(ctx.serverDB, ctx.userId),
       embeddingModel: new EmbeddingModel(ctx.serverDB, ctx.userId),
       fileModel: new FileModel(ctx.serverDB, ctx.userId),
       fileService: new FileService(ctx.serverDB, ctx.userId),

src/server/routers/async/ragEval.ts

@@ -25,11 +25,11 @@ const ragEvalProcedure = asyncAuthedProcedure.use(async (opts) => {
   return opts.next({
     ctx: {
       chunkModel: new ChunkModel(ctx.serverDB, ctx.userId),
-      chunkService: new ChunkService(ctx.userId),
-      datasetRecordModel: new EvalDatasetRecordModel(ctx.userId),
+      chunkService: new ChunkService(ctx.serverDB, ctx.userId),
+      datasetRecordModel: new EvalDatasetRecordModel(ctx.serverDB, ctx.userId),
       embeddingModel: new EmbeddingModel(ctx.serverDB, ctx.userId),
-      evalRecordModel: new EvaluationRecordModel(ctx.userId),
-      evaluationModel: new EvalEvaluationModel(ctx.userId),
+      evalRecordModel: new EvaluationRecordModel(ctx.serverDB, ctx.userId),
+      evaluationModel: new EvalEvaluationModel(ctx.serverDB, ctx.userId),
       fileModel: new FileModel(ctx.serverDB, ctx.userId),
     },
   });

src/server/routers/lambda/chunk.ts

@@ -26,7 +26,7 @@ const chunkProcedure = authedProcedure
       ctx: {
         asyncTaskModel: new AsyncTaskModel(ctx.serverDB, ctx.userId),
         chunkModel: new ChunkModel(ctx.serverDB, ctx.userId),
-        chunkService: new ChunkService(ctx.userId),
+        chunkService: new ChunkService(ctx.serverDB, ctx.userId),
         embeddingModel: new EmbeddingModel(ctx.serverDB, ctx.userId),
         fileModel: new FileModel(ctx.serverDB, ctx.userId),
         messageModel: new MessageModel(ctx.serverDB, ctx.userId),

src/server/routers/lambda/ragEval.ts

@@ -35,11 +35,11 @@ const ragEvalProcedure = authedProcedure
 
     return opts.next({
       ctx: {
-        datasetModel: new EvalDatasetModel(ctx.userId),
+        datasetModel: new EvalDatasetModel(ctx.serverDB, ctx.userId),
         fileModel: new FileModel(ctx.serverDB, ctx.userId),
-        datasetRecordModel: new EvalDatasetRecordModel(ctx.userId),
-        evaluationModel: new EvalEvaluationModel(ctx.userId),
-        evaluationRecordModel: new EvaluationRecordModel(ctx.userId),
+        datasetRecordModel: new EvalDatasetRecordModel(ctx.serverDB, ctx.userId),
+        evaluationModel: new EvalEvaluationModel(ctx.serverDB, ctx.userId),
+        evaluationRecordModel: new EvaluationRecordModel(ctx.serverDB, ctx.userId),
         fileService: new FileService(ctx.serverDB, ctx.userId),
       },
     });

src/server/routers/lambda/user.ts

@@ -58,7 +58,7 @@ export const userRouter = router({
           if (enableClerk) {
             const user = await ctx.clerkAuth.getCurrentUser();
             if (user) {
-              const userService = new UserService();
+              const userService = new UserService(ctx.serverDB);
 
               await userService.createUser(user.id, {
                 created_at: user.createdAt,

src/server/services/chunk/index.ts

@@ -1,7 +1,7 @@
 import { ClientSecretPayload } from '@/const/auth';
 import { AsyncTaskModel } from '@/database/models/asyncTask';
 import { FileModel } from '@/database/models/file';
-import { serverDB } from '@/database/server';
+import { LobeChatDatabase } from '@/database/type';
 import { ChunkContentParams, ContentChunk } from '@/server/modules/ContentChunk';
 import { createAsyncCaller } from '@/server/routers/async';
 import {
@@ -17,7 +17,7 @@ export class ChunkService {
   private fileModel: FileModel;
   private asyncTaskModel: AsyncTaskModel;
 
-  constructor(userId: string) {
+  constructor(serverDB: LobeChatDatabase, userId: string) {
     this.userId = userId;
 
     this.chunkClient = new ContentChunk();

src/server/services/nextAuthUser/index.test.ts

@@ -24,7 +24,7 @@ describe('NextAuthUserService', () => {
 
   beforeEach(async () => {
     vi.clearAllMocks();
-    service = new NextAuthUserService();
+    service = new NextAuthUserService(serverDB);
   });
 
   describe('safeUpdateUser', () => {

src/server/services/nextAuthUser/index.ts

@@ -2,15 +2,17 @@ import { NextResponse } from 'next/server';
 
 import { UserModel } from '@/database/models/user';
 import { UserItem } from '@/database/schemas';
-import { serverDB } from '@/database/server';
+import { LobeChatDatabase } from '@/database/type';
 import { pino } from '@/libs/logger';
 import { LobeNextAuthDbAdapter } from '@/libs/next-auth/adapter';
 
 export class NextAuthUserService {
   adapter;
+  private db: LobeChatDatabase;
 
-  constructor() {
-    this.adapter = LobeNextAuthDbAdapter(serverDB);
+  constructor(db: LobeChatDatabase) {
+    this.db = db;
+    this.adapter = LobeNextAuthDbAdapter(db);
   }
 
   safeUpdateUser = async (
@@ -27,7 +29,7 @@ export class NextAuthUserService {
 
     // 2. If found, Update user data from provider
     if (user?.id) {
-      const userModel = new UserModel(serverDB, user.id);
+      const userModel = new UserModel(this.db, user.id);
 
       // Perform update
       await userModel.updateUser({

src/server/services/user/index.test.ts

@@ -3,6 +3,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { UserModel } from '@/database/models/user';
 import { UserItem } from '@/database/schemas';
+import { LobeChatDatabase } from '@/database/type';
 import { pino } from '@/libs/logger';
 import { AgentService } from '@/server/services/agent';
 
@@ -46,6 +47,7 @@ vi.mock('@/server/services/agent', () => ({
 
 let service: UserService;
 const mockUserId = 'test-user-id';
+const mockDB = {} as LobeChatDatabase;
 
 // Mock user data
 const mockUserJSON: UserJSON = {
@@ -62,7 +64,7 @@ const mockUserJSON: UserJSON = {
 } as unknown as UserJSON;
 
 beforeEach(() => {
-  service = new UserService();
+  service = new UserService(mockDB);
   vi.clearAllMocks();
 });
 

src/server/services/user/index.ts

@@ -1,7 +1,7 @@
 import { UserJSON } from '@clerk/backend';
 
 import { UserModel } from '@/database/models/user';
-import { serverDB } from '@/database/server';
+import { LobeChatDatabase } from '@/database/type';
 import { initializeServerAnalytics } from '@/libs/analytics';
 import { pino } from '@/libs/logger';
 import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt';
@@ -9,9 +9,15 @@ import { S3 } from '@/server/modules/S3';
 import { AgentService } from '@/server/services/agent';
 
 export class UserService {
+  private db: LobeChatDatabase;
+
+  constructor(db: LobeChatDatabase) {
+    this.db = db;
+  }
+
   createUser = async (id: string, params: UserJSON) => {
     // Check if user already exists
-    const res = await UserModel.findById(serverDB, id);
+    const res = await UserModel.findById(this.db, id);
 
     // If user already exists, skip creating a new user
     if (res)
@@ -33,7 +39,7 @@ export class UserService {
     /* ↑ cloud slot ↑ */
 
     // 2. create user in database
-    await UserModel.createUser(serverDB, {
+    await UserModel.createUser(this.db, {
       avatar: params.image_url,
       clerkCreatedAt: new Date(params.created_at),
       email: email?.email_address,
@@ -45,7 +51,7 @@ export class UserService {
     });
 
     // 3. Create an inbox session for the user
-    const agentService = new AgentService(serverDB, id);
+    const agentService = new AgentService(this.db, id);
     await agentService.createInbox();
 
     /* ↓ cloud slot ↓ */
@@ -73,14 +79,14 @@ export class UserService {
   };
 
   deleteUser = async (id: string) => {
-    await UserModel.deleteUser(serverDB, id);
+    await UserModel.deleteUser(this.db, id);
   };
 
   updateUser = async (id: string, params: UserJSON) => {
-    const userModel = new UserModel(serverDB, id);
+    const userModel = new UserModel(this.db, id);
 
     // Check if user already exists
-    const res = await UserModel.findById(serverDB, id);
+    const res = await UserModel.findById(this.db, id);
 
     // If user not exists, skip update the user
     if (!res)
@@ -111,7 +117,7 @@ export class UserService {
   };
 
   getUserApiKeys = async (id: string) => {
-    return UserModel.getUserApiKeys(serverDB, id, KeyVaultsGateKeeper.getUserKeyVaults);
+    return UserModel.getUserApiKeys(this.db, id, KeyVaultsGateKeeper.getUserKeyVaults);
   };
 
   getUserAvatar = async (id: string, image: string) => {