⚡️ feat: support using env variable to set regions for OpenAI Edge Fu…

…nctions. (lobehub#473) * ⚡️ feat: support using env variable to set preferredRegion for OpenAI edge function * fix: fix the import problems * ⚡️ feat: add preferredRegion support for tts and stt * ✅ test: add more unit test for config.ts and server.ts --------- Co-authored-by: vophan <dev@vophan.day>
bryan31 · Nov 21, 2023 · de6b79e · de6b79e
1 parent 0c4ee9f
commit de6b79e
Show file tree

Hide file tree

Showing 9 changed files with 84 additions and 10 deletions.
diff --git a/README.md b/README.md
@@ -390,11 +390,12 @@ $ docker run -d -p 3210:3210 \
 
 This project provides some additional configuration items set with environment variables:
 
-| Environment Variable | Required | Description                                                                                                                                   | Example                                                                                   |
-| -------------------- | -------- | --------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------- |
-| `OPENAI_API_KEY`     | Yes      | This is the API key you apply on the OpenAI account page                                                                                      | `sk-xxxxxx...xxxxxx`                                                                      |
-| `OPENAI_PROXY_URL`   | No       | If you manually configure the OpenAI interface proxy, you can use this configuration item to override the default OpenAI API request base URL | `https://api.chatanywhere.cn/v1`<br/>The default value is<br/>`https://api.openai.com/v1` |
-| `ACCESS_CODE`        | No       | Add a password to access this service; the password should be a 6-digit number or letter                                                      | `awCT74` or `e3@09!`                                                                      |
+| Environment Variable      | Required | Description                                                                                                                                                                                                                  | Example                                                                                   |
+| ------------------------- | -------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------- |
+| `OPENAI_API_KEY`          | Yes      | This is the API key you apply on the OpenAI account page                                                                                                                                                                     | `sk-xxxxxx...xxxxxx`                                                                      |
+| `OPENAI_PROXY_URL`        | No       | If you manually configure the OpenAI interface proxy, you can use this configuration item to override the default OpenAI API request base URL                                                                                | `https://api.chatanywhere.cn/v1`<br/>The default value is<br/>`https://api.openai.com/v1` |
+| `OPENAI_FUNCTION_REGIONS` | No       | When you deploy Lobe-Chat using Vercel and need to specify the region for the Edge Function that handles requests to the OpenAI API, you can use this configuration. The value should be a comma-separated array of strings. | `iad1,sfo1`                                                                               |
+| `ACCESS_CODE`             | No       | Add a password to access this service; the password should be a 6-digit number or letter                                                                                                                                     | `awCT74` or `e3@09!`                                                                      |
 
 > \[!NOTE]
 >

diff --git a/README.zh-CN.md b/README.zh-CN.md
@@ -364,11 +364,12 @@ $ docker run -d -p 3210:3210 \
 
 本项目提供了一些额外的配置项，使用环境变量进行设置：
 
-| 环境变量           | 类型 | 描述                                                                                   | 示例                                                                         |
-| ------------------ | ---- | -------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------- |
-| `OPENAI_API_KEY`   | 必选 | 这是你在 OpenAI 账户页面申请的 API 密钥                                                | `sk-xxxxxx...xxxxxx`                                                         |
-| `OPENAI_PROXY_URL` | 可选 | 如果你手动配置了 OpenAI 接口代理，可以使用此配置项来覆盖默认的 OpenAI API 请求基础 URL | `https://api.chatanywhere.cn/v1`<br/>默认值:<br/>`https://api.openai.com/v1` |
-| `ACCESS_CODE`      | 可选 | 添加访问此服务的密码，密码应为 6 位数字或字母                                          | `awCT74` 或 `e3@09!`                                                         |
+| 环境变量                  | 类型 | 描述                                                                                                                                                            | 示例                                                                         |
+| ------------------------- | ---- | --------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---------------------------------------------------------------------------- |
+| `OPENAI_API_KEY`          | 必选 | 这是你在 OpenAI 账户页面申请的 API 密钥                                                                                                                         | `sk-xxxxxx...xxxxxx`                                                         |
+| `OPENAI_PROXY_URL`        | 可选 | 如果你手动配置了 OpenAI 接口代理，可以使用此配置项来覆盖默认的 OpenAI API 请求基础 URL                                                                          | `https://api.chatanywhere.cn/v1`<br/>默认值:<br/>`https://api.openai.com/v1` |
+| `OPENAI_FUNCTION_REGIONS` | 可选 | 当你使用 Vercel 部署 Lobe-Chat，而且有需求指定响应调用 OpenAI 接口的请求的 Edge Function 的 Region 时，可以使用该配置进行配置，该值的类型为逗号分隔的字符串数组 | `iad1,sfo1`                                                                  |
+| `ACCESS_CODE`             | 可选 | 添加访问此服务的密码，密码应为 6 位数字或字母                                                                                                                   | `awCT74` 或 `e3@09!`                                                         |
 
 > \[!NOTE]
 >

diff --git a/src/app/api/config.test.ts b/src/app/api/config.test.ts
@@ -0,0 +1,38 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import { getPreferredRegion } from './config';
+
+// 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('getPreferredRegion', () => {
+  beforeEach(() => {
+    // Reset environment variables before each test case
+    vi.restoreAllMocks();
+  });
+
+  it('returns default value when get config error', () => {
+    const originalProcess = global.process;
+    // @ts-ignore
+    global.process = undefined;
+    const preferredRegion = getPreferredRegion();
+    expect(preferredRegion).toBe('auto');
+
+    global.process = originalProcess;
+  });
+
+  it('return default value when preferredRegion is empty', () => {
+    process.env.OPENAI_FUNCTION_REGIONS = '';
+    const preferredRegion = getPreferredRegion();
+    expect(preferredRegion).toBe('auto');
+  });
+
+  it('return correct list values when preferredRegion is correctly passed', () => {
+    process.env.OPENAI_FUNCTION_REGIONS = 'ida1,sfo1';
+    const preferredRegion = getPreferredRegion();
+    expect(preferredRegion).toStrictEqual(['ida1', 'sfo1']);
+  });
+});
diff --git a/src/app/api/config.ts b/src/app/api/config.ts
@@ -0,0 +1,15 @@
+import { getServerConfig } from '@/config/server';
+
+export const getPreferredRegion = () => {
+  try {
+    const cfg = getServerConfig();
+    if (cfg.OPENAI_FUNCTION_REGIONS.length <= 0) {
+      return 'auto';
+    }
+
+    return cfg.OPENAI_FUNCTION_REGIONS;
+  } catch (error) {
+    console.error('get server config failed, error:', error);
+    return 'auto';
+  }
+};
diff --git a/src/app/api/openai/chat/route.ts b/src/app/api/openai/chat/route.ts
@@ -1,9 +1,11 @@
 import { OpenAIChatStreamPayload } from '@/types/openai/chat';
 
+import { getPreferredRegion } from '../../config';
 import { createBizOpenAI } from '../createBizOpenAI';
 import { createChatCompletion } from './createChatCompletion';
 
 export const runtime = 'edge';
+export const preferredRegion = getPreferredRegion();
 
 export const POST = async (req: Request) => {
   const payload = (await req.json()) as OpenAIChatStreamPayload;

diff --git a/src/app/api/openai/stt/route.ts b/src/app/api/openai/stt/route.ts
@@ -1,9 +1,11 @@
 import { OpenAISTTPayload } from '@lobehub/tts';
 import { createOpenaiAudioTranscriptions } from '@lobehub/tts/server';
 
+import { getPreferredRegion } from '../../config';
 import { createBizOpenAI } from '../createBizOpenAI';
 
 export const runtime = 'edge';
+export const preferredRegion = getPreferredRegion();
 
 export const POST = async (req: Request) => {
   const formData = await req.formData();

diff --git a/src/app/api/openai/tts/route.ts b/src/app/api/openai/tts/route.ts
@@ -1,9 +1,11 @@
 import { OpenAITTSPayload } from '@lobehub/tts';
 import { createOpenaiAudioSpeech } from '@lobehub/tts/server';
 
+import { getPreferredRegion } from '../../config';
 import { createBizOpenAI } from '../createBizOpenAI';
 
 export const runtime = 'edge';
+export const preferredRegion = getPreferredRegion();
 
 export const POST = async (req: Request) => {
   const payload = (await req.json()) as OpenAITTSPayload;

diff --git a/src/config/__tests__/server.test.ts b/src/config/__tests__/server.test.ts
@@ -39,6 +39,12 @@ describe('getServerConfig', () => {
     expect(config.USE_AZURE_OPENAI).toBe(false);
   });
 
+  it('correctly handles values for OPENAI_FUNCTION_REGIONS', () => {
+    process.env.OPENAI_FUNCTION_REGIONS = 'iad1,sfo1';
+    const config = getServerConfig();
+    expect(config.OPENAI_FUNCTION_REGIONS).toStrictEqual(['iad1', 'sfo1']);
+  });
+
   it('returns default IMGUR_CLIENT_ID when no environment variable is set', () => {
     const config = getServerConfig();
     expect(config.IMGUR_CLIENT_ID).toBe('e415f320d6e24f9');

diff --git a/src/config/server.ts b/src/config/server.ts
@@ -27,11 +27,18 @@ export const getServerConfig = () => {
     throw new Error('[Server Config] you are importing a nodejs-only module outside of nodejs');
   }
 
+  // region format: iad1,sfo1
+  let regions: string[] = [];
+  if (process.env.OPENAI_FUNCTION_REGIONS) {
+    regions = process.env.OPENAI_FUNCTION_REGIONS.split(',');
+  }
+
   return {
     ACCESS_CODE: process.env.ACCESS_CODE,
 
     OPENAI_API_KEY: process.env.OPENAI_API_KEY,
     OPENAI_PROXY_URL: process.env.OPENAI_PROXY_URL,
+    OPENAI_FUNCTION_REGIONS: regions,
 
     AZURE_API_KEY: process.env.AZURE_API_KEY,
     AZURE_API_VERSION: process.env.AZURE_API_VERSION,