diff --git a/.env.example b/.env.example index 80f55efd9c4d..7d139f383995 100644 --- a/.env.example +++ b/.env.example @@ -4,6 +4,9 @@ # add your custom model name, multi model separate by comma. for example gpt-3.5-1106,gpt-4-1106 # CUSTOM_MODELS=model1,model2,model3 +# Specify your API Key selection method, currently supporting `random` and `turn`. +# API_KEY_SELECT_MODE=random + # ---- only choose one from OpenAI Service and Azure OpenAI Service ---- # ######################################## diff --git a/Dockerfile b/Dockerfile index 798b7caae085..7303130f305d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -60,6 +60,8 @@ ENV PORT=3210 ENV ACCESS_CODE "" ENV CUSTOM_MODELS "" +ENV API_KEY_SELECT_MODE "" + # OpenAI ENV OPENAI_API_KEY "" ENV OPENAI_PROXY_URL "" diff --git a/docs/self-hosting/environment-variables/model-provider.mdx b/docs/self-hosting/environment-variables/model-provider.mdx index eb8f72f0e254..8cc19a681509 100644 --- a/docs/self-hosting/environment-variables/model-provider.mdx +++ b/docs/self-hosting/environment-variables/model-provider.mdx @@ -54,6 +54,17 @@ You can find all current model names in [modelProviders](https://github.com/lobe If you need to use Azure OpenAI to provide model services, you can refer to the [Deploying with Azure OpenAI](../Deployment/Deploy-with-Azure-OpenAI.en-US.md) section for detailed steps. Here, we will list the environment variables related to Azure OpenAI. +### `API_KEY_SELECT_MODE` + +- Type：Optional +- Description：Controls the mode for selecting the API Key when multiple API Keys are available. Currently supports `random` and `turn`. +- Default：`random` +- Example：`random` or `turn` + +When using the `random` mode, a random API Key will be selected from the available multiple API Keys. + +When using the `turn` mode, the API Keys will be retrieved in a round-robin manner according to the specified order. + ### `USE_AZURE_OPENAI` - Type: Optional diff --git a/docs/self-hosting/environment-variables/model-provider.zh-CN.mdx b/docs/self-hosting/environment-variables/model-provider.zh-CN.mdx index 91ab5911af60..2d5349a6086d 100644 --- a/docs/self-hosting/environment-variables/model-provider.zh-CN.mdx +++ b/docs/self-hosting/environment-variables/model-provider.zh-CN.mdx @@ -48,6 +48,17 @@ LobeChat 在部署时提供了丰富的模型服务商相关的环境变量， 你可以在 [modelProviders](https://github.com/lobehub/lobe-chat/tree/main/src/config/modelProviders) 查找到当前的所有模型名。 +### `API_KEY_SELECT_MODE` + +- 类型：可选 +- 描述：用于控制多个API Keys时，选择Key的模式，当前支持 `random` 和 `turn` +- 默认值：`random` +- 示例：`random` 或 `turn` + +使用 `random` 模式下，将在多个API Keys中随机获取一个API Key。 + +使用 `turn` 模式下，将按照填写的顺序，轮训获取得到API Key。 + ## Azure OpenAI 如果你需要使用 Azure OpenAI 来提供模型服务，可以查阅 [使用 Azure OpenAI 部署](../Deployment/Deploy-with-Azure-OpenAI.zh-CN.md) 章节查看详细步骤，这里将列举和 Azure OpenAI 相关的环境变量。 diff --git a/src/app/api/chat/[provider]/agentRuntime.ts b/src/app/api/chat/[provider]/agentRuntime.ts index 23d2b88a2a3f..8d541a401d16 100644 --- a/src/app/api/chat/[provider]/agentRuntime.ts +++ b/src/app/api/chat/[provider]/agentRuntime.ts @@ -14,6 +14,8 @@ import { ModelProvider, } from '@/libs/agent-runtime'; +import apiKeyManager from '../apiKeyManager'; + interface AzureOpenAIParams { apiVersion?: string; model: string; @@ -88,15 +90,17 @@ class AgentRuntime { private static initOpenAI(payload: JWTPayload, azureOpenAI?: AzureOpenAIParams) { const { OPENAI_API_KEY, OPENAI_PROXY_URL, AZURE_API_VERSION, AZURE_API_KEY, USE_AZURE_OPENAI } = getServerConfig(); - const apiKey = payload?.apiKey || OPENAI_API_KEY; + const openaiApiKey = payload?.apiKey || OPENAI_API_KEY; const baseURL = payload?.endpoint || OPENAI_PROXY_URL; const azureApiKey = payload.apiKey || AZURE_API_KEY; const useAzure = azureOpenAI?.useAzure || USE_AZURE_OPENAI; const apiVersion = azureOpenAI?.apiVersion || AZURE_API_VERSION; + const apiKey = apiKeyManager.pick(useAzure ? azureApiKey : openaiApiKey); + return new LobeOpenAI({ - apiKey: useAzure ? azureApiKey : apiKey, + apiKey, azureOptions: { apiVersion, model: azureOpenAI?.model, @@ -108,7 +112,7 @@ class AgentRuntime { private static initAzureOpenAI(payload: JWTPayload) { const { AZURE_API_KEY, AZURE_API_VERSION, AZURE_ENDPOINT } = getServerConfig(); - const apiKey = payload?.apiKey || AZURE_API_KEY; + const apiKey = apiKeyManager.pick(payload?.apiKey || AZURE_API_KEY); const endpoint = payload?.endpoint || AZURE_ENDPOINT; const apiVersion = payload?.azureApiVersion || AZURE_API_VERSION; @@ -117,21 +121,21 @@ class AgentRuntime { private static async initZhipu(payload: JWTPayload) { const { ZHIPU_API_KEY } = getServerConfig(); - const apiKey = payload?.apiKey || ZHIPU_API_KEY; + const apiKey = apiKeyManager.pick(payload?.apiKey || ZHIPU_API_KEY); return LobeZhipuAI.fromAPIKey(apiKey); } private static initMoonshot(payload: JWTPayload) { const { MOONSHOT_API_KEY, MOONSHOT_PROXY_URL } = getServerConfig(); - const apiKey = payload?.apiKey || MOONSHOT_API_KEY; + const apiKey = apiKeyManager.pick(payload?.apiKey || MOONSHOT_API_KEY); return new LobeMoonshotAI(apiKey, MOONSHOT_PROXY_URL); } private static initGoogle(payload: JWTPayload) { const { GOOGLE_API_KEY } = getServerConfig(); - const apiKey = payload?.apiKey || GOOGLE_API_KEY; + const apiKey = apiKeyManager.pick(payload?.apiKey || GOOGLE_API_KEY); return new LobeGoogleAI(apiKey); } @@ -161,7 +165,7 @@ class AgentRuntime { private static initPerplexity(payload: JWTPayload) { const { PERPLEXITY_API_KEY } = getServerConfig(); - const apiKey = payload?.apiKey || PERPLEXITY_API_KEY; + const apiKey = apiKeyManager.pick(payload?.apiKey || PERPLEXITY_API_KEY); return new LobePerplexityAI(apiKey); } diff --git a/src/app/api/chat/apiKeyManager.test.ts b/src/app/api/chat/apiKeyManager.test.ts new file mode 100644 index 000000000000..27bf71fc6417 --- /dev/null +++ b/src/app/api/chat/apiKeyManager.test.ts @@ -0,0 +1,132 @@ +import { nanoid } from 'nanoid'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { getServerConfig } from '@/config/server'; + +import { ApiKeyManager } from './apiKeyManager'; + +function generateKeys(count: number = 1) { + return new Array(count) + .fill('') + .map(() => { + return `sk-${nanoid()}`; + }) + .join(','); +} + +// Stub the global process object to safely mock environment variables +vi.stubGlobal('process', { + ...process, // Preserve the original process object + env: { ...process.env }, // Clone the environment variables object for modification +}); + +describe('apiKeyManager', () => { + beforeEach(() => { + vi.restoreAllMocks(); + }); + + describe('API Key unset or empty', () => { + it('should return an empty string when API_KEY_SELECT_MODE is unset', () => { + const apiKeyManager = new ApiKeyManager(); + + expect(apiKeyManager.pick('')).toBe(''); + expect(apiKeyManager.pick()).toBe(''); + }); + + it('should return an empty string when API_KEY_SELECT_MODE is "random"', () => { + process.env.API_KEY_SELECT_MODE = 'random'; + const apiKeyManager = new ApiKeyManager(); + + expect(apiKeyManager.pick('')).toBe(''); + expect(apiKeyManager.pick()).toBe(''); + }); + + it('should return an empty string when API_KEY_SELECT_MODE is "turn"', () => { + process.env.API_KEY_SELECT_MODE = 'turn'; + const apiKeyManager = new ApiKeyManager(); + + expect(apiKeyManager.pick('')).toBe(''); + expect(apiKeyManager.pick()).toBe(''); + }); + }); + + describe('single API Key', () => { + it('should return the only API Key when API_KEY_SELECT_MODE is unset', () => { + const apiKeyManager = new ApiKeyManager(); + const apiKeyStr = generateKeys(1); + + expect(apiKeyManager.pick(apiKeyStr)).toBe(apiKeyStr); + }); + + it('should return the only API when API_KEY_SELECT_MODE is "random"', () => { + process.env.API_KEY_SELECT_MODE = 'random'; + const apiKeyStr = generateKeys(1); + const apiKeyManager = new ApiKeyManager(); + + expect(apiKeyManager.pick(apiKeyStr)).toBe(apiKeyStr); + // multiple + expect(apiKeyManager.pick(apiKeyStr)).toBe(apiKeyStr); + expect(apiKeyManager.pick(apiKeyStr)).toBe(apiKeyStr); + }); + + it('should return the only API when API_KEY_SELECT_MODE is "turn"', () => { + process.env.API_KEY_SELECT_MODE = 'turn'; + const apiKeyStr = generateKeys(1); + const apiKeyManager = new ApiKeyManager(); + + expect(apiKeyManager.pick(apiKeyStr)).toBe(apiKeyStr); + // multiple + expect(apiKeyManager.pick(apiKeyStr)).toBe(apiKeyStr); + }); + }); + + describe('multiple API Keys', () => { + it('should return a random API Key when API_KEY_SELECT_MODE is unset', () => { + const apiKeyStr = generateKeys(5); + const apiKeys = apiKeyStr.split(','); + const apiKeyManager = new ApiKeyManager(); + const keyLen = apiKeys.length * 2; // multiple round + + for (let i = 0; i < keyLen; i++) { + expect(apiKeys).toContain(apiKeyManager.pick(apiKeyStr)); + } + }); + + it('should return a random API Key when environment variable of API_KEY_SELECT_MODE is "random"', () => { + process.env.API_KEY_SELECT_MODE = 'random'; + const apiKeyStr = generateKeys(5); + const apiKeys = apiKeyStr.split(','); + const apiKeyManager = new ApiKeyManager(); + const keyLen = apiKeys.length * 2; // multiple round + + for (let i = 0; i < keyLen; i++) { + expect(apiKeys).toContain(apiKeyManager.pick(apiKeyStr)); + } + }); + + it('should return API Keys sequentially when environment variable of API_KEY_SELECT_MODE is "turn"', () => { + process.env.API_KEY_SELECT_MODE = 'turn'; + const apiKeyStr = generateKeys(5); + const apiKeys = apiKeyStr.split(','); + const apiKeyManager = new ApiKeyManager(); + + const total = apiKeys.length; + const rounds = total * 2; + for (let i = 0; i < total; i++) { + expect(apiKeyManager.pick(apiKeyStr)).toBe(apiKeys[i % total]); + } + }); + + it('should return a random API Key when API_KEY_SELECT_MODE is anything other than "random" or "turn"', () => { + process.env.API_KEY_SELECT_MODE = nanoid(); + const apiKeyStr = generateKeys(5); + const apiKeys = apiKeyStr.split(','); + const apiKeyManager = new ApiKeyManager(); + const keyLen = apiKeys.length * 2; // multiple round + + for (let i = 0; i < keyLen; i++) { + expect(apiKeys).toContain(apiKeyManager.pick(apiKeyStr)); + } + }); + }); +}); diff --git a/src/app/api/chat/apiKeyManager.ts b/src/app/api/chat/apiKeyManager.ts new file mode 100644 index 000000000000..b16e170a8df8 --- /dev/null +++ b/src/app/api/chat/apiKeyManager.ts @@ -0,0 +1,46 @@ +import { getServerConfig } from '@/config/server'; + +interface KeyStore { + index: number; + keyLen: number; + keys: string[]; +} + +export class ApiKeyManager { + private _cache: Map = new Map(); + + private _mode: string; + + constructor() { + const { API_KEY_SELECT_MODE: mode = 'random' } = getServerConfig(); + + this._mode = mode; + } + + private getKeyStore(apiKeys: string) { + let store = this._cache.get(apiKeys); + + if (!store) { + const keys = apiKeys.split(',').filter((_) => !!_.trim()); + + store = { index: 0, keyLen: keys.length, keys } as KeyStore; + this._cache.set(apiKeys, store); + } + + return store; + } + + pick(apiKeys: string = '') { + if (!apiKeys) return ''; + + const store = this.getKeyStore(apiKeys); + let index = 0; + + if (this._mode === 'turn') index = store.index++ % store.keyLen; + if (this._mode === 'random') index = Math.floor(Math.random() * store.keyLen); + + return store.keys[index]; + } +} + +export default new ApiKeyManager(); diff --git a/src/app/api/chat/google/route.ts b/src/app/api/chat/google/route.ts index 553bf9635880..379dcab904c2 100644 --- a/src/app/api/chat/google/route.ts +++ b/src/app/api/chat/google/route.ts @@ -11,6 +11,7 @@ import { import { ChatErrorType } from '@/types/fetch'; import { ChatStreamPayload } from '@/types/openai/chat'; +import apiKeyManager from '../apiKeyManager'; import { checkAuthMethod, getJWTPayload } from '../auth'; // due to the Chinese region does not support accessing Google @@ -58,7 +59,7 @@ export const POST = async (req: Request) => { checkAuthMethod(payload.accessCode, payload.apiKey, oauthAuthorized); const { GOOGLE_API_KEY } = getServerConfig(); - const apiKey = payload?.apiKey || GOOGLE_API_KEY; + const apiKey = apiKeyManager.pick(payload?.apiKey || GOOGLE_API_KEY); agentRuntime = new LobeGoogleAI(apiKey); } catch (e) { diff --git a/src/config/server/provider.ts b/src/config/server/provider.ts index 27fd700b019f..4a926ee7c1ce 100644 --- a/src/config/server/provider.ts +++ b/src/config/server/provider.ts @@ -6,6 +6,8 @@ declare global { interface ProcessEnv { CUSTOM_MODELS?: string; + API_KEY_SELECT_MODE?: string; + // OpenAI Provider OPENAI_API_KEY?: string; OPENAI_PROXY_URL?: string; @@ -67,6 +69,8 @@ export const getProviderConfig = () => { return { CUSTOM_MODELS: process.env.CUSTOM_MODELS, + API_KEY_SELECT_MODE: process.env.API_KEY_SELECT_MODE, + OPENAI_API_KEY: process.env.OPENAI_API_KEY, OPENAI_PROXY_URL: process.env.OPENAI_PROXY_URL, OPENAI_FUNCTION_REGIONS: regions,