✨ feat: support multiple API Keys (#1345)

* ✨ feat: Support multiple API keys [RFC 027]

* 📝 docs: add env variable API_KEY_SELECT_MODE

* 🔧 chore: Adjust the parameter API_KEY_SELECT_MODE to be optional

* 🔧 fix: Adjustments made according to Code Review requirements

* ✅ test: add test for ApiKeyManager

* 🔧 fix: Support for multiple API Keys from user input on the client side

* 🔧 chore: handle Perplexity API Key

* 🔧 chore: update OpenAI or Azure API Key select

* 🔧 chore: update OpenAI or Azure API Key select

.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 ---- #
 
 ########################################

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 ""

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

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 相关的环境变量。

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);
   }

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));
+      }
+    });
+  });
+});

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<string, KeyStore> = 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();

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) {

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,