support setting openai api key on the client side

backend/image_generation.py

@@ -5,8 +5,8 @@ from openai import AsyncOpenAI
 from bs4 import BeautifulSoup
 
 
-async def process_tasks(prompts):
-    tasks = [generate_image(prompt) for prompt in prompts]
+async def process_tasks(prompts, api_key):
+    tasks = [generate_image(prompt, api_key) for prompt in prompts]
     results = await asyncio.gather(*tasks, return_exceptions=True)
 
     processed_results = []
@@ -20,8 +20,8 @@ async def process_tasks(prompts):
     return processed_results
 
 
-async def generate_image(prompt):
-    client = AsyncOpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
+async def generate_image(prompt, api_key):
+    client = AsyncOpenAI(api_key=api_key)
     image_params = {
         "model": "dall-e-3",
         "quality": "standard",
@@ -60,7 +60,7 @@ def create_alt_url_mapping(code):
     return mapping
 
 
-async def generate_images(code, image_cache):
+async def generate_images(code, api_key, image_cache):
     # Find all images
     soup = BeautifulSoup(code, "html.parser")
     images = soup.find_all("img")
@@ -87,7 +87,7 @@ async def generate_images(code, image_cache):
         return code
 
     # Generate images
-    results = await process_tasks(prompts)
+    results = await process_tasks(prompts, api_key)
 
     # Create a dict mapping alt text to image URL
     mapped_image_urls = dict(zip(prompts, results))

backend/llm.py

@@ -4,10 +4,12 @@ from openai import AsyncOpenAI
 
 MODEL_GPT_4_VISION = "gpt-4-vision-preview"
 
-client = AsyncOpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
 
+async def stream_openai_response(
+    messages, api_key: str, callback: Callable[[str], Awaitable[None]]
+):
+    client = AsyncOpenAI(api_key=api_key)
 
-async def stream_openai_response(messages, callback: Callable[[str], Awaitable[None]]):
     model = MODEL_GPT_4_VISION
 
     # Base parameters

backend/main.py

@@ -41,6 +41,25 @@ async def stream_code_test(websocket: WebSocket):
 
     params = await websocket.receive_json()
 
+    # Get the OpenAI API key from the request. Fall back to environment variable if not provided.
+    # If neither is provided, we throw an error.
+    if params["openAiApiKey"]:
+        openai_api_key = params["openAiApiKey"]
+        print("Using OpenAI API key from client-side settings dialog")
+    else:
+        openai_api_key = os.environ.get("OPENAI_API_KEY")
+        print("Using OpenAI API key from environment variable")
+
+    if not openai_api_key:
+        print("OpenAI API key not found")
+        await websocket.send_json(
+            {
+                "type": "error",
+                "value": "OpenAI API key found. Please add your API key in the settings dialog or add it to backend/.env file.",
+            }
+        )
+        return
+
     should_generate_images = (
         params["isImageGenerationEnabled"]
         if "isImageGenerationEnabled" in params
@@ -73,7 +92,8 @@ async def stream_code_test(websocket: WebSocket):
     else:
         completion = await stream_openai_response(
             prompt_messages,
-            lambda x: process_chunk(x),
+            api_key=openai_api_key,
+            callback=lambda x: process_chunk(x),
         )
 
     # Write the messages dict into a log so that we can debug later
@@ -84,7 +104,9 @@ async def stream_code_test(websocket: WebSocket):
             await websocket.send_json(
                 {"type": "status", "value": "Generating images..."}
             )
-            updated_html = await generate_images(completion, image_cache=image_cache)
+            updated_html = await generate_images(
+                completion, api_key=openai_api_key, image_cache=image_cache
+            )
         else:
             updated_html = completion
         await websocket.send_json({"type": "setCode", "value": updated_html})

frontend/src/App.tsx

@@ -29,6 +29,7 @@ function App() {
   const [updateInstruction, setUpdateInstruction] = useState("");
   const [history, setHistory] = useState<string[]>([]);
   const [settings, setSettings] = useState<Settings>({
+    openAiApiKey: null,
     isImageGenerationEnabled: true,
   });
 

frontend/src/components/SettingsDialog.tsx

@@ -1,6 +1,8 @@
 import {
   Dialog,
+  DialogClose,
   DialogContent,
+  DialogFooter,
   DialogHeader,
   DialogTitle,
   DialogTrigger,
@@ -9,6 +11,7 @@ import { FaCog } from "react-icons/fa";
 import { Settings } from "../types";
 import { Switch } from "./ui/switch";
 import { Label } from "./ui/label";
+import { Input } from "./ui/input";
 
 interface Props {
   settings: Settings;
@@ -24,6 +27,7 @@ function SettingsDialog({ settings, setSettings }: Props) {
       <DialogContent>
         <DialogHeader>
           <DialogTitle className="mb-4">Settings</DialogTitle>
+        </DialogHeader>
         <div className="flex items-center space-x-2">
           <Label htmlFor="image-generation">
             <div>DALL-E Placeholder Image Generation</div>
@@ -42,7 +46,29 @@ function SettingsDialog({ settings, setSettings }: Props) {
             }
           />
         </div>
-        </DialogHeader>
+        <div className="flex flex-col space-y-4">
+          <Label htmlFor="openai-api-key">
+            <div>OpenAI API key</div>
+            <div className="font-light mt-2">
+              Never stored. Overrides your .env config.
+            </div>
+          </Label>
+
+          <Input
+            id="openai-api-key"
+            placeholder="OpenAI API key"
+            value={settings.openAiApiKey || ""}
+            onChange={(e) =>
+              setSettings((s) => ({
+                ...s,
+                openAiApiKey: e.target.value,
+              }))
+            }
+          />
+        </div>
+        <DialogFooter>
+          <DialogClose>Save</DialogClose>
+        </DialogFooter>
       </DialogContent>
     </Dialog>
   );

frontend/src/components/ui/input.tsx

@@ -0,0 +1,25 @@
+import * as React from "react"
+
+import { cn } from "@/lib/utils"
+
+export interface InputProps
+  extends React.InputHTMLAttributes<HTMLInputElement> {}
+
+const Input = React.forwardRef<HTMLInputElement, InputProps>(
+  ({ className, type, ...props }, ref) => {
+    return (
+      <input
+        type={type}
+        className={cn(
+          "flex h-9 w-full rounded-md border border-input bg-transparent px-3 py-1 text-sm shadow-sm transition-colors file:border-0 file:bg-transparent file:text-sm file:font-medium placeholder:text-muted-foreground focus-visible:outline-none focus-visible:ring-1 focus-visible:ring-ring disabled:cursor-not-allowed disabled:opacity-50",
+          className
+        )}
+        ref={ref}
+        {...props}
+      />
+    )
+  }
+)
+Input.displayName = "Input"
+
+export { Input }

frontend/src/generateCode.ts

@@ -36,6 +36,9 @@ export function generateCode(
       onStatusUpdate(response.value);
     } else if (response.type === "setCode") {
       onSetCode(response.value);
+    } else if (response.type === "error") {
+      console.error("Error generating code", response.value);
+      toast.error(response.value);
     }
   });
 

frontend/src/types.ts

@@ -1,3 +1,4 @@
 export interface Settings {
+  openAiApiKey: string | null;
   isImageGenerationEnabled: boolean;
 }