add experimental support for SVG

sud0su · Dec 14, 2023 · 9b728d0 · 9b728d0
1 parent f676151
commit 9b728d0
Show file tree

Hide file tree

Showing 8 changed files with 129 additions and 8 deletions.
diff --git a/backend/eval.py b/backend/eval.py
@@ -49,7 +49,7 @@ async def main():
     for filename in evals:
         filepath = os.path.join(INPUT_DIR, filename)
         data_url = await image_to_data_url(filepath)
-        task = generate_code_core(data_url, "html_tailwind")
+        task = generate_code_core(data_url, "svg")
         tasks.append(task)
 
     results = await asyncio.gather(*tasks)

diff --git a/backend/imported_code_prompts.py b/backend/imported_code_prompts.py
@@ -78,3 +78,15 @@
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
+
+IMPORTED_CODE_SVG_SYSTEM_PROMPT = """
+You are an expert at building SVGs.
+
+- Do not add comments in the code such as "<!-- Add other navigation links as needed -->" and "<!-- ... other news items ... -->" in place of writing the full code. WRITE THE FULL CODE.
+- Repeat elements as needed to match the screenshot. For example, if there are 15 items, the code should have 15 items. DO NOT LEAVE comments like "<!-- Repeat for each news item -->" or bad things will happen.
+- For images, use placeholder images from https://placehold.co and include a detailed description of the image in the alt text so that an image generation AI can generate the image later.
+- You can use Google Fonts
+
+Return only the full code in <svg></svg> tags.
+Do not include markdown "```" or "```svg" at the start or end.
+"""
diff --git a/backend/prompts.py b/backend/prompts.py
@@ -7,19 +7,25 @@
     IMPORTED_CODE_IONIC_TAILWIND_SYSTEM_PROMPT,
     IMPORTED_CODE_REACT_TAILWIND_SYSTEM_PROMPT,
     IMPORTED_CODE_TAILWIND_SYSTEM_PROMPT,
+    IMPORTED_CODE_SVG_SYSTEM_PROMPT,
 )
 from screenshot_system_prompts import (
     BOOTSTRAP_SYSTEM_PROMPT,
     IONIC_TAILWIND_SYSTEM_PROMPT,
     REACT_TAILWIND_SYSTEM_PROMPT,
     TAILWIND_SYSTEM_PROMPT,
+    SVG_SYSTEM_PROMPT,
 )
 
 
 USER_PROMPT = """
 Generate code for a web page that looks exactly like this.
 """
 
+SVG_USER_PROMPT = """
+Generate code for a SVG that looks exactly like this.
+"""
+
 
 def assemble_imported_code_prompt(
     code: str, stack: str, result_image_data_url: Union[str, None] = None
@@ -33,17 +39,24 @@ def assemble_imported_code_prompt(
         system_content = IMPORTED_CODE_BOOTSTRAP_SYSTEM_PROMPT
     elif stack == "ionic_tailwind":
         system_content = IMPORTED_CODE_IONIC_TAILWIND_SYSTEM_PROMPT
+    elif stack == "svg":
+        system_content = IMPORTED_CODE_SVG_SYSTEM_PROMPT
     else:
         raise Exception("Code config is not one of available options")
 
+    user_content = (
+        "Here is the code of the app: " + code
+        if stack != "svg"
+        else "Here is the code of the SVG: " + code
+    )
     return [
         {
             "role": "system",
             "content": system_content,
         },
         {
             "role": "user",
-            "content": "Here is the code of the app: " + code,
+            "content": user_content,
         },
     ]
     # TODO: Use result_image_data_url
@@ -64,17 +77,21 @@ def assemble_prompt(
         system_content = BOOTSTRAP_SYSTEM_PROMPT
     elif generated_code_config == "ionic_tailwind":
         system_content = IONIC_TAILWIND_SYSTEM_PROMPT
+    elif generated_code_config == "svg":
+        system_content = SVG_SYSTEM_PROMPT
     else:
         raise Exception("Code config is not one of available options")
 
+    user_prompt = USER_PROMPT if generated_code_config != "svg" else SVG_USER_PROMPT
+
     user_content: List[ChatCompletionContentPartParam] = [
         {
             "type": "image_url",
             "image_url": {"url": image_data_url, "detail": "high"},
         },
         {
             "type": "text",
-            "text": USER_PROMPT,
+            "text": user_prompt,
         },
     ]
 

diff --git a/backend/screenshot_system_prompts.py b/backend/screenshot_system_prompts.py
@@ -110,3 +110,21 @@
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
+
+
+SVG_SYSTEM_PROMPT = """
+You are an expert at building SVGs.
+You take screenshots of a reference web page from the user, and then build a SVG that looks exactly like the screenshot.
+
+- Make sure the SVG looks exactly like the screenshot.
+- Pay close attention to background color, text color, font size, font family, 
+padding, margin, border, etc. Match the colors and sizes exactly.
+- Use the exact text from the screenshot.
+- Do not add comments in the code such as "<!-- Add other navigation links as needed -->" and "<!-- ... other news items ... -->" in place of writing the full code. WRITE THE FULL CODE.
+- Repeat elements as needed to match the screenshot. For example, if there are 15 items, the code should have 15 items. DO NOT LEAVE comments like "<!-- Repeat for each news item -->" or bad things will happen.
+- For images, use placeholder images from https://placehold.co and include a detailed description of the image in the alt text so that an image generation AI can generate the image later.
+- You can use Google Fonts
+
+Return only the full code in <svg></svg> tags.
+Do not include markdown "```" or "```svg" at the start or end.
+"""
diff --git a/backend/test_prompts.py b/backend/test_prompts.py
@@ -113,6 +113,23 @@
 Do not include markdown "```" or "```html" at the start or end.
 """
 
+SVG_SYSTEM_PROMPT = """
+You are an expert at building SVGs.
+You take screenshots of a reference web page from the user, and then build a SVG that looks exactly like the screenshot.
+
+- Make sure the SVG looks exactly like the screenshot.
+- Pay close attention to background color, text color, font size, font family, 
+padding, margin, border, etc. Match the colors and sizes exactly.
+- Use the exact text from the screenshot.
+- Do not add comments in the code such as "<!-- Add other navigation links as needed -->" and "<!-- ... other news items ... -->" in place of writing the full code. WRITE THE FULL CODE.
+- Repeat elements as needed to match the screenshot. For example, if there are 15 items, the code should have 15 items. DO NOT LEAVE comments like "<!-- Repeat for each news item -->" or bad things will happen.
+- For images, use placeholder images from https://placehold.co and include a detailed description of the image in the alt text so that an image generation AI can generate the image later.
+- You can use Google Fonts
+
+Return only the full code in <svg></svg> tags.
+Do not include markdown "```" or "```svg" at the start or end.
+"""
+
 IMPORTED_CODE_TAILWIND_SYSTEM_PROMPT = """
 You are an expert Tailwind developer.
 
@@ -194,27 +211,55 @@
 Do not include markdown "```" or "```html" at the start or end.
 """
 
+IMPORTED_CODE_SVG_SYSTEM_PROMPT = """
+You are an expert at building SVGs.
+
+- Do not add comments in the code such as "<!-- Add other navigation links as needed -->" and "<!-- ... other news items ... -->" in place of writing the full code. WRITE THE FULL CODE.
+- Repeat elements as needed to match the screenshot. For example, if there are 15 items, the code should have 15 items. DO NOT LEAVE comments like "<!-- Repeat for each news item -->" or bad things will happen.
+- For images, use placeholder images from https://placehold.co and include a detailed description of the image in the alt text so that an image generation AI can generate the image later.
+- You can use Google Fonts
+
+Return only the full code in <svg></svg> tags.
+Do not include markdown "```" or "```svg" at the start or end.
+"""
+
+USER_PROMPT = """
+Generate code for a web page that looks exactly like this.
+"""
+
+SVG_USER_PROMPT = """
+Generate code for a SVG that looks exactly like this.
+"""
+
 
 def test_prompts():
     tailwind_prompt = assemble_prompt(
         "image_data_url", "html_tailwind", "result_image_data_url"
     )
     assert tailwind_prompt[0]["content"] == TAILWIND_SYSTEM_PROMPT
+    assert tailwind_prompt[1]["content"][2]["text"] == USER_PROMPT  # type: ignore
 
     react_tailwind_prompt = assemble_prompt(
         "image_data_url", "react_tailwind", "result_image_data_url"
     )
     assert react_tailwind_prompt[0]["content"] == REACT_TAILWIND_SYSTEM_PROMPT
+    assert react_tailwind_prompt[1]["content"][2]["text"] == USER_PROMPT  # type: ignore
 
     bootstrap_prompt = assemble_prompt(
         "image_data_url", "bootstrap", "result_image_data_url"
     )
     assert bootstrap_prompt[0]["content"] == BOOTSTRAP_SYSTEM_PROMPT
+    assert bootstrap_prompt[1]["content"][2]["text"] == USER_PROMPT  # type: ignore
 
     ionic_tailwind = assemble_prompt(
         "image_data_url", "ionic_tailwind", "result_image_data_url"
     )
     assert ionic_tailwind[0]["content"] == IONIC_TAILWIND_SYSTEM_PROMPT
+    assert ionic_tailwind[1]["content"][2]["text"] == USER_PROMPT  # type: ignore
+
+    svg_prompt = assemble_prompt("image_data_url", "svg", "result_image_data_url")
+    assert svg_prompt[0]["content"] == SVG_SYSTEM_PROMPT
+    assert svg_prompt[1]["content"][2]["text"] == SVG_USER_PROMPT  # type: ignore
 
 
 def test_imported_code_prompts():
@@ -253,3 +298,10 @@ def test_imported_code_prompts():
         {"role": "user", "content": "Here is the code of the app: code"},
     ]
     assert ionic_tailwind == expected_ionic_tailwind
+
+    svg = assemble_imported_code_prompt("code", "svg", "result_image_data_url")
+    expected_svg = [
+        {"role": "system", "content": IMPORTED_CODE_SVG_SYSTEM_PROMPT},
+        {"role": "user", "content": "Here is the code of the SVG: code"},
+    ]
+    assert svg == expected_svg
diff --git a/frontend/src/components/OutputSettingsSection.tsx b/frontend/src/components/OutputSettingsSection.tsx
@@ -6,6 +6,7 @@ import {
   SelectTrigger,
 } from "./ui/select";
 import { GeneratedCodeConfig } from "../types";
+import { Badge } from "./ui/badge";
 
 function generateDisplayComponent(config: GeneratedCodeConfig) {
   switch (config) {
@@ -36,9 +37,16 @@ function generateDisplayComponent(config: GeneratedCodeConfig) {
           <span className="font-semibold">Tailwind</span>
         </div>
       );
-    default:
-      // TODO: Should never reach this out. Error out
-      return config;
+    case GeneratedCodeConfig.SVG:
+      return (
+        <div>
+          <span className="font-semibold">SVG</span>
+        </div>
+      );
+    default: {
+      const exhaustiveCheck: never = config;
+      throw new Error(`Unhandled case: ${exhaustiveCheck}`);
+    }
   }
 }
 
@@ -83,7 +91,20 @@ function OutputSettingsSection({
                 {generateDisplayComponent(GeneratedCodeConfig.BOOTSTRAP)}
               </SelectItem>
               <SelectItem value={GeneratedCodeConfig.IONIC_TAILWIND}>
-                {generateDisplayComponent(GeneratedCodeConfig.IONIC_TAILWIND)}
+                <div className="flex items-center">
+                  {generateDisplayComponent(GeneratedCodeConfig.IONIC_TAILWIND)}
+                  <Badge className="ml-2" variant="secondary">
+                    Beta
+                  </Badge>
+                </div>
+              </SelectItem>
+              <SelectItem value={GeneratedCodeConfig.SVG}>
+                <div className="flex items-center">
+                  {generateDisplayComponent(GeneratedCodeConfig.SVG)}
+                  <Badge className="ml-2" variant="secondary">
+                    Beta
+                  </Badge>
+                </div>
               </SelectItem>
             </SelectGroup>
           </SelectContent>

diff --git a/frontend/src/components/Preview.tsx b/frontend/src/components/Preview.tsx
@@ -1,6 +1,6 @@
 import { useEffect, useRef } from "react";
 import classNames from "classnames";
-import useThrottle from "../hooks/useThrottle";
+// import useThrottle from "../hooks/useThrottle";
 
 interface Props {
   code: string;

diff --git a/frontend/src/types.ts b/frontend/src/types.ts
@@ -9,6 +9,7 @@ export enum GeneratedCodeConfig {
   REACT_TAILWIND = "react_tailwind",
   BOOTSTRAP = "bootstrap",
   IONIC_TAILWIND = "ionic_tailwind",
+  SVG = "svg",
 }
 
 export interface Settings {