Chat api key security (#253)

llm-ui-kit · May 31, 2024 · d93a48e · d93a48e
1 parent 50c88d8
commit d93a48e
Show file tree

Hide file tree

Showing 3 changed files with 91 additions and 39 deletions.
diff --git a/apps/www/src/components/Chat.tsx b/apps/www/src/components/Chat.tsx
@@ -9,6 +9,7 @@ import {
   SelectTrigger,
 } from "@/components/ui/Select";
 import { Icons } from "@/icons";
+import { getMaskedKey } from "@/lib/maskKey";
 import { cn } from "@/lib/utils";
 import { nanoid, type Message } from "ai";
 import { useChat } from "ai/react";
@@ -57,6 +58,7 @@ export const Chat = () => {
   const [selectedChatGptModel, setSelectedChatGptModel] =
     React.useState<string>(CHAT_GPT_MODELS[0]);
   const [systemMessage, setSystemMessage] = React.useState<string>("");
+  const [error, setError] = React.useState<Error>();
   const { messages, input, handleInputChange, handleSubmit, isLoading } =
     useChat({
       api: "/api/chat",
@@ -74,6 +76,9 @@ export const Chat = () => {
         apiKey: currentApiKey,
         model: selectedChatGptModel,
       },
+      onError: (error: Error) => {
+        setError(JSON.parse(error.message));
+      },
     });
 
   const scrollToBottom = () => {
@@ -85,13 +90,16 @@ export const Chat = () => {
   const handleUpdateApiKey = (e: React.ChangeEvent<HTMLInputElement>) => {
     const newApiKey = e.target.value;
     setCurrentApiKey(newApiKey);
+
+    e.currentTarget.blur();
   };
 
   const handleUpdateChatGptModel = (value: string) => {
     setSelectedChatGptModel(value);
   };
 
   const onSubmit = (e: React.FormEvent<HTMLFormElement>) => {
+    setError(undefined);
     storage?.setItem(CHAT_OPENAI_API_KEY, currentApiKey);
     scrollToBottom();
     handleSubmit(e);
@@ -119,10 +127,24 @@ export const Chat = () => {
           </SelectContent>
         </Select>
         <Input
-          value={currentApiKey}
+          value={getMaskedKey(currentApiKey)}
+          onKeyDown={(e) => {
+            if (!((e.ctrlKey || e.metaKey) && e.key === "v")) {
+              e.preventDefault();
+            }
+          }}
+          onFocus={(e) => {
+            e.currentTarget.select();
+          }}
           className="focus-within:border-white"
-          placeholder="Enter Your API Key"
+          placeholder="Paste Your API Key"
           onChange={handleUpdateApiKey}
+          onDragStart={(e) => e.preventDefault()}
+          onDragOver={(e) => e.preventDefault()}
+          onMouseDown={(e) => {
+            e.preventDefault();
+            e.currentTarget.focus();
+          }}
         />
       </div>
       <form
@@ -178,33 +200,40 @@ export const Chat = () => {
             </div>
           </div>
         </div>
-        <div className="bg-background w-full flex flex-row overflow-hidden focus-within:border-white relative px-3 py-1 shadow-lg mb-6 sm:rounded-xl sm:border md:py-3 max-w-2xl mx-auto">
-          <AutosizeTextarea
-            onKeyDown={(e) => {
-              if (isLoading) {
-                return;
-              }
-              if (e.key === "Enter" && !e.shiftKey) {
-                e.preventDefault();
-                onSubmit(e as unknown as React.FormEvent<HTMLFormElement>);
-              }
-            }}
-            placeholder="Message ChatGPT"
-            value={input}
-            rows={1}
-            style={{ height: 42 }}
-            minHeight={42}
-            maxHeight={200}
-            onChange={handleInputChange}
-            className="focus-visible:ring-0 pr-0 resize-none bg-transparent focus-within:outline-none sm:text-base border-none"
-          />
-          <Button
-            disabled={isLoading || !input}
-            className="self-end"
-            type="submit"
-          >
-            Run <Icons.return className="size-4 ml-2" />
-          </Button>
+        <div className="flex flex-col w-full max-w-2xl mx-auto">
+          {error && (
+            <div className="text-destructive text-center mb-4 font-bold p-2">
+              <p>{error?.message}</p>
+            </div>
+          )}
+          <div className="bg-background flex flex-row overflow-hidden focus-within:border-white px-3 py-1 shadow-lg mb-6 sm:rounded-xl sm:border md:py-3">
+            <AutosizeTextarea
+              onKeyDown={(e) => {
+                if (isLoading) {
+                  return;
+                }
+                if (e.key === "Enter" && !e.shiftKey) {
+                  e.preventDefault();
+                  onSubmit(e as unknown as React.FormEvent<HTMLFormElement>);
+                }
+              }}
+              placeholder="Message ChatGPT"
+              value={input}
+              rows={1}
+              style={{ height: 42 }}
+              minHeight={42}
+              maxHeight={200}
+              onChange={handleInputChange}
+              className="focus-visible:ring-0 pr-0 resize-none bg-transparent focus-within:outline-none sm:text-base border-none"
+            />
+            <Button
+              disabled={isLoading || !input}
+              className="self-end"
+              type="submit"
+            >
+              Run <Icons.return className="size-4 ml-2" />
+            </Button>
+          </div>
         </div>
       </form>
     </div>

diff --git a/apps/www/src/lib/maskKey.ts b/apps/www/src/lib/maskKey.ts
@@ -0,0 +1,5 @@
+export const getMaskedKey = (value: string): string => {
+  const charsToShow = 3;
+  const maskedLength = Math.max(0, value.length - charsToShow);
+  return `${value.substring(0, charsToShow)}${"*".repeat(maskedLength)}`;
+};
diff --git a/apps/www/src/pages/api/chat.ts b/apps/www/src/pages/api/chat.ts
@@ -1,3 +1,5 @@
+/* eslint-disable  @typescript-eslint/no-explicit-any */
+
 import { OpenAIStream, StreamingTextResponse } from "ai";
 import type { APIRoute } from "astro";
 import OpenAI from "openai";
@@ -11,6 +13,13 @@ const chatRequestSchema = z.object({
   model: z.string(),
 });
 
+const getApiKey = (apiKey: string, model: string) => {
+  if (model === "gpt-3.5-turbo" && apiKey.length === 0) {
+    return process.env.OPENAI_API_KEY;
+  }
+  return apiKey;
+};
+
 export const POST: APIRoute = async ({ request }) => {
   const jsonBody = await request.json();
   const result = chatRequestSchema.safeParse(jsonBody);
@@ -23,15 +32,24 @@ export const POST: APIRoute = async ({ request }) => {
     );
   }
   const { apiKey, messages, model } = result.data;
-  const openai = new OpenAI({
-    apiKey,
-  });
-  const completion = await openai.chat.completions.create({
-    model,
-    messages: messages,
-    stream: true,
-  });
-  const stream = OpenAIStream(completion);
+  try {
+    const openai = new OpenAI({
+      apiKey: getApiKey(apiKey, model),
+    });
+    const completion = await openai.chat.completions.create({
+      model,
+      messages: messages,
+      stream: true,
+    });
+    const stream = OpenAIStream(completion);
 
-  return new StreamingTextResponse(stream);
+    return new StreamingTextResponse(stream);
+  } catch (error: any) {
+    return new Response(
+      JSON.stringify({
+        message: `${error.message}`,
+      }),
+      { status: 500 },
+    );
+  }
 };