@lobehub/lobehub 2.0.0-next.200 → 2.0.0-next.202

package/src/server/routers/lambda/chunk.ts

@@ -2,7 +2,6 @@ import { DEFAULT_FILE_EMBEDDING_MODEL_ITEM } from '@lobechat/const';
 import {
   type ChatSemanticSearchChunk,
   type FileSearchResult,
-  type ProviderConfig,
   SemanticSearchSchema,
 } from '@lobechat/types';
 import { TRPCError } from '@trpc/server';
@@ -17,13 +16,11 @@ import { DocumentModel } from '@/database/models/document';
 import { EmbeddingModel } from '@/database/models/embedding';
 import { FileModel } from '@/database/models/file';
 import { MessageModel } from '@/database/models/message';
-import { AiInfraRepos } from '@/database/repositories/aiInfra';
 import { knowledgeBaseFiles } from '@/database/schemas';
 import { authedProcedure, router } from '@/libs/trpc/lambda';
 import { keyVaults, serverDatabase } from '@/libs/trpc/lambda/middleware';
-import { getServerDefaultFilesConfig, getServerGlobalConfig } from '@/server/globalConfig';
-import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt';
-import { initModelRuntimeWithUserPayload } from '@/server/modules/ModelRuntime';
+import { getServerDefaultFilesConfig } from '@/server/globalConfig';
+import { initModelRuntimeFromDB } from '@/server/modules/ModelRuntime';
 import { ChunkService } from '@/server/services/chunk';
 import { DocumentService } from '@/server/services/document';
 
@@ -32,15 +29,9 @@ const chunkProcedure = authedProcedure
   .use(keyVaults)
   .use(async (opts) => {
     const { ctx } = opts;
-    const { aiProvider } = await getServerGlobalConfig();
 
     return opts.next({
       ctx: {
-        aiInfraRepos: new AiInfraRepos(
-          ctx.serverDB,
-          ctx.userId,
-          aiProvider as Record<string, ProviderConfig>,
-        ),
         asyncTaskModel: new AsyncTaskModel(ctx.serverDB, ctx.userId),
         chunkModel: new ChunkModel(ctx.serverDB, ctx.userId),
         chunkService: new ChunkService(ctx.serverDB, ctx.userId),
@@ -105,7 +96,7 @@ export const chunkRouter = router({
       }),
     )
     .mutation(async ({ ctx, input }) => {
-      const asyncTaskId = await ctx.chunkService.asyncEmbeddingFileChunks(input.id, ctx.jwtPayload);
+      const asyncTaskId = await ctx.chunkService.asyncEmbeddingFileChunks(input.id);
 
       return { id: asyncTaskId, success: true };
     }),
@@ -118,11 +109,7 @@ export const chunkRouter = router({
       }),
     )
     .mutation(async ({ ctx, input }) => {
-      const asyncTaskId = await ctx.chunkService.asyncParseFileToChunks(
-        input.id,
-        ctx.jwtPayload,
-        input.skipExist,
-      );
+      const asyncTaskId = await ctx.chunkService.asyncParseFileToChunks(input.id, input.skipExist);
 
       return { id: asyncTaskId, success: true };
     }),
@@ -223,7 +210,7 @@ export const chunkRouter = router({
       }
 
       // 2. create a new asyncTask for chunking
-      const asyncTaskId = await ctx.chunkService.asyncParseFileToChunks(input.id, ctx.jwtPayload);
+      const asyncTaskId = await ctx.chunkService.asyncParseFileToChunks(input.id);
 
       return { id: asyncTaskId, success: true };
     }),
@@ -239,7 +226,8 @@ export const chunkRouter = router({
     .mutation(async ({ ctx, input }) => {
       const { model, provider } =
         getServerDefaultFilesConfig().embeddingModel || DEFAULT_FILE_EMBEDDING_MODEL_ITEM;
-      const agentRuntime = await initModelRuntimeWithUserPayload(provider, ctx.jwtPayload);
+      // Read user's provider config from database
+      const agentRuntime = await initModelRuntimeFromDB(ctx.serverDB, ctx.userId, provider);
 
       const embeddings = await agentRuntime.embeddings({
         dimensions: 1024,
@@ -262,15 +250,8 @@ export const chunkRouter = router({
           getServerDefaultFilesConfig().embeddingModel || DEFAULT_FILE_EMBEDDING_MODEL_ITEM;
         let embedding: number[];
 
-        const providerDetail = await ctx.aiInfraRepos.getAiProviderDetail(
-          provider,
-          KeyVaultsGateKeeper.getUserKeyVaults,
-        );
-
-        const modelRuntime = initModelRuntimeWithUserPayload(
-          provider,
-          providerDetail.keyVaults || {},
-        );
+        // Read user's provider config from database
+        const modelRuntime = await initModelRuntimeFromDB(ctx.serverDB, ctx.userId, provider);
 
         // slice content to make sure in the context window limit
         const query = input.query.length > 8000 ? input.query.slice(0, 8000) : input.query;

package/src/server/routers/lambda/image.ts

@@ -228,15 +228,9 @@ export const imageRouter = router({
 
     try {
       log('Creating unified async caller for userId: %s', userId);
-      log(
-        'Lambda context - userId: %s, jwtPayload keys: %O',
-        ctx.userId,
-        Object.keys(ctx.jwtPayload || {}),
-      );
 
-      // 使用统一的 caller 工厂创建 caller
+      // Async router will read keyVaults from DB, no need to pass jwtPayload
       const asyncCaller = await createAsyncCaller({
-        jwtPayload: ctx.jwtPayload,
         userId: ctx.userId,
       });
 

package/src/server/routers/lambda/ragEval.ts

@@ -201,9 +201,9 @@ export const ragEvalRouter = router({
         })),
       );
 
+      // Async router will read keyVaults from DB, no need to pass jwtPayload
       const asyncCaller = await createAsyncCaller({
         userId: ctx.userId,
-        jwtPayload: ctx.jwtPayload,
       });
 
       await ctx.evaluationModel.update(input.id, { status: EvalEvaluationStatus.Processing });

package/src/server/routers/lambda/userMemories/reembed.ts

@@ -49,7 +49,10 @@ export const reembedRouter = router({
     .mutation(async ({ ctx, input }) => {
       try {
         const options = input ?? {};
-        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(ctx.jwtPayload);
+        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(
+          ctx.serverDB,
+          ctx.userId,
+        );
         const concurrency = options.concurrency ?? 10;
         const shouldProcess = (key: ReEmbedTableKey) =>
           !options.only || options.only.length === 0 || options.only.includes(key);

package/src/server/routers/lambda/userMemories/search.ts

@@ -1,18 +1,18 @@
-import { ModelProvider } from 'model-bank';
+import { type LobeChatDatabase } from '@lobechat/database';
 import { type z } from 'zod';
 
 import { DEFAULT_FILE_EMBEDDING_MODEL_ITEM } from '@/const/settings/knowledge';
 import { type UserMemoryModel } from '@/database/models/userMemory';
 import { getServerDefaultFilesConfig } from '@/server/globalConfig';
-import { initModelRuntimeWithUserPayload } from '@/server/modules/ModelRuntime';
-import { type ClientSecretPayload } from '@/types/auth';
+import { initModelRuntimeFromDB } from '@/server/modules/ModelRuntime';
 import { type SearchMemoryResult, searchMemorySchema } from '@/types/userMemory';
 
 import { EMBEDDING_VECTOR_DIMENSION, memoryProcedure, router } from './shared';
 
 type MemorySearchContext = {
-  jwtPayload: ClientSecretPayload;
   memoryModel: UserMemoryModel;
+  serverDB: LobeChatDatabase;
+  userId: string;
 };
 
 type MemorySearchResult = Awaited<ReturnType<UserMemoryModel['searchWithEmbedding']>>;
@@ -80,10 +80,10 @@ export const searchUserMemories = async (
   ctx: MemorySearchContext,
   input: z.infer<typeof searchMemorySchema>,
 ): Promise<SearchMemoryResult> => {
-  const agentRuntime = await initModelRuntimeWithUserPayload(ModelProvider.OpenAI, ctx.jwtPayload);
-
-  const { model: embeddingModel } =
+  const { provider, model: embeddingModel } =
     getServerDefaultFilesConfig().embeddingModel || DEFAULT_FILE_EMBEDDING_MODEL_ITEM;
+  // Read user's provider config from database
+  const agentRuntime = await initModelRuntimeFromDB(ctx.serverDB, ctx.userId, provider);
 
   const queryEmbeddings = await agentRuntime.embeddings({
     dimensions: EMBEDDING_VECTOR_DIMENSION,

package/src/server/routers/lambda/userMemories/shared.ts

@@ -1,15 +1,12 @@
+import { type LobeChatDatabase } from '@lobechat/database';
 import { type SQL, and } from 'drizzle-orm';
-import { ModelProvider } from 'model-bank';
 
 import { DEFAULT_FILE_EMBEDDING_MODEL_ITEM } from '@/const/settings/knowledge';
 import { UserMemoryModel } from '@/database/models/userMemory';
-import { authedProcedure,  } from '@/libs/trpc/lambda';
+import { authedProcedure } from '@/libs/trpc/lambda';
 import { keyVaults, serverDatabase } from '@/libs/trpc/lambda/middleware';
 import { getServerDefaultFilesConfig } from '@/server/globalConfig';
-import { initModelRuntimeWithUserPayload } from '@/server/modules/ModelRuntime';
-import { type ClientSecretPayload } from '@/types/auth';
-
-
+import { initModelRuntimeFromDB } from '@/server/modules/ModelRuntime';
 
 export const EMBEDDING_VECTOR_DIMENSION = 1024;
 
@@ -25,10 +22,11 @@ export const memoryProcedure = authedProcedure
     });
   });
 
-export const getEmbeddingRuntime = async (jwtPayload: ClientSecretPayload) => {
-  const agentRuntime = await initModelRuntimeWithUserPayload(ModelProvider.OpenAI, jwtPayload);
-  const { model: embeddingModel } =
+export const getEmbeddingRuntime = async (serverDB: LobeChatDatabase, userId: string) => {
+  const { provider, model: embeddingModel } =
     getServerDefaultFilesConfig().embeddingModel || DEFAULT_FILE_EMBEDDING_MODEL_ITEM;
+  // Read user's provider config from database
+  const agentRuntime = await initModelRuntimeFromDB(serverDB, userId, provider);
 
   return { agentRuntime, embeddingModel };
 };
@@ -62,4 +60,4 @@ export const normalizeEmbeddable = (value?: string | null): string | undefined =
   return trimmed.length > 0 ? trimmed : undefined;
 };
 
-export {router} from '@/libs/trpc/lambda';
+export { router } from '@/libs/trpc/lambda';

package/src/server/routers/lambda/userMemories/tools.ts

@@ -1,13 +1,18 @@
-import { type IdentityEntryBasePayload, type IdentityEntryPayload, UserMemoryModel } from '@/database/models/userMemory';
-import { searchMemorySchema } from '@/types/userMemory';
 import {
+  AddIdentityActionSchema,
   ContextMemoryItemSchema,
   ExperienceMemoryItemSchema,
   PreferenceMemoryItemSchema,
-  AddIdentityActionSchema,
-  UpdateIdentityActionSchema,
   RemoveIdentityActionSchema,
-} from '@lobechat/memory-user-memory'
+  UpdateIdentityActionSchema,
+} from '@lobechat/memory-user-memory';
+
+import {
+  type IdentityEntryBasePayload,
+  type IdentityEntryPayload,
+  UserMemoryModel,
+} from '@/database/models/userMemory';
+import { searchMemorySchema } from '@/types/userMemory';
 
 import { searchUserMemories } from './search';
 import { createEmbedder, getEmbeddingRuntime, memoryProcedure, router } from './shared';
@@ -17,7 +22,10 @@ export const toolsRouter = router({
     .input(ContextMemoryItemSchema)
     .mutation(async ({ input, ctx }) => {
       try {
-        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(ctx.jwtPayload);
+        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(
+          ctx.serverDB,
+          ctx.userId,
+        );
         const embed = createEmbedder(agentRuntime, embeddingModel);
 
         const summaryEmbedding = await embed(input.summary);
@@ -26,8 +34,10 @@ export const toolsRouter = router({
 
         const { context, memory } = await ctx.memoryModel.createContextMemory({
           context: {
-            associatedObjects: UserMemoryModel.parseAssociatedObjects(input.withContext.associatedObjects) ?? null,
-            associatedSubjects: UserMemoryModel.parseAssociatedSubjects(input.withContext.associatedSubjects) ?? null,
+            associatedObjects:
+              UserMemoryModel.parseAssociatedObjects(input.withContext.associatedObjects) ?? null,
+            associatedSubjects:
+              UserMemoryModel.parseAssociatedSubjects(input.withContext.associatedSubjects) ?? null,
             currentStatus: input.withContext.currentStatus ?? null,
             description: input.withContext.description ?? null,
             descriptionVector: contextDescriptionEmbedding ?? null,
@@ -67,7 +77,10 @@ export const toolsRouter = router({
     .input(ExperienceMemoryItemSchema)
     .mutation(async ({ input, ctx }) => {
       try {
-        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(ctx.jwtPayload);
+        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(
+          ctx.serverDB,
+          ctx.userId,
+        );
         const embed = createEmbedder(agentRuntime, embeddingModel);
 
         const summaryEmbedding = await embed(input.summary);
@@ -120,7 +133,10 @@ export const toolsRouter = router({
     .input(AddIdentityActionSchema)
     .mutation(async ({ input, ctx }) => {
       try {
-        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(ctx.jwtPayload);
+        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(
+          ctx.serverDB,
+          ctx.userId,
+        );
         const embed = createEmbedder(agentRuntime, embeddingModel);
 
         const summaryEmbedding = await embed(input.summary);
@@ -185,7 +201,10 @@ export const toolsRouter = router({
     .input(PreferenceMemoryItemSchema)
     .mutation(async ({ input, ctx }) => {
       try {
-        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(ctx.jwtPayload);
+        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(
+          ctx.serverDB,
+          ctx.userId,
+        );
         const embed = createEmbedder(agentRuntime, embeddingModel);
 
         const summaryEmbedding = await embed(input.summary);
@@ -277,111 +296,114 @@ export const toolsRouter = router({
 
   updateIdentityMemory: memoryProcedure
     .input(UpdateIdentityActionSchema)
-        .mutation(async ({ input, ctx }) => {
-          try {
-            const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(ctx.jwtPayload);
-            const embed = createEmbedder(agentRuntime, embeddingModel);
-
-            let summaryVector1024: number[] | null | undefined;
-            if (input.set.summary !== undefined) {
-              const vector = await embed(input.set.summary);
-              summaryVector1024 = vector ?? null;
-            }
-
-            let detailsVector1024: number[] | null | undefined;
-            if (input.set.details !== undefined) {
-              const vector = await embed(input.set.details);
-              detailsVector1024 = vector ?? null;
-            }
-
-            let descriptionVector: number[] | null | undefined;
-            if (input.set.withIdentity.description !== undefined) {
-              const vector = await embed(input.set.withIdentity.description);
-              descriptionVector = vector ?? null;
-            }
-
-            const metadataUpdates: Record<string, unknown> = {};
-            if (Object.hasOwn(input.set.withIdentity, 'scoreConfidence')) {
-              metadataUpdates.scoreConfidence = input.set.withIdentity.scoreConfidence ?? null;
-            }
-            if (Object.hasOwn(input.set.withIdentity, 'sourceEvidence')) {
-              metadataUpdates.sourceEvidence = input.set.withIdentity.sourceEvidence ?? null;
-            }
-
-            const identityPayload: Partial<IdentityEntryPayload> = {};
-            if (input.set.withIdentity.description !== undefined) {
-              identityPayload.description = input.set.withIdentity.description;
-              identityPayload.descriptionVector = descriptionVector;
-            }
-            if (input.set.withIdentity.episodicDate !== undefined) {
-              identityPayload.episodicDate = input.set.withIdentity.episodicDate;
-            }
-            if (input.set.withIdentity.relationship !== undefined) {
-              identityPayload.relationship = input.set.withIdentity.relationship;
-            }
-            if (input.set.withIdentity.role !== undefined) {
-              identityPayload.role = input.set.withIdentity.role;
-            }
-            if (input.set.tags !== undefined) {
-              identityPayload.tags = input.set.tags;
-            }
-            if (input.set.withIdentity.type !== undefined) {
-              identityPayload.type = input.set.withIdentity.type;
-            }
-            if (Object.keys(metadataUpdates).length > 0) {
-              identityPayload.metadata = metadataUpdates;
-            }
-
-            const basePayload: Partial<IdentityEntryBasePayload> = {};
-            if (input.set.details !== undefined) {
-              basePayload.details = input.set.details;
-              basePayload.detailsVector1024 = detailsVector1024;
-            }
-            if (input.set.memoryCategory !== undefined) {
-              basePayload.memoryCategory = input.set.memoryCategory;
-            }
-            if (input.set.memoryType !== undefined) {
-              basePayload.memoryType = input.set.memoryType;
-            }
-            if (input.set.summary !== undefined) {
-              basePayload.summary = input.set.summary;
-              basePayload.summaryVector1024 = summaryVector1024;
-            }
-            if (input.set.tags !== undefined) {
-              basePayload.tags = input.set.tags;
-            }
-            if (input.set.title !== undefined) {
-              basePayload.title = input.set.title;
-            }
-            if (Object.keys(metadataUpdates).length > 0) {
-              basePayload.metadata = metadataUpdates;
-            }
-
-            const updated = await ctx.memoryModel.updateIdentityEntry({
-              base: Object.keys(basePayload).length > 0 ? basePayload : undefined,
-              identity: Object.keys(identityPayload).length > 0 ? identityPayload : undefined,
-              identityId: input.id,
-              mergeStrategy: input.mergeStrategy,
-            });
-
-            if (!updated) {
-              return {
-                message: 'Identity memory not found',
-                success: false,
-              };
-            }
-
-            return {
-              identityId: input.id,
-              message: 'Identity memory updated successfully',
-              success: true,
-            };
-          } catch (error) {
-            console.error('Failed to update identity memory:', error);
-            return {
-              message: `Failed to update identity memory: ${(error as Error).message}`,
-              success: false,
-            };
-          }
-        }),
+    .mutation(async ({ input, ctx }) => {
+      try {
+        const { agentRuntime, embeddingModel } = await getEmbeddingRuntime(
+          ctx.serverDB,
+          ctx.userId,
+        );
+        const embed = createEmbedder(agentRuntime, embeddingModel);
+
+        let summaryVector1024: number[] | null | undefined;
+        if (input.set.summary !== undefined) {
+          const vector = await embed(input.set.summary);
+          summaryVector1024 = vector ?? null;
+        }
+
+        let detailsVector1024: number[] | null | undefined;
+        if (input.set.details !== undefined) {
+          const vector = await embed(input.set.details);
+          detailsVector1024 = vector ?? null;
+        }
+
+        let descriptionVector: number[] | null | undefined;
+        if (input.set.withIdentity.description !== undefined) {
+          const vector = await embed(input.set.withIdentity.description);
+          descriptionVector = vector ?? null;
+        }
+
+        const metadataUpdates: Record<string, unknown> = {};
+        if (Object.hasOwn(input.set.withIdentity, 'scoreConfidence')) {
+          metadataUpdates.scoreConfidence = input.set.withIdentity.scoreConfidence ?? null;
+        }
+        if (Object.hasOwn(input.set.withIdentity, 'sourceEvidence')) {
+          metadataUpdates.sourceEvidence = input.set.withIdentity.sourceEvidence ?? null;
+        }
+
+        const identityPayload: Partial<IdentityEntryPayload> = {};
+        if (input.set.withIdentity.description !== undefined) {
+          identityPayload.description = input.set.withIdentity.description;
+          identityPayload.descriptionVector = descriptionVector;
+        }
+        if (input.set.withIdentity.episodicDate !== undefined) {
+          identityPayload.episodicDate = input.set.withIdentity.episodicDate;
+        }
+        if (input.set.withIdentity.relationship !== undefined) {
+          identityPayload.relationship = input.set.withIdentity.relationship;
+        }
+        if (input.set.withIdentity.role !== undefined) {
+          identityPayload.role = input.set.withIdentity.role;
+        }
+        if (input.set.tags !== undefined) {
+          identityPayload.tags = input.set.tags;
+        }
+        if (input.set.withIdentity.type !== undefined) {
+          identityPayload.type = input.set.withIdentity.type;
+        }
+        if (Object.keys(metadataUpdates).length > 0) {
+          identityPayload.metadata = metadataUpdates;
+        }
+
+        const basePayload: Partial<IdentityEntryBasePayload> = {};
+        if (input.set.details !== undefined) {
+          basePayload.details = input.set.details;
+          basePayload.detailsVector1024 = detailsVector1024;
+        }
+        if (input.set.memoryCategory !== undefined) {
+          basePayload.memoryCategory = input.set.memoryCategory;
+        }
+        if (input.set.memoryType !== undefined) {
+          basePayload.memoryType = input.set.memoryType;
+        }
+        if (input.set.summary !== undefined) {
+          basePayload.summary = input.set.summary;
+          basePayload.summaryVector1024 = summaryVector1024;
+        }
+        if (input.set.tags !== undefined) {
+          basePayload.tags = input.set.tags;
+        }
+        if (input.set.title !== undefined) {
+          basePayload.title = input.set.title;
+        }
+        if (Object.keys(metadataUpdates).length > 0) {
+          basePayload.metadata = metadataUpdates;
+        }
+
+        const updated = await ctx.memoryModel.updateIdentityEntry({
+          base: Object.keys(basePayload).length > 0 ? basePayload : undefined,
+          identity: Object.keys(identityPayload).length > 0 ? identityPayload : undefined,
+          identityId: input.id,
+          mergeStrategy: input.mergeStrategy,
+        });
+
+        if (!updated) {
+          return {
+            message: 'Identity memory not found',
+            success: false,
+          };
+        }
+
+        return {
+          identityId: input.id,
+          message: 'Identity memory updated successfully',
+          success: true,
+        };
+      } catch (error) {
+        console.error('Failed to update identity memory:', error);
+        return {
+          message: `Failed to update identity memory: ${(error as Error).message}`,
+          success: false,
+        };
+      }
+    }),
 });

package/src/server/routers/lambda/userMemories.test.ts

@@ -4,7 +4,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { getServerDB } from '@/database/core/db-adaptor';
 import { UserMemoryModel } from '@/database/models/userMemory';
-import { initModelRuntimeWithUserPayload } from '@/server/modules/ModelRuntime';
+import { initModelRuntimeFromDB } from '@/server/modules/ModelRuntime';
 
 import { userMemoriesRouter } from './userMemories';
 
@@ -19,11 +19,7 @@ vi.mock('@/server/globalConfig', () => ({
 }));
 
 vi.mock('@/server/modules/ModelRuntime', () => ({
-  initModelRuntimeWithUserPayload: vi.fn(),
-}));
-
-vi.mock('@lobechat/utils/server', () => ({
-  getXorPayload: vi.fn().mockReturnValue({ userId: 'test-user' }),
+  initModelRuntimeFromDB: vi.fn(),
 }));
 
 vi.mock('@/database/models/userMemory', async (importOriginal) => {
@@ -47,7 +43,7 @@ beforeEach(() => {
     return items.map((_, index) => [index + 1]);
   });
 
-  vi.mocked(initModelRuntimeWithUserPayload).mockResolvedValue({
+  vi.mocked(initModelRuntimeFromDB).mockResolvedValue({
     embeddings: embeddingsMock,
   } as any);
 });