Merge pull request #48 from Yidadaa/custom-token

v1.4 Custom Api Key & Copy Code Button

app/api/chat-stream/route.ts

@@ -2,19 +2,25 @@ import type { ChatRequest } from "../chat/typing";
 import { createParser } from "eventsource-parser";
 import { NextRequest } from "next/server";
 
-const apiKey = process.env.OPENAI_API_KEY;
-
-async function createStream(payload: ReadableStream<Uint8Array>) {
+async function createStream(req: NextRequest) {
   const encoder = new TextEncoder();
   const decoder = new TextDecoder();
 
+  let apiKey = process.env.OPENAI_API_KEY;
+
+  const userApiKey = req.headers.get("token");
+  if (userApiKey) {
+    apiKey = userApiKey;
+    console.log("[Stream] using user api key");
+  }
+
   const res = await fetch("https://api.openai.com/v1/chat/completions", {
     headers: {
       "Content-Type": "application/json",
       Authorization: `Bearer ${apiKey}`,
     },
     method: "POST",
-    body: payload,
+    body: req.body,
   });
 
   const stream = new ReadableStream({
@@ -49,7 +55,7 @@ async function createStream(payload: ReadableStream<Uint8Array>) {
 
 export async function POST(req: NextRequest) {
   try {
-    const stream = await createStream(req.body!);
+    const stream = await createStream(req);
     return new Response(stream);
   } catch (error) {
     console.error("[Chat Stream]", error);

app/api/chat/route.ts

@@ -1,23 +1,26 @@
 import { OpenAIApi, Configuration } from "openai";
 import { ChatRequest } from "./typing";
 
-const apiKey = process.env.OPENAI_API_KEY;
-
-const openai = new OpenAIApi(
-  new Configuration({
-    apiKey,
-  })
-);
-
 export async function POST(req: Request) {
   try {
-    const requestBody = (await req.json()) as ChatRequest;
-    const completion = await openai!.createChatCompletion(
-      {
-        ...requestBody,
-      }
+    let apiKey = process.env.OPENAI_API_KEY;
+
+    const userApiKey = req.headers.get("token");
+    if (userApiKey) {
+      apiKey = userApiKey;
+    }
+
+    const openai = new OpenAIApi(
+      new Configuration({
+        apiKey,
+      })
     );
 
+    const requestBody = (await req.json()) as ChatRequest;
+    const completion = await openai!.createChatCompletion({
+      ...requestBody,
+    });
+
     return new Response(JSON.stringify(completion.data));
   } catch (e) {
     console.error("[Chat] ", e);

app/components/markdown.tsx

@@ -4,15 +4,36 @@ import RemarkMath from "remark-math";
 import RehypeKatex from "rehype-katex";
 import RemarkGfm from "remark-gfm";
 import RehypePrsim from "rehype-prism-plus";
+import { useRef } from "react";
+import { copyToClipboard } from "../utils";
+
+export function PreCode(props: { children: any }) {
+  const ref = useRef<HTMLPreElement>(null);
+
+  return (
+    <pre ref={ref}>
+      <span
+        className="copy-code-button"
+        onClick={() => {
+          if (ref.current) {
+            const code = ref.current.innerText;
+            copyToClipboard(code);
+          }
+        }}
+      ></span>
+      {props.children}
+    </pre>
+  );
+}
 
 export function Markdown(props: { content: string }) {
   return (
     <ReactMarkdown
       remarkPlugins={[RemarkMath, RemarkGfm]}
-      rehypePlugins={[
-        RehypeKatex,
-        [RehypePrsim, { ignoreMissing: true }],
-      ]}
+      rehypePlugins={[RehypeKatex, [RehypePrsim, { ignoreMissing: true }]]}
+      components={{
+        pre: PreCode,
+      }}
     >
       {props.content}
     </ReactMarkdown>

app/components/settings.tsx

@@ -257,6 +257,20 @@ export function Settings(props: { closeSettings: () => void }) {
             <></>
           )}
 
+          <SettingItem
+            title={Locale.Settings.Token.Title}
+            subTitle={Locale.Settings.Token.SubTitle}
+          >
+            <input
+              value={accessStore.token}
+              type="text"
+              placeholder={Locale.Settings.Token.Placeholder}
+              onChange={(e) => {
+                accessStore.updateToken(e.currentTarget.value);
+              }}
+            ></input>
+          </SettingItem>
+
           <SettingItem
             title={Locale.Settings.HistoryCount.Title}
             subTitle={Locale.Settings.HistoryCount.SubTitle}

app/locales/cn.ts

@@ -69,6 +69,11 @@ const cn = {
       Title: "历史消息长度压缩阈值",
       SubTitle: "当未压缩的历史消息超过该值时，将进行压缩",
     },
+    Token: {
+      Title: "API Key",
+      SubTitle: "使用自己的 Key 可绕过受控访问限制",
+      Placeholder: "OpenAI API Key",
+    },
     AccessCode: {
       Title: "访问码",
       SubTitle: "现在是受控访问状态",

app/locales/en.ts

@@ -74,6 +74,11 @@ const en: LocaleType = {
       SubTitle:
         "Will compress if uncompressed messages length exceeds the value",
     },
+    Token: {
+      Title: "API Key",
+      SubTitle: "Use your key to ignore access code limit",
+      Placeholder: "OpenAI API Key",
+    },
     AccessCode: {
       Title: "Access Code",
       SubTitle: "Access control enabled",

app/page.tsx

@@ -1,5 +1,5 @@
 import { Analytics } from "@vercel/analytics/react";
-import { Home } from './components/home'
+import { Home } from "./components/home";
 
 export default function App() {
   return (

app/requests.ts

@@ -35,6 +35,10 @@ function getHeaders() {
     headers["access-code"] = accessStore.accessCode;
   }
 
+  if (accessStore.token && accessStore.token.length > 0) {
+    headers["token"] = accessStore.token;
+  }
+
   return headers;
 }
 

app/store/access.ts

@@ -4,7 +4,9 @@ import { queryMeta } from "../utils";
 
 export interface AccessControlStore {
   accessCode: string;
+  token: string;
 
+  updateToken: (_: string) => void;
   updateCode: (_: string) => void;
   enabledAccessControl: () => boolean;
 }
@@ -14,6 +16,7 @@ export const ACCESS_KEY = "access-control";
 export const useAccessStore = create<AccessControlStore>()(
   persist(
     (set, get) => ({
+      token: "",
       accessCode: "",
       enabledAccessControl() {
         return queryMeta("access") === "enabled";
@@ -21,6 +24,9 @@ export const useAccessStore = create<AccessControlStore>()(
       updateCode(code: string) {
         set((state) => ({ accessCode: code }));
       },
+      updateToken(token: string) {
+        set((state) => ({ token }));
+      },
     }),
     {
       name: ACCESS_KEY,

app/store/app.ts

@@ -49,22 +49,24 @@ export interface ChatConfig {
 
 export type ModelConfig = ChatConfig["modelConfig"];
 
+const ENABLE_GPT4 = true;
+
 export const ALL_MODELS = [
   {
     name: "gpt-4",
-    available: false,
+    available: ENABLE_GPT4,
   },
   {
     name: "gpt-4-0314",
-    available: false,
+    available: ENABLE_GPT4,
   },
   {
     name: "gpt-4-32k",
-    available: false,
+    available: ENABLE_GPT4,
   },
   {
     name: "gpt-4-32k-0314",
-    available: false,
+    available: ENABLE_GPT4,
   },
   {
     name: "gpt-3.5-turbo",

app/styles/globals.scss

@@ -206,3 +206,36 @@ div.math {
     text-decoration: underline;
   }
 }
+
+pre {
+  position: relative;
+
+  &:hover .copy-code-button {
+    pointer-events: all;
+    transform: translateX(0px);
+    opacity: 0.5;
+  }
+
+  .copy-code-button {
+    position: absolute;
+    right: 10px;
+    cursor: pointer;
+    padding: 0px 5px;
+    background-color: var(--black);
+    color: var(--white);
+    border: var(--border-in-light);
+    border-radius: 10px;
+    transform: translateX(10px);
+    pointer-events: none;
+    opacity: 0;
+    transition: all ease 0.3s;
+
+    &:after {
+      content: "copy";
+    }
+
+    &:hover {
+      opacity: 1;
+    }
+  }
+}

middleware.ts

@@ -8,13 +8,14 @@ export const config = {
 
 export function middleware(req: NextRequest, res: NextResponse) {
   const accessCode = req.headers.get("access-code");
+  const token = req.headers.get("token");
   const hashedCode = md5.hash(accessCode ?? "").trim();
 
   console.log("[Auth] allowed hashed codes: ", [...ACCESS_CODES]);
   console.log("[Auth] got access code:", accessCode);
   console.log("[Auth] hashed access code:", hashedCode);
 
-  if (ACCESS_CODES.size > 0 && !ACCESS_CODES.has(hashedCode)) {
+  if (ACCESS_CODES.size > 0 && !ACCESS_CODES.has(hashedCode) && !token) {
     return NextResponse.json(
       {
         needAccessCode: true,