diff --git a/backend/prompts.py b/backend/prompts.py new file mode 100644 index 0000000..222cc20 --- /dev/null +++ b/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, + }, + ] diff --git a/backend/routes/generate_code.py b/backend/routes/generate_code.py index 19ef382..02be7f1 100644 --- a/backend/routes/generate_code.py +++ b/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( { diff --git a/frontend/src/App.tsx b/frontend/src/App.tsx index e3b5b64..548ae66 100644 --- a/frontend/src/App.tsx +++ b/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 */}
- {referenceImages.length > 0 && ( -
-
- {inputMode === "image" && ( - Reference - )} - {inputMode === "video" && ( -
-
- {inputMode === "video" - ? "Original Video" - : "Original Screenshot"} -
+
+
+ {referenceImages.length > 1 ? ( +
+ {referenceImages.map((image, i) => ( + {`Reference + ))} +
+ ) : ( + + )}
- )} +
+ Original Screenshot{referenceImages.length > 1 ? "s" : ""} +
+

Console diff --git a/frontend/src/components/ImageUpload.tsx b/frontend/src/components/ImageUpload.tsx index bbecd92..e862dc2 100644 --- a/frontend/src/components/ImageUpload.tsx +++ b/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"], diff --git a/frontend/src/types.ts b/frontend/src/types.ts index dbe10ad..9f6d837 100644 --- a/frontend/src/types.ts +++ b/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;