@lobehub/chat 1.45.17 → 1.46.1

package/src/server/routers/async/file.ts

@@ -5,15 +5,15 @@ import { z } from 'zod';
 
 import { serverDBEnv } from '@/config/db';
 import { fileEnv } from '@/config/file';
-import { DEFAULT_EMBEDDING_MODEL } from '@/const/settings';
+import { DEFAULT_FILE_EMBEDDING_MODEL_ITEM } from '@/const/settings/knowledge';
 import { NewChunkItem, NewEmbeddingsItem } from '@/database/schemas';
 import { serverDB } from '@/database/server';
 import { ASYNC_TASK_TIMEOUT, AsyncTaskModel } from '@/database/server/models/asyncTask';
 import { ChunkModel } from '@/database/server/models/chunk';
 import { EmbeddingModel } from '@/database/server/models/embedding';
 import { FileModel } from '@/database/server/models/file';
-import { ModelProvider } from '@/libs/agent-runtime';
 import { asyncAuthedProcedure, asyncRouter as router } from '@/libs/trpc/async';
+import { getServerDefaultFilesConfig } from '@/server/globalConfig';
 import { initAgentRuntimeWithUserPayload } from '@/server/modules/AgentRuntime';
 import { S3 } from '@/server/modules/S3';
 import { ChunkService } from '@/server/services/chunk';
@@ -44,7 +44,6 @@ export const fileRouter = router({
     .input(
       z.object({
         fileId: z.string(),
-        model: z.string().default(DEFAULT_EMBEDDING_MODEL),
         taskId: z.string(),
       }),
     )
@@ -57,6 +56,9 @@ export const fileRouter = router({
 
       const asyncTask = await ctx.asyncTaskModel.findById(input.taskId);
 
+      const { model, provider } =
+        getServerDefaultFilesConfig().embeddingModel || DEFAULT_FILE_EMBEDDING_MODEL_ITEM;
+
       if (!asyncTask) throw new TRPCError({ code: 'BAD_REQUEST', message: 'Async Task not found' });
 
       try {
@@ -84,13 +86,12 @@ export const fileRouter = router({
 
           const chunks = await ctx.chunkModel.getChunksTextByFileId(input.fileId);
           const requestArray = chunk(chunks, CHUNK_SIZE);
-
           try {
             await pMap(
               requestArray,
               async (chunks, index) => {
                 const agentRuntime = await initAgentRuntimeWithUserPayload(
-                  ModelProvider.OpenAI,
+                  provider,
                   ctx.jwtPayload,
                 );
 
@@ -98,11 +99,10 @@ export const fileRouter = router({
                 console.log(`执行第 ${number} 个任务`);
 
                 console.time(`任务[${number}]: embeddings`);
-
                 const embeddings = await agentRuntime.embeddings({
                   dimensions: 1024,
                   input: chunks.map((c) => c.text),
-                  model: input.model,
+                  model,
                 });
                 console.timeEnd(`任务[${number}]: embeddings`);
 
@@ -111,7 +111,7 @@ export const fileRouter = router({
                     chunkId: chunks[idx].id,
                     embeddings: e,
                     fileId: input.fileId,
-                    model: input.model,
+                    model,
                   })) || [];
 
                 console.time(`任务[${number}]: insert db`);

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

@@ -1,7 +1,7 @@
 import { inArray } from 'drizzle-orm/expressions';
 import { z } from 'zod';
 
-import { DEFAULT_EMBEDDING_MODEL } from '@/const/settings';
+import { DEFAULT_FILE_EMBEDDING_MODEL_ITEM } from '@/const/settings/knowledge';
 import { knowledgeBaseFiles } from '@/database/schemas';
 import { serverDB } from '@/database/server';
 import { AsyncTaskModel } from '@/database/server/models/asyncTask';
@@ -9,9 +9,9 @@ import { ChunkModel } from '@/database/server/models/chunk';
 import { EmbeddingModel } from '@/database/server/models/embedding';
 import { FileModel } from '@/database/server/models/file';
 import { MessageModel } from '@/database/server/models/message';
-import { ModelProvider } from '@/libs/agent-runtime';
 import { authedProcedure, router } from '@/libs/trpc';
 import { keyVaults } from '@/libs/trpc/middleware/keyVaults';
+import { getServerDefaultFilesConfig } from '@/server/globalConfig';
 import { initAgentRuntimeWithUserPayload } from '@/server/modules/AgentRuntime';
 import { ChunkService } from '@/server/services/chunk';
 import { SemanticSearchSchema } from '@/types/rag';
@@ -101,21 +101,18 @@ export const chunkRouter = router({
     .input(
       z.object({
         fileIds: z.array(z.string()).optional(),
-        model: z.string().default(DEFAULT_EMBEDDING_MODEL),
         query: z.string(),
       }),
     )
     .mutation(async ({ ctx, input }) => {
-      console.time('embedding');
-      const agentRuntime = await initAgentRuntimeWithUserPayload(
-        ModelProvider.OpenAI,
-        ctx.jwtPayload,
-      );
+      const { model, provider } =
+        getServerDefaultFilesConfig().embeddingModel || DEFAULT_FILE_EMBEDDING_MODEL_ITEM;
+      const agentRuntime = await initAgentRuntimeWithUserPayload(provider, ctx.jwtPayload);
 
       const embeddings = await agentRuntime.embeddings({
         dimensions: 1024,
         input: input.query,
-        model: input.model,
+        model,
       });
       console.timeEnd('embedding');
 
@@ -130,27 +127,25 @@ export const chunkRouter = router({
     .input(SemanticSearchSchema)
     .mutation(async ({ ctx, input }) => {
       const item = await ctx.messageModel.findMessageQueriesById(input.messageId);
+      const { model, provider } =
+        getServerDefaultFilesConfig().embeddingModel || DEFAULT_FILE_EMBEDDING_MODEL_ITEM;
       let embedding: number[];
       let ragQueryId: string;
-
       // if there is no message rag or it's embeddings, then we need to create one
       if (!item || !item.embeddings) {
         // TODO: need to support customize
-        const agentRuntime = await initAgentRuntimeWithUserPayload(
-          ModelProvider.OpenAI,
-          ctx.jwtPayload,
-        );
+        const agentRuntime = await initAgentRuntimeWithUserPayload(provider, ctx.jwtPayload);
 
         const embeddings = await agentRuntime.embeddings({
           dimensions: 1024,
           input: input.rewriteQuery,
-          model: input.model || DEFAULT_EMBEDDING_MODEL,
+          model,
         });
 
         embedding = embeddings![0];
         const embeddingsId = await ctx.embeddingModel.create({
           embeddings: embedding,
-          model: input.model,
+          model,
         });
 
         const result = await ctx.messageModel.createMessageQuery({
@@ -182,6 +177,7 @@ export const chunkRouter = router({
         fileIds: finalFileIds,
         query: input.rewriteQuery,
       });
+      // TODO: need to rerank the chunks
       console.timeEnd('semanticSearch');
 
       return { chunks, queryId: ragQueryId };

package/src/types/knowledgeBase/index.ts

@@ -1,3 +1,5 @@
+import { FilesConfigItem } from '../user/settings/filesConfig';
+
 export enum KnowledgeBaseTabs {
   Files = 'files',
   Settings = 'Settings',
@@ -43,3 +45,9 @@ export interface KnowledgeItem {
   name: string;
   type: KnowledgeType;
 }
+
+export interface SystemEmbeddingConfig {
+  embeddingModel: FilesConfigItem;
+  queryModel: string;
+  rerankerModel: FilesConfigItem;
+}

package/src/types/user/settings/filesConfig.ts

@@ -0,0 +1,9 @@
+export interface FilesConfigItem {
+  model: string;
+  provider: string;
+}
+export interface FilesConfig {
+  embeddingModel: FilesConfigItem;
+  queryModel: string;
+  rerankerModel: FilesConfigItem;
+}

package/src/types/user/settings/keyVaults.ts

@@ -44,6 +44,7 @@ export interface UserKeyVaults {
   huggingface?: OpenAICompatibleKeyVault;
   hunyuan?: OpenAICompatibleKeyVault;
   internlm?: OpenAICompatibleKeyVault;
+  lmstudio?: OpenAICompatibleKeyVault;
   lobehub?: any;
   minimax?: OpenAICompatibleKeyVault;
   mistral?: OpenAICompatibleKeyVault;