From 7744518079f7e777ff65d0c738ddaa9766880ba0 Mon Sep 17 00:00:00 2001
From: clean99 <xff9924@gmail.com>
Date: Sun, 26 Nov 2023 15:13:39 +0800
Subject: [PATCH] feat: optimaze init flow

---
 backend/main.py                             |  8 ++++++-
 backend/prompts.py                          |  2 +-
 frontend/src/App.tsx                        |  2 +-
 frontend/src/components/ReactCodeEditor.tsx | 26 ++++++++++++++++-----
 4 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/backend/main.py b/backend/main.py
index 93d2e9a..0a30a4c 100644
--- a/backend/main.py
+++ b/backend/main.py
@@ -100,11 +100,17 @@ async def stream_code(websocket: WebSocket):
     if params.get("resultImage") and params["resultImage"]:
         prompt_messages = assemble_prompt(params["image"], params["resultImage"])
     else:
-        prompt_messages = assemble_prompt(params["image"])
+        if params.get("codeType") and params["codeType"] == "react":
+            prompt_messages = assemble_prompt(
+                params["image"], is_react=True
+            )
+        else:
+            prompt_messages = assemble_prompt(params["image"])
 
     # Image cache for updates so that we don't have to regenerate images
     image_cache = {}
 
+
     if params["generationType"] == "update":
         # Transform into message format
         # TODO: Move this to frontend
diff --git a/backend/prompts.py b/backend/prompts.py
index 9b5858e..3628897 100644
--- a/backend/prompts.py
+++ b/backend/prompts.py
@@ -39,7 +39,7 @@ padding, margin, border, etc. Match the colors and sizes exactly.
 In terms of libraries, assume that all libraries are already install, you can import it if you want to use it.
 
 Strictly return only the full jsx code,
-Do not include markdown "```" or "```jsx" at the start or end, do not include other information or explanation.
+Strictly do not return markdown "```" or "```jsx" at the start or end, do not return other information or explanation.
 """
 
 USER_PROMPT = """
diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx
index fdd7ad6..7dd6ca2 100644
--- a/frontend/src/App.tsx
+++ b/frontend/src/App.tsx
@@ -302,7 +302,7 @@ function App() {
                     Code
                   </TabsTrigger>
                 </TabsList>
-                <ReactCodeEditor doGenerateCode={doGenerateCode} referenceImage={referenceImages[0]} />
+                <ReactCodeEditor doGenerateCode={doGenerateCode} referenceImage={referenceImages[0]} appState={appState} />
               </div>
               <TabsContent value="desktop">
                 <Preview code={generatedCode} device="desktop" />
diff --git a/frontend/src/components/ReactCodeEditor.tsx b/frontend/src/components/ReactCodeEditor.tsx
index 829ce5c..131a7ae 100644
--- a/frontend/src/components/ReactCodeEditor.tsx
+++ b/frontend/src/components/ReactCodeEditor.tsx
@@ -13,7 +13,7 @@ import { useState } from "react";
 import { FaReact } from "react-icons/fa"
 import CodeMirror from "./CodeMirror";
 import { usePersistedState } from "@/hooks/usePersistedState";
-import { EditorTheme, Settings } from "@/types";
+import { AppState, EditorTheme, Settings } from "@/types";
 import { doCopyCode } from "@/lib/utils";
 import { Textarea } from "./ui/textarea";
 
@@ -21,9 +21,10 @@ import { Textarea } from "./ui/textarea";
 interface IProps {
     doGenerateCode: (params: CodeGenerationParams, setCode: (value: React.SetStateAction<string>) => void) => void;
     referenceImage: string;
+    appState: AppState;
 }
 
-export const ReactCodeEditor: React.FC<IProps> = ({ doGenerateCode, referenceImage }) => {
+export const ReactCodeEditor: React.FC<IProps> = ({ doGenerateCode, referenceImage, appState }) => {
     const [generatedReactCode, setGeneratedReactCode] = useState("");
     const [updateInstruction, setUpdateInstruction] = useState("");
     const [reactHistory, setReactHistory] = useState<string[]>([]);
@@ -51,19 +52,32 @@ export const ReactCodeEditor: React.FC<IProps> = ({ doGenerateCode, referenceIma
         setUpdateInstruction("");
     }
 
+    async function init() {
+        doGenerateCode({
+            generationType: "create",
+            codeType: CodeType.REACT,
+            image: referenceImage,
+            history: [],
+        }, setGeneratedReactCode);
+    }
+
     return (
-    <Dialog>
+    <Dialog onOpenChange={(open: boolean) => {
+        if (open && generatedReactCode === "") {
+            init();
+        }
+    }}>
         <DialogTrigger asChild>
         <Button className="flex items-center gap-x-2 ml-4"><FaReact /> Generate React Code</Button>
         </DialogTrigger>
-        <DialogContent className="sm:min-w-[425px] lg:min-w-[1240px] min-h-[600px] flex flex-col justify-between">
+        <DialogContent className="sm:min-w-[425px] lg:min-w-[1240px] h-[660px] flex flex-col justify-between">
             <DialogHeader>
                 <DialogTitle>Generate React Code</DialogTitle>
                 <DialogDescription>
                 Notice: The code below may shows different results than the HTML code.
                 </DialogDescription>
             </DialogHeader>
-            <div className="gap-2 flex-1">
+            <div className="gap-2 flex-1 overflow-scroll">
                 <CodeMirror
                     editorTheme={settings.editorTheme}
                     code={generatedReactCode}
@@ -77,7 +91,7 @@ export const ReactCodeEditor: React.FC<IProps> = ({ doGenerateCode, referenceIma
                 value={updateInstruction}
             />
             <DialogFooter>
-                <Button onClick={doUpdateReact}>Update</Button>
+                <Button disabled={appState === AppState.CODING} onClick={doUpdateReact}>Update</Button>
             </DialogFooter>
         </DialogContent>
     </Dialog>