diff --git a/backend/main.py b/backend/main.py index 72b4dd7..c97067d 100644 --- a/backend/main.py +++ b/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 = {} diff --git a/backend/prompts.py b/backend/prompts.py index 3761404..1ba254d 100644 --- a/backend/prompts.py +++ b/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 -update it to look more like the reference image. +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(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): +def assemble_prompt(image_data_url, result_image_data_url=None): + 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": [ - { - "type": "image_url", - "image_url": {"url": image_data_url, "detail": "high"}, - }, - { - "type": "text", - "text": USER_PROMPT, - }, - ], + "content": content, }, ] diff --git a/frontend/package.json b/frontend/package.json index 912d0b7..ef840e7 100644 --- a/frontend/package.json +++ b/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", diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index 5c946dc..4c843d6 100644 --- a/frontend/src/App.tsx +++ b/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(false); const wsRef = useRef(null); + const takeScreenshot = async (): Promise => { + 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]; - - doGenerateCode({ - generationType: "update", - image: referenceImages[0], - history: updatedHistory, - }); + if (shouldIncludeResultImage) { + const resultImage = await takeScreenshot(); + doGenerateCode({ + generationType: "update", + 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} /> +
+
+ Include screenshot of current version? +
+ +
diff --git a/frontend/src/components/Preview.tsx b/frontend/src/components/Preview.tsx index 3a603ff..24f33c0 100644 --- a/frontend/src/components/Preview.tsx +++ b/frontend/src/components/Preview.tsx @@ -23,6 +23,7 @@ function Preview({ code, device }: Props) { return (