User to add keys via frontend UI

client/src/App.tsx

@@ -8,6 +8,8 @@ const App = () => {
   const [models, setModels] = useState<[]>();
   const [model, setModel] = useState("");
   const [message, setMessage] = useState("");
+  const [apiKey, setApiKey] = useState("");
+  const [keyAlert, setKeyAlert] = useState(false);
   const [imageFile, setImageFile] = useState("./blur.jpg");
 
   // Populates the model dropdown
@@ -25,12 +27,15 @@ const App = () => {
 
   const submitPrompt = async () => {
     try {
+      const userApiKey = sessionStorage.getItem("userApiKey");
+      if (!userApiKey) throw new Error("Please enter your API key");
       if (!model) throw new Error("Please select a model");
       if (!message) throw new Error("Empty prompt");
       const res = await fetch("http://localhost:8000/image", {
         method: "POST",
         headers: {
           "Content-Type": "application/json",
+          "X-User-Api-Key": userApiKey,
         },
         body: JSON.stringify({ model, message }),
       });
@@ -43,6 +48,15 @@ const App = () => {
     }
   };
 
+  const submitKey = (e: React.MouseEvent<HTMLButtonElement>) => {
+    e.preventDefault();
+    if (!apiKey) throw new Error("Please enter API key");
+    sessionStorage.setItem("userApiKey", apiKey);
+    setApiKey("");
+    setKeyAlert(true);
+    setTimeout(() => setKeyAlert(false), 3000);
+  };
+
   useEffect(() => {
     fetchModels();
   }, []);
@@ -54,6 +68,10 @@ const App = () => {
           models={models}
           selectedModel={model}
           handleModel={setModel}
+          apiKey={apiKey}
+          handleKey={setApiKey}
+          submitKey={submitKey}
+          keyAlert={keyAlert}
         />
         <Prompt
           message={message}

client/src/components/Header/header.css

@@ -11,3 +11,9 @@
 .header-button {
   visibility: hidden;
 }
+
+.key-button {
+  background-color: #000;
+  color: #fff;
+  border: none;
+}

client/src/components/ModelSelector/ModelSelector.css

@@ -1,5 +1,10 @@
-.model-selector {
+.model-selector,
+.api-key-container {
   width: 80%;
   display: block;
   margin: auto;
 }
+
+.api-key {
+  margin: 1vh;
+}

client/src/components/ModelSelector/ModelSelector.tsx

@@ -4,31 +4,57 @@ interface ModelSelectorProps {
   models: [] | undefined;
   selectedModel: string;
   handleModel: (value: string) => void;
+  apiKey: string;
+  handleKey: React.Dispatch<React.SetStateAction<string>>;
+  submitKey: (value: React.MouseEvent<HTMLButtonElement>) => void;
+  keyAlert: boolean;
 }
 
 const ModelSelector = ({
   models,
   selectedModel,
   handleModel,
+  apiKey,
+  handleKey,
+  submitKey,
+  keyAlert,
 }: ModelSelectorProps) => {
   return (
-    <select
-      className="model-selector"
-      name="model-selector"
-      id="model-selector"
-      value={selectedModel}
-      onChange={(e) => handleModel(e.target.value)}
-    >
-      <option key="">Select a model . . .</option>
-      {models &&
-        models.map((model: { id: number; name: string }) => {
-          return (
-            <option key={model.id} value={model.id}>
-              {model.name}
-            </option>
-          );
-        })}
-    </select>
+    <>
+      <div className="api-key-container">
+        <label htmlFor="api-key">API Key:</label>
+        <input
+          className="api-key"
+          type="password"
+          name="api-key"
+          id="api-key"
+          value={apiKey}
+          onChange={(e) => handleKey(e.target.value)}
+          placeholder="Enter your API key here..."
+        />
+        <button onClick={submitKey}>Save key</button>
+      </div>
+      <p style={{ visibility: keyAlert ? "visible" : "hidden" }}>
+        Key saved successfully
+      </p>
+      <select
+        className="model-selector"
+        name="model-selector"
+        id="model-selector"
+        value={selectedModel}
+        onChange={(e) => handleModel(e.target.value)}
+      >
+        <option key="">Select a model . . .</option>
+        {models &&
+          models.map((model: { id: number; name: string }) => {
+            return (
+              <option key={model.id} value={model.id}>
+                {model.name}
+              </option>
+            );
+          })}
+      </select>
+    </>
   );
 };
 

client/src/index.css

@@ -2,3 +2,7 @@ body {
   background-color: #000;
   color: #fff;
 }
+
+button:hover {
+  cursor: pointer;
+}

server/main.py

@@ -1,9 +1,7 @@
 import requests
 import base64
 import datetime
-import os
-from dotenv import load_dotenv
-from fastapi import FastAPI
+from fastapi import FastAPI, Header, Depends, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 from pydantic import BaseModel, Field
@@ -24,9 +22,10 @@
     allow_headers=["*"],
 )
 
-load_dotenv()
-
-API_KEY_REF = os.getenv("API_KEY_REF")
+async def get_user_api_key(x_user_api_key: str = Header(...)):
+    if not x_user_api_key:
+        raise HTTPException(status_code=401, detail="API key required")
+    return x_user_api_key
 
 # Mounts image directory as route
 app.mount("/images", StaticFiles(directory="images"), name="images")
@@ -39,7 +38,7 @@ class Prompt(BaseModel):
 
 
 @app.post("/image", status_code=200)
-async def image(prompt: Prompt):
+async def image(prompt: Prompt, API_KEY_REF: str = Depends(get_user_api_key)):
     file_name = f"./images/{datetime.datetime.now()}.png"
     url = "https://openrouter.ai/api/v1/chat/completions"