@lobehub/chat 1.45.17 → 1.46.1

package/src/config/modelProviders/lmstudio.ts

@@ -0,0 +1,25 @@
+import { ModelProviderCard } from '@/types/llm';
+
+// ref: https://ollama.com/library
+const LMStudio: ModelProviderCard = {
+  chatModels: [],
+  description: 'LM Studio 是一个用于在您的计算机上开发和实验 LLMs 的桌面应用程序。',
+  id: 'lmstudio',
+  modelsUrl: 'https://lmstudio.ai/models',
+  name: 'LM Studio',
+  settings: {
+    defaultShowBrowserRequest: true,
+    proxyUrl: {
+      placeholder: 'http://127.0.0.1:1234/v1',
+    },
+    showApiKey: false,
+    showModelFetcher: true,
+    smoothing: {
+      speed: 2,
+      text: true,
+    },
+  },
+  url: 'https://lmstudio.ai',
+};
+
+export default LMStudio;

package/src/const/settings/knowledge.ts

@@ -0,0 +1,25 @@
+import { FilesConfig, FilesConfigItem } from '@/types/user/settings/filesConfig';
+
+import {
+  DEFAULT_EMBEDDING_MODEL,
+  DEFAULT_PROVIDER,
+  DEFAULT_RERANK_MODEL,
+  DEFAULT_RERANK_PROVIDER,
+  DEFAULT_RERANK_QUERY_MODE,
+} from './llm';
+
+export const DEFAULT_FILE_EMBEDDING_MODEL_ITEM: FilesConfigItem = {
+  model: DEFAULT_EMBEDDING_MODEL,
+  provider: DEFAULT_PROVIDER,
+};
+
+export const DEFAULT_FILE_RERANK_MODEL_ITEM: FilesConfigItem = {
+  model: DEFAULT_RERANK_MODEL,
+  provider: DEFAULT_RERANK_PROVIDER,
+};
+
+export const DEFAULT_FILES_CONFIG: FilesConfig = {
+  embeddingModel: DEFAULT_FILE_EMBEDDING_MODEL_ITEM,
+  queryModel: DEFAULT_RERANK_QUERY_MODE,
+  rerankerModel: DEFAULT_FILE_RERANK_MODEL_ITEM,
+};

package/src/const/settings/llm.ts

@@ -2,6 +2,9 @@ import { ModelProvider } from '@/libs/agent-runtime';
 import { genUserLLMConfig } from '@/utils/genUserLLMConfig';
 
 export const DEFAULT_LLM_CONFIG = genUserLLMConfig({
+  lmstudio: {
+    fetchOnClient: true,
+  },
   ollama: {
     enabled: true,
     fetchOnClient: true,
@@ -12,6 +15,12 @@ export const DEFAULT_LLM_CONFIG = genUserLLMConfig({
 });
 
 export const DEFAULT_MODEL = 'gpt-4o-mini';
+
 export const DEFAULT_EMBEDDING_MODEL = 'text-embedding-3-small';
+export const DEFAULT_EMBEDDING_PROVIDER = ModelProvider.OpenAI;
+
+export const DEFAULT_RERANK_MODEL = 'rerank-english-v3.0';
+export const DEFAULT_RERANK_PROVIDER = 'cohere';
+export const DEFAULT_RERANK_QUERY_MODE = 'full_text';
 
 export const DEFAULT_PROVIDER = ModelProvider.OpenAI;

package/src/database/schemas/ragEvals.ts

@@ -1,7 +1,7 @@
 /* eslint-disable sort-keys-fix/sort-keys-fix  */
 import { integer, jsonb, pgTable, text, uuid } from 'drizzle-orm/pg-core';
 
-import { DEFAULT_EMBEDDING_MODEL, DEFAULT_MODEL } from '@/const/settings';
+import { DEFAULT_MODEL } from '@/const/settings';
 import { EvalEvaluationStatus } from '@/types/eval';
 
 import { timestamps } from './_helpers';
@@ -60,7 +60,7 @@ export const evalEvaluation = pgTable('rag_eval_evaluations', {
     onDelete: 'cascade',
   }),
   languageModel: text('language_model').$defaultFn(() => DEFAULT_MODEL),
-  embeddingModel: text('embedding_model').$defaultFn(() => DEFAULT_EMBEDDING_MODEL),
+  embeddingModel: text('embedding_model'),
 
   userId: text('user_id').references(() => users.id, { onDelete: 'cascade' }),
   ...timestamps,

package/src/libs/agent-runtime/AgentRuntime.ts

@@ -20,6 +20,7 @@ import { LobeHigressAI } from './higress';
 import { LobeHuggingFaceAI } from './huggingface';
 import { LobeHunyuanAI } from './hunyuan';
 import { LobeInternLMAI } from './internlm';
+import { LobeLMStudioAI } from './lmstudio';
 import { LobeMinimaxAI } from './minimax';
 import { LobeMistralAI } from './mistral';
 import { LobeMoonshotAI } from './moonshot';
@@ -147,6 +148,7 @@ class AgentRuntime {
       huggingface: { apiKey?: string; baseURL?: string };
       hunyuan: Partial<ClientOptions>;
       internlm: Partial<ClientOptions>;
+      lmstudio: Partial<ClientOptions>;
       minimax: Partial<ClientOptions>;
       mistral: Partial<ClientOptions>;
       moonshot: Partial<ClientOptions>;
@@ -207,6 +209,11 @@ class AgentRuntime {
         break;
       }
 
+      case ModelProvider.LMStudio: {
+        runtimeModel = new LobeLMStudioAI(params.lmstudio);
+        break;
+      }
+
       case ModelProvider.Ollama: {
         runtimeModel = new LobeOllamaAI(params.ollama);
         break;

package/src/libs/agent-runtime/bedrock/index.ts

@@ -1,12 +1,20 @@
 import {
   BedrockRuntimeClient,
+  InvokeModelCommand,
   InvokeModelWithResponseStreamCommand,
 } from '@aws-sdk/client-bedrock-runtime';
 import { experimental_buildLlama2Prompt } from 'ai/prompts';
 
 import { LobeRuntimeAI } from '../BaseAI';
 import { AgentRuntimeErrorType } from '../error';
-import { ChatCompetitionOptions, ChatStreamPayload, ModelProvider } from '../types';
+import {
+  ChatCompetitionOptions,
+  ChatStreamPayload,
+  Embeddings,
+  EmbeddingsOptions,
+  EmbeddingsPayload,
+  ModelProvider,
+} from '../types';
 import { buildAnthropicMessages, buildAnthropicTools } from '../utils/anthropicHelpers';
 import { AgentRuntimeError } from '../utils/createError';
 import { debugStream } from '../utils/debugStream';
@@ -32,9 +40,7 @@ export class LobeBedrockAI implements LobeRuntimeAI {
   constructor({ region, accessKeyId, accessKeySecret, sessionToken }: LobeBedrockAIParams = {}) {
     if (!(accessKeyId && accessKeySecret))
       throw AgentRuntimeError.createError(AgentRuntimeErrorType.InvalidBedrockCredentials);
-
     this.region = region ?? 'us-east-1';
-
     this.client = new BedrockRuntimeClient({
       credentials: {
         accessKeyId: accessKeyId,
@@ -50,6 +56,61 @@ export class LobeBedrockAI implements LobeRuntimeAI {
 
     return this.invokeClaudeModel(payload, options);
   }
+  /**
+   * Supports the Amazon Titan Text models series.
+   * Cohere Embed models are not supported
+   * because the current text size per request
+   * exceeds the maximum 2048 characters limit
+   * for a single request for this series of models.
+   * [bedrock embed guide] https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-embed.html
+   */
+  async embeddings(payload: EmbeddingsPayload, options?: EmbeddingsOptions): Promise<Embeddings[]> {
+    const input = Array.isArray(payload.input) ? payload.input : [payload.input];
+    const promises = input.map((inputText: string) =>
+      this.invokeEmbeddingModel(
+        {
+          dimensions: payload.dimensions,
+          input: inputText,
+          model: payload.model,
+        },
+        options,
+      ),
+    );
+    return Promise.all(promises);
+  }
+
+  private invokeEmbeddingModel = async (
+    payload: EmbeddingsPayload,
+    options?: EmbeddingsOptions,
+  ): Promise<Embeddings> => {
+    const command = new InvokeModelCommand({
+      accept: 'application/json',
+      body: JSON.stringify({
+        dimensions: payload.dimensions,
+        inputText: payload.input,
+        normalize: true,
+      }),
+      contentType: 'application/json',
+      modelId: payload.model,
+    });
+    try {
+      const res = await this.client.send(command, { abortSignal: options?.signal });
+      const responseBody = JSON.parse(new TextDecoder().decode(res.body));
+      return responseBody.embedding;
+    } catch (e) {
+      const err = e as Error & { $metadata: any };
+      throw AgentRuntimeError.chat({
+        error: {
+          body: err.$metadata,
+          message: err.message,
+          type: err.name,
+        },
+        errorType: AgentRuntimeErrorType.ProviderBizError,
+        provider: ModelProvider.Bedrock,
+        region: this.region,
+      });
+    }
+  };
 
   private invokeClaudeModel = async (
     payload: ChatStreamPayload,

package/src/libs/agent-runtime/lmstudio/index.test.ts

@@ -0,0 +1,255 @@
+// @vitest-environment node
+import OpenAI from 'openai';
+import { Mock, afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import {
+  ChatStreamCallbacks,
+  LobeOpenAICompatibleRuntime,
+  ModelProvider,
+} from '@/libs/agent-runtime';
+
+import * as debugStreamModule from '../utils/debugStream';
+import { LobeLMStudioAI } from './index';
+
+const provider = ModelProvider.LMStudio;
+const defaultBaseURL = 'http://localhost:1234/v1';
+
+const bizErrorType = 'ProviderBizError';
+const invalidErrorType = 'InvalidProviderAPIKey';
+
+// Mock the console.error to avoid polluting test output
+vi.spyOn(console, 'error').mockImplementation(() => {});
+
+let instance: LobeOpenAICompatibleRuntime;
+
+beforeEach(() => {
+  instance = new LobeLMStudioAI({ apiKey: 'test' });
+
+  // 使用 vi.spyOn 来模拟 chat.completions.create 方法
+  vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
+    new ReadableStream() as any,
+  );
+});
+
+afterEach(() => {
+  vi.clearAllMocks();
+});
+
+describe('LobeLMStudioAI', () => {
+  describe('init', () => {
+    it('should correctly initialize with an API key', async () => {
+      const instance = new LobeLMStudioAI({ apiKey: 'test_api_key' });
+      expect(instance).toBeInstanceOf(LobeLMStudioAI);
+      expect(instance.baseURL).toEqual(defaultBaseURL);
+    });
+  });
+
+  describe('chat', () => {
+    describe('Error', () => {
+      it('should return OpenAIBizError with an openai error response when OpenAI.APIError is thrown', async () => {
+        // Arrange
+        const apiError = new OpenAI.APIError(
+          400,
+          {
+            status: 400,
+            error: {
+              message: 'Bad Request',
+            },
+          },
+          'Error message',
+          {},
+        );
+
+        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(apiError);
+
+        // Act
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'deepseek-chat',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: defaultBaseURL,
+            error: {
+              error: { message: 'Bad Request' },
+              status: 400,
+            },
+            errorType: bizErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should throw AgentRuntimeError with NoOpenAIAPIKey if no apiKey is provided', async () => {
+        try {
+          new LobeLMStudioAI({});
+        } catch (e) {
+          expect(e).toEqual({ errorType: invalidErrorType });
+        }
+      });
+
+      it('should return OpenAIBizError with the cause when OpenAI.APIError is thrown with cause', async () => {
+        // Arrange
+        const errorInfo = {
+          stack: 'abc',
+          cause: {
+            message: 'api is undefined',
+          },
+        };
+        const apiError = new OpenAI.APIError(400, errorInfo, 'module error', {});
+
+        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(apiError);
+
+        // Act
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'deepseek-chat',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: defaultBaseURL,
+            error: {
+              cause: { message: 'api is undefined' },
+              stack: 'abc',
+            },
+            errorType: bizErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should return OpenAIBizError with an cause response with desensitize Url', async () => {
+        // Arrange
+        const errorInfo = {
+          stack: 'abc',
+          cause: { message: 'api is undefined' },
+        };
+        const apiError = new OpenAI.APIError(400, errorInfo, 'module error', {});
+
+        instance = new LobeLMStudioAI({
+          apiKey: 'test',
+
+          baseURL: 'https://api.abc.com/v1',
+        });
+
+        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(apiError);
+
+        // Act
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'deepseek-chat',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: 'https://api.***.com/v1',
+            error: {
+              cause: { message: 'api is undefined' },
+              stack: 'abc',
+            },
+            errorType: bizErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should throw an InvalidDeepSeekAPIKey error type on 401 status code', async () => {
+        // Mock the API call to simulate a 401 error
+        const error = new Error('Unauthorized') as any;
+        error.status = 401;
+        vi.mocked(instance['client'].chat.completions.create).mockRejectedValue(error);
+
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'deepseek-chat',
+            temperature: 0,
+          });
+        } catch (e) {
+          // Expect the chat method to throw an error with InvalidDeepSeekAPIKey
+          expect(e).toEqual({
+            endpoint: defaultBaseURL,
+            error: new Error('Unauthorized'),
+            errorType: invalidErrorType,
+            provider,
+          });
+        }
+      });
+
+      it('should return AgentRuntimeError for non-OpenAI errors', async () => {
+        // Arrange
+        const genericError = new Error('Generic Error');
+
+        vi.spyOn(instance['client'].chat.completions, 'create').mockRejectedValue(genericError);
+
+        // Act
+        try {
+          await instance.chat({
+            messages: [{ content: 'Hello', role: 'user' }],
+            model: 'deepseek-chat',
+            temperature: 0,
+          });
+        } catch (e) {
+          expect(e).toEqual({
+            endpoint: defaultBaseURL,
+            errorType: 'AgentRuntimeError',
+            provider,
+            error: {
+              name: genericError.name,
+              cause: genericError.cause,
+              message: genericError.message,
+              stack: genericError.stack,
+            },
+          });
+        }
+      });
+    });
+
+    describe('DEBUG', () => {
+      it('should call debugStream and return StreamingTextResponse when DEBUG_LMSTUDIO_CHAT_COMPLETION is 1', async () => {
+        // Arrange
+        const mockProdStream = new ReadableStream() as any; // 模拟的 prod 流
+        const mockDebugStream = new ReadableStream({
+          start(controller) {
+            controller.enqueue('Debug stream content');
+            controller.close();
+          },
+        }) as any;
+        mockDebugStream.toReadableStream = () => mockDebugStream; // 添加 toReadableStream 方法
+
+        // 模拟 chat.completions.create 返回值，包括模拟的 tee 方法
+        (instance['client'].chat.completions.create as Mock).mockResolvedValue({
+          tee: () => [mockProdStream, { toReadableStream: () => mockDebugStream }],
+        });
+
+        // 保存原始环境变量值
+        const originalDebugValue = process.env.DEBUG_LMSTUDIO_CHAT_COMPLETION;
+
+        // 模拟环境变量
+        process.env.DEBUG_LMSTUDIO_CHAT_COMPLETION = '1';
+        vi.spyOn(debugStreamModule, 'debugStream').mockImplementation(() => Promise.resolve());
+
+        // 执行测试
+        // 运行你的测试函数，确保它会在条件满足时调用 debugStream
+        // 假设的测试函数调用，你可能需要根据实际情况调整
+        await instance.chat({
+          messages: [{ content: 'Hello', role: 'user' }],
+          model: 'deepseek-chat',
+          stream: true,
+          temperature: 0,
+        });
+
+        // 验证 debugStream 被调用
+        expect(debugStreamModule.debugStream).toHaveBeenCalled();
+
+        // 恢复原始环境变量值
+        process.env.DEBUG_LMSTUDIO_CHAT_COMPLETION = originalDebugValue;
+      });
+    });
+  });
+});

package/src/libs/agent-runtime/lmstudio/index.ts

@@ -0,0 +1,11 @@
+import { ModelProvider } from '../types';
+import { LobeOpenAICompatibleFactory } from '../utils/openaiCompatibleFactory';
+
+export const LobeLMStudioAI = LobeOpenAICompatibleFactory({
+  apiKey: 'placeholder-to-avoid-error',
+  baseURL: 'http://localhost:1234/v1',
+  debug: {
+    chatCompletion: () => process.env.DEBUG_LMSTUDIO_CHAT_COMPLETION === '1',
+  },
+  provider: ModelProvider.LMStudio,
+});

package/src/libs/agent-runtime/ollama/index.ts

@@ -6,7 +6,13 @@ import { ChatModelCard } from '@/types/llm';
 
 import { LobeRuntimeAI } from '../BaseAI';
 import { AgentRuntimeErrorType } from '../error';
-import { ChatCompetitionOptions, ChatStreamPayload, ModelProvider } from '../types';
+import {
+  ChatCompetitionOptions,
+  ChatStreamPayload,
+  Embeddings,
+  EmbeddingsPayload,
+  ModelProvider,
+} from '../types';
 import { AgentRuntimeError } from '../utils/createError';
 import { debugStream } from '../utils/debugStream';
 import { StreamingResponse } from '../utils/response';
@@ -84,6 +90,18 @@ export class LobeOllamaAI implements LobeRuntimeAI {
     }
   }
 
+  async embeddings(payload: EmbeddingsPayload): Promise<Embeddings[]> {
+    const input = Array.isArray(payload.input) ? payload.input : [payload.input];
+    const promises = input.map((inputText: string) =>
+      this.invokeEmbeddingModel({
+        dimensions: payload.dimensions,
+        input: inputText,
+        model: payload.model,
+      }),
+    );
+    return await Promise.all(promises);
+  }
+
   async models(): Promise<ChatModelCard[]> {
     const list = await this.client.list();
     return list.models.map((model) => ({
@@ -91,6 +109,24 @@ export class LobeOllamaAI implements LobeRuntimeAI {
     }));
   }
 
+  private invokeEmbeddingModel = async (payload: EmbeddingsPayload): Promise<Embeddings> => {
+    try {
+      const responseBody = await this.client.embeddings({
+        model: payload.model,
+        prompt: payload.input as string,
+      });
+      return responseBody.embedding;
+    } catch (error) {
+      const e = error as { message: string; name: string; status_code: number };
+
+      throw AgentRuntimeError.chat({
+        error: { message: e.message, name: e.name, status_code: e.status_code },
+        errorType: AgentRuntimeErrorType.OllamaBizError,
+        provider: ModelProvider.Ollama,
+      });
+    }
+  };
+
   private buildOllamaMessages(messages: OpenAIChatMessage[]) {
     return messages.map((message) => this.convertContentToOllamaMessage(message));
   }

package/src/libs/agent-runtime/types/type.ts

@@ -39,6 +39,7 @@ export enum ModelProvider {
   HuggingFace = 'huggingface',
   Hunyuan = 'hunyuan',
   InternLM = 'internlm',
+  LMStudio = 'lmstudio',
   Minimax = 'minimax',
   Mistral = 'mistral',
   Moonshot = 'moonshot',

package/src/server/globalConfig/index.ts

@@ -1,6 +1,7 @@
 import { appEnv, getAppConfig } from '@/config/app';
 import { authEnv } from '@/config/auth';
 import { fileEnv } from '@/config/file';
+import { knowledgeEnv } from '@/config/knowledge';
 import { langfuseEnv } from '@/config/langfuse';
 import { enableNextAuth } from '@/const/auth';
 import { parseSystemAgent } from '@/server/globalConfig/parseSystemAgent';
@@ -9,6 +10,7 @@ import { GlobalServerConfig } from '@/types/serverConfig';
 import { genServerLLMConfig } from './_deprecated';
 import { genServerAiProvidersConfig } from './genServerAiProviderConfig';
 import { parseAgentConfig } from './parseDefaultAgent';
+import { parseFilesConfig } from './parseFilesConfig';
 
 export const getServerGlobalConfig = () => {
   const { ACCESS_CODES, DEFAULT_AGENT_CONFIG } = getAppConfig();
@@ -73,3 +75,7 @@ export const getServerDefaultAgentConfig = () => {
 
   return parseAgentConfig(DEFAULT_AGENT_CONFIG) || {};
 };
+
+export const getServerDefaultFilesConfig = () => {
+  return parseFilesConfig(knowledgeEnv.DEFAULT_FILES_CONFIG);
+};

package/src/server/globalConfig/parseFilesConfig.test.ts

@@ -0,0 +1,17 @@
+import { describe, expect, it } from 'vitest';
+
+import { parseFilesConfig } from './parseFilesConfig';
+
+describe('parseFilesConfig', () => {
+  // 测试embeddings配置是否被正确解析
+  it('parses embeddings configuration correctly', () => {
+    const envStr =
+      'embedding_model=openai/embedding-text-3-large,reranker_model=cohere/rerank-english-v3.0,query_model=full_text';
+    const expected = {
+      embeddingModel: { provider: 'openai', model: 'embedding-text-3-large' },
+      rerankerModel: { provider: 'cohere', model: 'rerank-english-v3.0' },
+      queryModel: 'full_text',
+    };
+    expect(parseFilesConfig(envStr)).toEqual(expected);
+  });
+});

package/src/server/globalConfig/parseFilesConfig.ts

@@ -0,0 +1,57 @@
+import { DEFAULT_FILES_CONFIG } from '@/const/settings/knowledge';
+import { SystemEmbeddingConfig } from '@/types/knowledgeBase';
+import { FilesConfig } from '@/types/user/settings/filesConfig';
+
+const protectedKeys = Object.keys({
+  embedding_model: null,
+  query_model: null,
+  reranker_model: null,
+});
+
+export const parseFilesConfig = (envString: string = ''): SystemEmbeddingConfig => {
+  if (!envString) return DEFAULT_FILES_CONFIG;
+  const config: FilesConfig = {} as any;
+
+  // 处理全角逗号和多余空格
+  let envValue = envString.replaceAll('，', ',').trim();
+
+  const pairs = envValue.split(',');
+
+  for (const pair of pairs) {
+    const [key, value] = pair.split('=').map((s) => s.trim());
+
+    if (key && value) {
+      const [provider, ...modelParts] = value.split('/');
+      const model = modelParts.join('/');
+
+      if ((!provider || !model) && key !== 'query_model') {
+        throw new Error('Missing model or provider value');
+      }
+
+      if (key === 'query_model' && value === '') {
+        throw new Error('Missing query mode value');
+      }
+
+      if (protectedKeys.includes(key)) {
+        switch (key) {
+          case 'embedding_model': {
+            config.embeddingModel = { model: model.trim(), provider: provider.trim() };
+            break;
+          }
+          case 'reranker_model': {
+            config.rerankerModel = { model: model.trim(), provider: provider.trim() };
+            break;
+          }
+          case 'query_model': {
+            config.queryModel = value;
+            break;
+          }
+        }
+      }
+    } else {
+      throw new Error('Invalid environment variable format');
+    }
+  }
+
+  return config;
+};