From 7e8def50aa9e4c95464a21b021a707a0ccd28ec3 Mon Sep 17 00:00:00 2001
From: Yidadaa <yidadaa@qq.com>
Date: Tue, 9 May 2023 23:20:03 +0800
Subject: [PATCH] feat: close #444 use env var to disable gpt-4

---
 README.md               |  8 +++++++-
 README_CN.md            |  6 +++++-
 app/api/config/route.ts |  3 ++-
 app/config/server.ts    |  2 ++
 app/store/access.ts     |  9 +++++++++
 app/store/config.ts     | 20 ++++++++++++++++++++
 6 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 7f6695bb..90ed7d42 100644
--- a/README.md
+++ b/README.md
@@ -168,7 +168,13 @@ Specify OpenAI organization ID.
 
 > Default: Empty
 
-If you do not want users to input their own API key, set this environment variable to 1.
+If you do not want users to input their own API key, set this value to 1.
+
+### `DISABLE_GPT4` (optional)
+
+> Default: Empty
+
+If you do not want users to use GPT-4, set this value to 1.
 
 ## Development
 
diff --git a/README_CN.md b/README_CN.md
index 1da68f65..9601e5fd 100644
--- a/README_CN.md
+++ b/README_CN.md
@@ -64,7 +64,7 @@ code1,code2,code3
 
 ## 环境变量
 
-> 本项目大多数配置项都通过环境变量来设置。
+> 本项目大多数配置项都通过环境变量来设置，教程：[如何修改 Vercel 环境变量](./docs/vercel-cn.md)。
 
 ### `OPENAI_API_KEY` （必填项）
 
@@ -94,6 +94,10 @@ OpenAI 接口代理 URL，如果你手动配置了 openai 接口代理，请填
 
 如果你不想让用户自行填入 API Key，将此环境变量设置为 1 即可。
 
+### `DISABLE_GPT4` （可选）
+
+如果你不想让用户使用 GPT-4，将此环境变量设置为 1 即可。
+
 ## 开发
 
 > 强烈不建议在本地进行开发或者部署，由于一些技术原因，很难在本地配置好 OpenAI API 代理，除非你能保证可以直连 OpenAI 服务器。
diff --git a/app/api/config/route.ts b/app/api/config/route.ts
index 62c8d60f..2b3bcbf2 100644
--- a/app/api/config/route.ts
+++ b/app/api/config/route.ts
@@ -1,4 +1,4 @@
-import { NextRequest, NextResponse } from "next/server";
+import { NextResponse } from "next/server";
 
 import { getServerSideConfig } from "../../config/server";
 
@@ -9,6 +9,7 @@ const serverConfig = getServerSideConfig();
 const DANGER_CONFIG = {
   needCode: serverConfig.needCode,
   hideUserApiKey: serverConfig.hideUserApiKey,
+  enableGPT4: serverConfig.enableGPT4,
 };
 
 declare global {
diff --git a/app/config/server.ts b/app/config/server.ts
index c1cf439b..23fec868 100644
--- a/app/config/server.ts
+++ b/app/config/server.ts
@@ -8,6 +8,7 @@ declare global {
       PROXY_URL?: string;
       VERCEL?: string;
       HIDE_USER_API_KEY?: string; // disable user's api key input
+      DISABLE_GPT4?: string; // allow user to use gpt-4 or not
     }
   }
 }
@@ -40,5 +41,6 @@ export const getServerSideConfig = () => {
     proxyUrl: process.env.PROXY_URL,
     isVercel: !!process.env.VERCEL,
     hideUserApiKey: !!process.env.HIDE_USER_API_KEY,
+    enableGPT4: !process.env.DISABLE_GPT4,
   };
 };
diff --git a/app/store/access.ts b/app/store/access.ts
index 51290d0a..79b7b990 100644
--- a/app/store/access.ts
+++ b/app/store/access.ts
@@ -2,6 +2,7 @@ import { create } from "zustand";
 import { persist } from "zustand/middleware";
 import { StoreKey } from "../constant";
 import { BOT_HELLO } from "./chat";
+import { ALL_MODELS } from "./config";
 
 export interface AccessControlStore {
   accessCode: string;
@@ -60,6 +61,14 @@ export const useAccessStore = create<AccessControlStore>()(
             console.log("[Config] got config from server", res);
             set(() => ({ ...res }));
 
+            if (!res.enableGPT4) {
+              ALL_MODELS.forEach((model) => {
+                if (model.name.startsWith("gpt-4")) {
+                  (model as any).available = false;
+                }
+              });
+            }
+
             if ((res as any).botHello) {
               BOT_HELLO.content = (res as any).botHello;
             }
diff --git a/app/store/config.ts b/app/store/config.ts
index 926c296f..1e960456 100644
--- a/app/store/config.ts
+++ b/app/store/config.ts
@@ -76,6 +76,26 @@ export const ALL_MODELS = [
     name: "gpt-3.5-turbo-0301",
     available: true,
   },
+  {
+    name: "qwen-v1", // 通义千问
+    available: false,
+  },
+  {
+    name: "ernie", // 文心一言
+    available: false,
+  },
+  {
+    name: "spark", // 讯飞星火
+    available: false,
+  },
+  {
+    name: "llama", // llama
+    available: false,
+  },
+  {
+    name: "chatglm", // chatglm-6b
+    available: false,
+  },
 ] as const;
 
 export type ModelType = (typeof ALL_MODELS)[number]["name"];