Merge branch 'main' into pr/79

README.md

@@ -62,7 +62,7 @@ The app will be up and running at http://localhost:5173. Note that you can't dev
 ## 🙋‍♂️ FAQs
 
 - **I'm running into an error when setting up the backend. How can I fix it?** [Try this](https://github.com/abi/screenshot-to-code/issues/3#issuecomment-1814777959). If that still doesn't work, open an issue.
-- **How do I get an OpenAI API key that has the GPT4 Vision model available?** Create an OpenAI account. And then, you need to buy at least $1 worth of credit on the [Billing dashboard](https://platform.openai.com/account/billing/overview).
+- **How do I get an OpenAI API key that has the GPT4 Vision model available?** Create an OpenAI account. And then, you need to buy at least $1 worth of credit on the [Billing dashboard](https://platform.openai.com/account/billing/overview). Also, see [OpenAI docs](https://help.openai.com/en/articles/7102672-how-can-i-access-gpt-4).
 - **How can I provide feedback?** For feedback, feature requests and bug reports, open an issue or ping me on [Twitter](https://twitter.com/_abi_).
 
 ## 📚 Examples

backend/main.py

@@ -96,7 +96,10 @@ async def stream_code(websocket: WebSocket):
     async def process_chunk(content):
         await websocket.send_json({"type": "chunk", "value": content})
 
-    prompt_messages = assemble_prompt(params["image"])
+    if params.get("resultImage") and params["resultImage"]:
+        prompt_messages = assemble_prompt(params["image"], params["resultImage"])
+    else:
+        prompt_messages = assemble_prompt(params["image"])
 
     # Image cache for updates so that we don't have to regenerate images
     image_cache = {}

backend/prompts.py

@@ -2,8 +2,8 @@ SYSTEM_PROMPT = """
 You are an expert Tailwind developer
 You take screenshots of a reference web page from the user, and then build single page apps 
 using Tailwind, HTML and JS.
-You might also be given a screenshot of a web page that you have already built, and asked to
+You might also be given a screenshot(The second image) of a web page that you have already built, and asked to
-update it to look more like the reference image.
+update it to look more like the reference image(The first image).
 
 - Make sure the app looks exactly like the screenshot.
 - Pay close attention to background color, text color, font size, font family, 
@@ -27,21 +27,26 @@ USER_PROMPT = """
 Generate code for a web page that looks exactly like this.
 """
 
-
+def assemble_prompt(image_data_url, result_image_data_url=None):
-def assemble_prompt(image_data_url):
+    content = [
+        {
+            "type": "image_url",
+            "image_url": {"url": image_data_url, "detail": "high"},
+        },
+        {
+            "type": "text",
+            "text": USER_PROMPT,
+        },
+    ]
+    if result_image_data_url:
+        content.insert(1, {
+            "type": "image_url",
+            "image_url": {"url": result_image_data_url, "detail": "high"},
+        })
     return [
         {"role": "system", "content": SYSTEM_PROMPT},
         {
             "role": "user",
-            "content": [
+            "content": content,
-                {
-                    "type": "image_url",
-                    "image_url": {"url": image_data_url, "detail": "high"},
-                },
-                {
-                    "type": "text",
-                    "text": USER_PROMPT,
-                },
-            ],
         },
     ]

frontend/package.json

@@ -22,6 +22,7 @@
     "classnames": "^2.3.2",
     "clsx": "^2.0.0",
     "codemirror": "^6.0.1",
+    "html2canvas": "^1.4.1",
     "copy-to-clipboard": "^3.3.3",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",

frontend/src/App.tsx

@@ -13,6 +13,8 @@ import {
   FaMobile,
   FaUndo,
 } from "react-icons/fa";
+
+import { Switch } from "./components/ui/switch";
 import copy from "copy-to-clipboard";
 import toast from "react-hot-toast";
 import { Button } from "@/components/ui/button";
@@ -27,6 +29,7 @@ import { OnboardingNote } from "./components/OnboardingNote";
 import { usePersistedState } from "./hooks/usePersistedState";
 import { UrlInputSection } from "./components/UrlInputSection";
 import TermsOfServiceDialog from "./components/TermsOfServiceDialog";
+import html2canvas from "html2canvas";
 import { USER_CLOSE_WEB_SOCKET_CODE } from "./constants";
 
 function App() {
@@ -46,8 +49,23 @@ function App() {
     },
     "setting"
   );
+  const [shouldIncludeResultImage, setShouldIncludeResultImage] =
+    useState<boolean>(false);
   const wsRef = useRef<WebSocket>(null);
 
+  const takeScreenshot = async (): Promise<string> => {
+    const iframeElement = document.querySelector(
+      "#preview-desktop"
+    ) as HTMLIFrameElement;
+    if (!iframeElement?.contentWindow?.document.body) {
+      return "";
+    }
+
+    const canvas = await html2canvas(iframeElement.contentWindow.document.body);
+    const png = canvas.toDataURL("image/png");
+    return png;
+  };
+
   const downloadCode = () => {
     // Create a blob from the generated code
     const blob = new Blob([generatedCode], { type: "text/html" });
@@ -75,6 +93,8 @@ function App() {
 
   const stop = () => {
     wsRef.current?.close?.(USER_CLOSE_WEB_SOCKET_CODE);
+    // make sure stop can correct the state even if the websocket is already closed
+    setAppState(AppState.CODE_READY);
   };
 
   function doGenerateCode(params: CodeGenerationParams) {
@@ -106,14 +126,23 @@ function App() {
   }
 
   // Subsequent updates
-  function doUpdate() {
+  async function doUpdate() {
     const updatedHistory = [...history, generatedCode, updateInstruction];
-
+    if (shouldIncludeResultImage) {
-    doGenerateCode({
+      const resultImage = await takeScreenshot();
-      generationType: "update",
+      doGenerateCode({
-      image: referenceImages[0],
+        generationType: "update",
-      history: updatedHistory,
+        image: referenceImages[0],
-    });
+        resultImage: resultImage,
+        history: updatedHistory,
+      });
+    } else {
+      doGenerateCode({
+        generationType: "update",
+        image: referenceImages[0],
+        history: updatedHistory,
+      });
+    }
 
     setHistory(updatedHistory);
     setGeneratedCode("");
@@ -183,6 +212,15 @@ function App() {
                       onChange={(e) => setUpdateInstruction(e.target.value)}
                       value={updateInstruction}
                     />
+                    <div className="flex justify-between items-center gap-x-2">
+                      <div className="font-500 text-xs text-slate-700">
+                        Include screenshot of current version?
+                      </div>
+                      <Switch
+                        checked={shouldIncludeResultImage}
+                        onCheckedChange={setShouldIncludeResultImage}
+                      />
+                    </div>
                     <Button onClick={doUpdate}>Update</Button>
                   </div>
                   <div className="flex items-center gap-x-2 mt-2">

frontend/src/components/Preview.tsx

@@ -23,6 +23,7 @@ function Preview({ code, device }: Props) {
   return (
     <div className="flex justify-center mx-2">
       <iframe
+        id={`preview-${device}`}
         ref={iframeRef}
         title="Preview"
         className={classNames(

frontend/src/generateCode.ts

@@ -10,6 +10,7 @@ const STOP_MESSAGE = "Code generation stopped";
 export interface CodeGenerationParams {
   generationType: "create" | "update";
   image: string;
+  resultImage?: string;
   history?: string[];
   // isImageGenerationEnabled: boolean; // TODO: Merge with Settings type in types.ts
 }
@@ -49,14 +50,11 @@ export function generateCode(
     console.log("Connection closed", event.code, event.reason);
     if (event.code === USER_CLOSE_WEB_SOCKET_CODE) {
       toast.success(STOP_MESSAGE);
-      onComplete();
+    } else if (event.code !== 1000) {
-    } else if (event.code === 1000) {
-      onComplete();
-    } else {
       console.error("WebSocket error code", event);
       toast.error(ERROR_MESSAGE);
-      onComplete();
     }
+    onComplete();
   });
 
   ws.addEventListener("error", (error) => {

frontend/yarn.lock

@@ -1107,6 +1107,11 @@ balanced-match@^1.0.0:
   resolved "https://registry.npmjs.org/balanced-match/-/balanced-match-1.0.2.tgz"
   integrity sha512-3oSeUO0TMV67hN1AmbXsK4yaqU7tjiHlbxRDZOpH0KW9+CeX4bRAaX0Anxt0tx2MrpRpWwQaPwIlISEJhYU5Pw==
 
+base64-arraybuffer@^1.0.2:
+  version "1.0.2"
+  resolved "https://registry.npmjs.org/base64-arraybuffer/-/base64-arraybuffer-1.0.2.tgz#1c37589a7c4b0746e34bd1feb951da2df01c1bdc"
+  integrity sha512-I3yl4r9QB5ZRY3XuJVEPfc2XhZO6YweFPI+UovAzn+8/hb3oJ6lnysaFcjVpkCPfVWFUDvoZ8kmVDP7WyRtYtQ==
+
 binary-extensions@^2.0.0:
   version "2.2.0"
   resolved "https://registry.npmjs.org/binary-extensions/-/binary-extensions-2.2.0.tgz"
@@ -1279,6 +1284,13 @@ cross-spawn@^7.0.2:
     shebang-command "^2.0.0"
     which "^2.0.1"
 
+css-line-break@^2.1.0:
+  version "2.1.0"
+  resolved "https://registry.npmjs.org/css-line-break/-/css-line-break-2.1.0.tgz#bfef660dfa6f5397ea54116bb3cb4873edbc4fa0"
+  integrity sha512-FHcKFCZcAha3LwfVBhCQbW2nCNbkZXn7KVUJcsT5/P8YmfsVja0FMPJr0B903j/E69HUphKiV9iQArX8SDYA4w==
+  dependencies:
+    utrie "^1.0.2"
+
 cssesc@^3.0.0:
   version "3.0.0"
   resolved "https://registry.npmjs.org/cssesc/-/cssesc-3.0.0.tgz"
@@ -1687,6 +1699,14 @@ hasown@^2.0.0:
   dependencies:
     function-bind "^1.1.2"
 
+html2canvas@^1.4.1:
+  version "1.4.1"
+  resolved "https://registry.npmjs.org/html2canvas/-/html2canvas-1.4.1.tgz#7cef1888311b5011d507794a066041b14669a543"
+  integrity sha512-fPU6BHNpsyIhr8yyMpTLLxAbkaK8ArIBcmZIRiBLiDhjeqvXolaEmDGmELFuX9I4xDcaKKcJl+TKZLqruBbmWA==
+  dependencies:
+    css-line-break "^2.1.0"
+    text-segmentation "^1.0.3"
+
 ignore@^5.2.0, ignore@^5.2.4:
   version "5.2.4"
   resolved "https://registry.npmjs.org/ignore/-/ignore-5.2.4.tgz"
@@ -2386,6 +2406,13 @@ tailwindcss@^3.3.5:
     resolve "^1.22.2"
     sucrase "^3.32.0"
 
+text-segmentation@^1.0.3:
+  version "1.0.3"
+  resolved "https://registry.npmjs.org/text-segmentation/-/text-segmentation-1.0.3.tgz#52a388159efffe746b24a63ba311b6ac9f2d7943"
+  integrity sha512-iOiPUo/BGnZ6+54OsWxZidGCsdU8YbE4PSpdPinp7DeMtUJNJBoJ/ouUSTJjHkh1KntHaltHl/gDs2FC4i5+Nw==
+  dependencies:
+    utrie "^1.0.2"
+
 text-table@^0.2.0:
   version "0.2.0"
   resolved "https://registry.npmjs.org/text-table/-/text-table-0.2.0.tgz"
@@ -2514,6 +2541,13 @@ util-deprecate@^1.0.2:
   resolved "https://registry.npmjs.org/util-deprecate/-/util-deprecate-1.0.2.tgz"
   integrity sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==
 
+utrie@^1.0.2:
+  version "1.0.2"
+  resolved "https://registry.npmjs.org/utrie/-/utrie-1.0.2.tgz#d42fe44de9bc0119c25de7f564a6ed1b2c87a645"
+  integrity sha512-1MLa5ouZiOmQzUbjbu9VmjLzn1QLXBhwpUa7kdLUQK+KQ5KA9I1vk5U4YHe/X2Ch7PYnJfWuWT+VbuxbGwljhw==
+  dependencies:
+    base64-arraybuffer "^1.0.2"
+
 vite-plugin-checker@^0.6.2:
   version "0.6.2"
   resolved "https://registry.npmjs.org/vite-plugin-checker/-/vite-plugin-checker-0.6.2.tgz"