Merge 8ff1bf1976 into 9f732c4f5d

backend/prompts.py

@@ -0,0 +1,98 @@
+from typing import List, Union
+
+from openai.types.chat import ChatCompletionMessageParam, ChatCompletionContentPartParam
+
+from imported_code_prompts import (
+    IMPORTED_CODE_BOOTSTRAP_SYSTEM_PROMPT,
+    IMPORTED_CODE_IONIC_TAILWIND_SYSTEM_PROMPT,
+    IMPORTED_CODE_REACT_TAILWIND_SYSTEM_PROMPT,
+    IMPORTED_CODE_TAILWIND_SYSTEM_PROMPT,
+)
+from screenshot_system_prompts import (
+    BOOTSTRAP_SYSTEM_PROMPT,
+    IONIC_TAILWIND_SYSTEM_PROMPT,
+    REACT_TAILWIND_SYSTEM_PROMPT,
+    TAILWIND_SYSTEM_PROMPT,
+)
+
+
+USER_PROMPT = """
+Generate code for a web page that looks exactly like these images. Seperate them into different sectionsif multiple images are provided.
+"""
+
+
+def assemble_imported_code_prompt(
+    code: str, stack: str, result_image_data_url: Union[str, None] = None
+) -> List[ChatCompletionMessageParam]:
+    system_content = IMPORTED_CODE_TAILWIND_SYSTEM_PROMPT
+    if stack == "html_tailwind":
+        system_content = IMPORTED_CODE_TAILWIND_SYSTEM_PROMPT
+    elif stack == "react_tailwind":
+        system_content = IMPORTED_CODE_REACT_TAILWIND_SYSTEM_PROMPT
+    elif stack == "bootstrap":
+        system_content = IMPORTED_CODE_BOOTSTRAP_SYSTEM_PROMPT
+    elif stack == "ionic_tailwind":
+        system_content = IMPORTED_CODE_IONIC_TAILWIND_SYSTEM_PROMPT
+    else:
+        raise Exception("Code config is not one of available options")
+
+    return [
+        {
+            "role": "system",
+            "content": system_content,
+        },
+        {
+            "role": "user",
+            "content": "Here is the code of the app: " + code,
+        },
+    ]
+    # TODO: Use result_image_data_url
+
+
+def assemble_prompt(
+    images: List[str],
+    generated_code_config: str,
+    result_image_data_url: Union[str, None] = None,
+) -> List[ChatCompletionMessageParam]:
+    # Set the system prompt based on the output settings
+    system_content = TAILWIND_SYSTEM_PROMPT
+    if generated_code_config == "html_tailwind":
+        system_content = TAILWIND_SYSTEM_PROMPT
+    elif generated_code_config == "react_tailwind":
+        system_content = REACT_TAILWIND_SYSTEM_PROMPT
+    elif generated_code_config == "bootstrap":
+        system_content = BOOTSTRAP_SYSTEM_PROMPT
+    elif generated_code_config == "ionic_tailwind":
+        system_content = IONIC_TAILWIND_SYSTEM_PROMPT
+    else:
+        raise Exception("Code config is not one of available options")
+
+    user_content: List[ChatCompletionContentPartParam]= []
+    for image_data_url in images:
+        user_content.append(
+            {
+                "type": "image_url",
+                "image_url": {"url": image_data_url, "detail": "high"},
+            },
+        )
+
+    if result_image_data_url:
+        user_content.insert(1, {
+            "type": "image_url",
+            "image_url": {"url": result_image_data_url, "detail": "high"},
+        })
+
+    user_content.append(
+        {
+            "type": "text",
+            "text": USER_PROMPT,
+        }
+    )
+
+    return [
+        {"role": "system", "content": system_content},
+        {
+            "role": "user",
+            "content": user_content,
+        },
+    ]

backend/routes/generate_code.py

@@ -200,10 +200,12 @@ async def stream_code(websocket: WebSocket):
         try:
             if params.get("resultImage") and params["resultImage"]:
                 prompt_messages = assemble_prompt(
-                    params["image"], valid_stack, params["resultImage"]
+                    params["images"], generated_code_config, params["resultImage"]
                 )
             else:
-                prompt_messages = assemble_prompt(params["image"], valid_stack)
+                prompt_messages = assemble_prompt(
+                    params["images"], generated_code_config
+                )
         except:
             await websocket.send_json(
                 {

frontend/src/App.tsx

@@ -288,8 +288,7 @@ function App() {
       doGenerateCode(
         {
           generationType: "create",
-          image: referenceImages[0],
-          inputMode,
+          images: referenceImages,
         },
         currentVersion
       );
@@ -340,8 +339,7 @@ function App() {
       doGenerateCode(
         {
           generationType: "update",
-          inputMode,
-          image: referenceImages[0],
+          images: referenceImages,
           resultImage: resultImage,
           history: updatedHistory,
           isImportedFromCode,
@@ -352,8 +350,7 @@ function App() {
       doGenerateCode(
         {
           generationType: "update",
-          inputMode,
-          image: referenceImages[0],
+          images: referenceImages,
           history: updatedHistory,
           isImportedFromCode,
         },
@@ -556,37 +553,34 @@ function App() {
 
               {/* Reference image display */}
               <div className="flex gap-x-2 mt-2">
-                {referenceImages.length > 0 && (
-                  <div className="flex flex-col">
-                    <div
-                      className={classNames({
-                        "scanning relative": appState === AppState.CODING,
-                      })}
-                    >
-                      {inputMode === "image" && (
-                        <img
-                          className="w-[340px] border border-gray-200 rounded-md"
-                          src={referenceImages[0]}
-                          alt="Reference"
-                        />
-                      )}
-                      {inputMode === "video" && (
-                        <video
-                          muted
-                          autoPlay
-                          loop
-                          className="w-[340px] border border-gray-200 rounded-md"
-                          src={referenceImages[0]}
-                        />
-                      )}
-                    </div>
-                    <div className="text-gray-400 uppercase text-sm text-center mt-1">
-                      {inputMode === "video"
-                        ? "Original Video"
-                        : "Original Screenshot"}
-                    </div>
+                <div className="flex flex-col">
+                  <div
+                    className={classNames({
+                      "scanning relative": appState === AppState.CODING,
+                    })}
+                  >
+                    {referenceImages.length > 1 ? (
+                      <div className="grid grid-cols-2 gap-2">
+                        {referenceImages.map((image, i) => (
+                          <img
+                            className="border max-h-40 w-full h-full object-contain border-gray-200 rounded-md"
+                            src={image}
+                            alt={`Reference image: ${i}`}
+                          />
+                        ))}
+                      </div>
+                    ) : (
+                      <img
+                        className="w-[340px] border border-gray-200 rounded-md"
+                        src={referenceImages[0]}
+                        alt-="Reference Image"
+                      />
+                    )}
                   </div>
-                )}
+                  <div className="text-gray-400 uppercase text-sm text-center mt-1">
+                    Original Screenshot{referenceImages.length > 1 ? "s" : ""}
+                  </div>
+                </div>
                 <div className="bg-gray-400 px-4 py-2 rounded text-sm hidden">
                   <h2 className="text-lg mb-4 border-b border-gray-800">
                     Console

frontend/src/components/ImageUpload.tsx

@@ -69,8 +69,7 @@ function ImageUpload({ setReferenceImages }: Props) {
 
   const { getRootProps, getInputProps, isFocused, isDragAccept, isDragReject } =
     useDropzone({
-      maxFiles: 1,
-      maxSize: 1024 * 1024 * 20, // 20 MB
+      maxSize: 1024 * 1024 * 5, // 5 MB
       accept: {
         // Image formats
         "image/png": [".png"],

frontend/src/types.ts

@@ -33,8 +33,7 @@ export enum ScreenRecorderState {
 
 export interface CodeGenerationParams {
   generationType: "create" | "update";
-  inputMode: "image" | "video";
-  image: string;
+  images: string[];
   resultImage?: string;
   history?: string[];
   isImportedFromCode?: boolean;