Merge pull request #177 from abi/import-from-code

Allow starting a new project from existing code
2023-12-10 15:58:03 -05:00 · 2023-12-10 15:58:03 -05:00 · bc64da750d
commit bc64da750d
parent 68a8d2788d dd65a6e1e4
24 changed files with 928 additions and 412 deletions
--- a/.vscode/settings.json
+++ b/.vscode/settings.json
@ -0,0 +1,3 @@
 {
  "python.analysis.typeCheckingMode": "strict"
 }
--- a/backend/README.md
+++ b/backend/README.md
@ -1,3 +1,3 @@
-Run tests
+# Run tests
-pytest test_prompts.py
+poetry run pytest
--- a/backend/config.py
+++ b/backend/config.py
@ -0,0 +1,11 @@
 # Useful for debugging purposes when you don't want to waste GPT4-Vision credits
 # Setting to True will stream a mock response instead of calling the OpenAI API
 # TODO: Should only be set to true when value is 'True', not any abitrary truthy value
 import os
 SHOULD_MOCK_AI_RESPONSE = bool(os.environ.get("MOCK", False))
 # Set to True when running in production (on the hosted version)
 # Used as a feature flag to enable or disable certain features
 IS_PROD = os.environ.get("IS_PROD", False)
--- a/backend/image_generation.py
+++ b/backend/image_generation.py
@ -1,15 +1,15 @@
 import asyncio
 import os
 import re
 from typing import Dict, List, Union
 from openai import AsyncOpenAI
 from bs4 import BeautifulSoup
-async def process_tasks(prompts, api_key, base_url):
+async def process_tasks(prompts: List[str], api_key: str, base_url: str):
    tasks = [generate_image(prompt, api_key, base_url) for prompt in prompts]
    results = await asyncio.gather(*tasks, return_exceptions=True)
-    processed_results = []
+    processed_results: List[Union[str, None]] = []
    for result in results:
        if isinstance(result, Exception):
            print(f"An exception occurred: {result}")
@ -20,9 +20,9 @@ async def process_tasks(prompts, api_key, base_url):
    return processed_results
-async def generate_image(prompt, api_key, base_url):
+async def generate_image(prompt: str, api_key: str, base_url: str):
    client = AsyncOpenAI(api_key=api_key, base_url=base_url)
-    image_params = {
+    image_params: Dict[str, Union[str, int]] = {
        "model": "dall-e-3",
        "quality": "standard",
        "style": "natural",
@ -35,7 +35,7 @@ async def generate_image(prompt, api_key, base_url):
    return res.data[0].url
-def extract_dimensions(url):
+def extract_dimensions(url: str):
    # Regular expression to match numbers in the format '300x200'
    matches = re.findall(r"(\d+)x(\d+)", url)
@ -48,11 +48,11 @@ def extract_dimensions(url):
        return (100, 100)
-def create_alt_url_mapping(code):
+def create_alt_url_mapping(code: str) -> Dict[str, str]:
    soup = BeautifulSoup(code, "html.parser")
    images = soup.find_all("img")
-    mapping = {}
+    mapping: Dict[str, str] = {}
    for image in images:
        if not image["src"].startswith("https://placehold.co"):
@ -61,7 +61,9 @@ def create_alt_url_mapping(code):
    return mapping
-async def generate_images(code, api_key, base_url, image_cache):
+async def generate_images(
    code: str, api_key: str, base_url: Union[str, None], image_cache: Dict[str, str]
 ):
    # Find all images
    soup = BeautifulSoup(code, "html.parser")
    images = soup.find_all("img")
--- a/backend/imported_code_prompts.py
+++ b/backend/imported_code_prompts.py
@ -0,0 +1,80 @@
 IMPORTED_CODE_TAILWIND_SYSTEM_PROMPT = """
 You are an expert Tailwind developer.
 - 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. 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.
 In terms of libraries,
 - Use this script to include Tailwind: <script src="https://cdn.tailwindcss.com"></script>
 - You can use Google Fonts
 - Font Awesome for icons: <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css"></link>
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 IMPORTED_CODE_REACT_TAILWIND_SYSTEM_PROMPT = """
 You are an expert React/Tailwind developer
 - 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. 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.
 In terms of libraries,
 - Use these script to include React so that it can run on a standalone page:
    <script src="https://unpkg.com/react/umd/react.development.js"></script>
    <script src="https://unpkg.com/react-dom/umd/react-dom.development.js"></script>
    <script src="https://unpkg.com/@babel/standalone/babel.js"></script>
 - Use this script to include Tailwind: <script src="https://cdn.tailwindcss.com"></script>
 - You can use Google Fonts
 - Font Awesome for icons: <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css"></link>
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 IMPORTED_CODE_BOOTSTRAP_SYSTEM_PROMPT = """
 You are an expert Bootstrap developer.
 - 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. 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.
 In terms of libraries,
 - Use this script to include Bootstrap: <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.2/dist/css/bootstrap.min.css" rel="stylesheet" integrity="sha384-T3c6CoIi6uLrA9TneNEoa7RxnatzjcDSCmG1MXxSR1GAsXEV/Dwwykc2MPK8M2HN" crossorigin="anonymous">
 - You can use Google Fonts
 - Font Awesome for icons: <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css"></link>
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 IMPORTED_CODE_IONIC_TAILWIND_SYSTEM_PROMPT = """
 You are an expert Ionic/Tailwind developer.
 - 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. 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.
 In terms of libraries,
 - Use these script to include Ionic so that it can run on a standalone page:
    <script type="module" src="https://cdn.jsdelivr.net/npm/@ionic/core/dist/ionic/ionic.esm.js"></script>
    <script nomodule src="https://cdn.jsdelivr.net/npm/@ionic/core/dist/ionic/ionic.js"></script>
    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/@ionic/core/css/ionic.bundle.css" />
 - Use this script to include Tailwind: <script src="https://cdn.tailwindcss.com"></script>
 - You can use Google Fonts
 - ionicons for icons, add the following <script > tags near the end of the page, right before the closing </body> tag:
    <script type="module">
        import ionicons from 'https://cdn.jsdelivr.net/npm/ionicons/+esm'
    </script>
    <script nomodule src="https://cdn.jsdelivr.net/npm/ionicons/dist/esm/ionicons.min.js"></script>
    <link href="https://cdn.jsdelivr.net/npm/ionicons/dist/collection/components/icon/icon.min.css" rel="stylesheet">
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
--- a/backend/llm.py
+++ b/backend/llm.py
@ -1,16 +1,16 @@
-import os
+from typing import Awaitable, Callable, List
 from typing import Awaitable, Callable
 from openai import AsyncOpenAI
 from openai.types.chat import ChatCompletionMessageParam, ChatCompletionChunk
 MODEL_GPT_4_VISION = "gpt-4-vision-preview"
 async def stream_openai_response(
-    messages,
+    messages: List[ChatCompletionMessageParam],
    api_key: str,
    base_url: str | None,
    callback: Callable[[str], Awaitable[None]],
-):
+) -> str:
    client = AsyncOpenAI(api_key=api_key, base_url=base_url)
    model = MODEL_GPT_4_VISION
@ -23,9 +23,10 @@ async def stream_openai_response(
        params["max_tokens"] = 4096
        params["temperature"] = 0
-    completion = await client.chat.completions.create(**params)
+    stream = await client.chat.completions.create(**params)  # type: ignore
    full_response = ""
-    async for chunk in completion:
+    async for chunk in stream:  # type: ignore
        assert isinstance(chunk, ChatCompletionChunk)
        content = chunk.choices[0].delta.content or ""
        full_response += content
        await callback(content)
--- a/backend/main.py
+++ b/backend/main.py
@ -1,25 +1,12 @@
 # Load environment variables first
 from dotenv import load_dotenv
 load_dotenv()
-import json
+from fastapi import FastAPI
 import os
 import traceback
 from datetime import datetime
 from fastapi import FastAPI, WebSocket
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import HTMLResponse
+from routes import screenshot, generate_code, home
 import openai
 from llm import stream_openai_response
 from mock import mock_completion
 from utils import pprint_prompt
 from image_generation import create_alt_url_mapping, generate_images
 from prompts import assemble_prompt
 from routes import screenshot
 from access_token import validate_access_token
 app = FastAPI(openapi_url=None, docs_url=None, redoc_url=None)
@ -32,231 +19,7 @@ app.add_middleware(
    allow_headers=["*"],
 )
-
+# Add routes
-# Useful for debugging purposes when you don't want to waste GPT4-Vision credits
+app.include_router(generate_code.router)
 # Setting to True will stream a mock response instead of calling the OpenAI API
 # TODO: Should only be set to true when value is 'True', not any abitrary truthy value
 SHOULD_MOCK_AI_RESPONSE = bool(os.environ.get("MOCK", False))
 # Set to True when running in production (on the hosted version)
 # Used as a feature flag to enable or disable certain features
 IS_PROD = os.environ.get("IS_PROD", False)
 app.include_router(screenshot.router)
-
+app.include_router(home.router)
@app.get("/")
 async def get_status():
    return HTMLResponse(
        content="<h3>Your backend is running correctly. Please open the front-end URL (default is http://localhost:5173) to use screenshot-to-code.</h3>"
    )
 def write_logs(prompt_messages, completion):
    # Get the logs path from environment, default to the current working directory
    logs_path = os.environ.get("LOGS_PATH", os.getcwd())
    # Create run_logs directory if it doesn't exist within the specified logs path
    logs_directory = os.path.join(logs_path, "run_logs")
    if not os.path.exists(logs_directory):
        os.makedirs(logs_directory)
    print("Writing to logs directory:", logs_directory)
    # Generate a unique filename using the current timestamp within the logs directory
    filename = datetime.now().strftime(f"{logs_directory}/messages_%Y%m%d_%H%M%S.json")
    # Write the messages dict into a new file for each run
    with open(filename, "w") as f:
        f.write(json.dumps({"prompt": prompt_messages, "completion": completion}))
@app.websocket("/generate-code")
 async def stream_code(websocket: WebSocket):
    await websocket.accept()
    print("Incoming websocket connection...")
    async def throw_error(
        message: str,
    ):
        await websocket.send_json({"type": "error", "value": message})
        await websocket.close()
    params = await websocket.receive_json()
    print("Received params")
    # Read the code config settings from the request. Fall back to default if not provided.
    generated_code_config = ""
    if "generatedCodeConfig" in params and params["generatedCodeConfig"]:
        generated_code_config = params["generatedCodeConfig"]
    print(f"Generating {generated_code_config} code")
    # Get the OpenAI API key from the request. Fall back to environment variable if not provided.
    # If neither is provided, we throw an error.
    openai_api_key = None
    if "accessCode" in params and params["accessCode"]:
        print("Access code - using platform API key")
        res = await validate_access_token(params["accessCode"])
        if res["success"]:
            openai_api_key = os.environ.get("PLATFORM_OPENAI_API_KEY")
        else:
            await websocket.send_json(
                {
                    "type": "error",
                    "value": res["failure_reason"],
                }
            )
            return
    else:
        if params["openAiApiKey"]:
            openai_api_key = params["openAiApiKey"]
            print("Using OpenAI API key from client-side settings dialog")
        else:
            openai_api_key = os.environ.get("OPENAI_API_KEY")
            if openai_api_key:
                print("Using OpenAI API key from environment variable")
    if not openai_api_key:
        print("OpenAI API key not found")
        await websocket.send_json(
            {
                "type": "error",
                "value": "No OpenAI API key found. Please add your API key in the settings dialog or add it to backend/.env file.",
            }
        )
        return
    # Get the OpenAI Base URL from the request. Fall back to environment variable if not provided.
    openai_base_url = None
    # Disable user-specified OpenAI Base URL in prod
    if not os.environ.get("IS_PROD"):
        if "openAiBaseURL" in params and params["openAiBaseURL"]:
            openai_base_url = params["openAiBaseURL"]
            print("Using OpenAI Base URL from client-side settings dialog")
        else:
            openai_base_url = os.environ.get("OPENAI_BASE_URL")
            if openai_base_url:
                print("Using OpenAI Base URL from environment variable")
    if not openai_base_url:
        print("Using official OpenAI URL")
    # Get the image generation flag from the request. Fall back to True if not provided.
    should_generate_images = (
        params["isImageGenerationEnabled"]
        if "isImageGenerationEnabled" in params
        else True
    )
    print("generating code...")
    await websocket.send_json({"type": "status", "value": "Generating code..."})
    async def process_chunk(content):
        await websocket.send_json({"type": "chunk", "value": content})
    # Assemble the prompt
    try:
        if params.get("resultImage") and params["resultImage"]:
            prompt_messages = assemble_prompt(
                params["image"], generated_code_config, params["resultImage"]
            )
        else:
            prompt_messages = assemble_prompt(params["image"], generated_code_config)
    except:
        await websocket.send_json(
            {
                "type": "error",
                "value": "Error assembling prompt. Contact support at support@picoapps.xyz",
            }
        )
        await websocket.close()
        return
    # 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
        for index, text in enumerate(params["history"]):
            prompt_messages += [
                {"role": "assistant" if index % 2 == 0 else "user", "content": text}
            ]
        image_cache = create_alt_url_mapping(params["history"][-2])
    if SHOULD_MOCK_AI_RESPONSE:
        completion = await mock_completion(process_chunk)
    else:
        try:
            completion = await stream_openai_response(
                prompt_messages,
                api_key=openai_api_key,
                base_url=openai_base_url,
                callback=lambda x: process_chunk(x),
            )
        except openai.AuthenticationError as e:
            print("[GENERATE_CODE] Authentication failed", e)
            error_message = (
                "Incorrect OpenAI key. Please make sure your OpenAI API key is correct, or create a new OpenAI API key on your OpenAI dashboard."
                + (
                    " Alternatively, you can purchase code generation credits directly on this website."
                    if IS_PROD
                    else ""
                )
            )
            return await throw_error(error_message)
        except openai.NotFoundError as e:
            print("[GENERATE_CODE] Model not found", e)
            error_message = (
                e.message
                + ". Please make sure you have followed the instructions correctly to obtain an OpenAI key with GPT vision access: https://github.com/abi/screenshot-to-code/blob/main/Troubleshooting.md"
                + (
                    " Alternatively, you can purchase code generation credits directly on this website."
                    if IS_PROD
                    else ""
                )
            )
            return await throw_error(error_message)
        except openai.RateLimitError as e:
            print("[GENERATE_CODE] Rate limit exceeded", e)
            error_message = (
                "OpenAI error - 'You exceeded your current quota, please check your plan and billing details.'"
                + (
                    " Alternatively, you can purchase code generation credits directly on this website."
                    if IS_PROD
                    else ""
                )
            )
            return await throw_error(error_message)
    # Write the messages dict into a log so that we can debug later
    write_logs(prompt_messages, completion)
    try:
        if should_generate_images:
            await websocket.send_json(
                {"type": "status", "value": "Generating images..."}
            )
            updated_html = await generate_images(
                completion,
                api_key=openai_api_key,
                base_url=openai_base_url,
                image_cache=image_cache,
            )
        else:
            updated_html = completion
        await websocket.send_json({"type": "setCode", "value": updated_html})
        await websocket.send_json(
            {"type": "status", "value": "Code generation complete."}
        )
    except Exception as e:
        traceback.print_exc()
        print("Image generation failed", e)
        await websocket.send_json(
            {"type": "status", "value": "Image generation failed but code is complete."}
        )
    await websocket.close()
--- a/backend/mock_llm.py
+++ b/backend/mock_llm.py
@ -1,7 +1,8 @@
 import asyncio
 from typing import Awaitable, Callable
-async def mock_completion(process_chunk):
+async def mock_completion(process_chunk: Callable[[str], Awaitable[None]]) -> str:
    code_to_return = NO_IMAGES_NYTIMES_MOCK_CODE
    for i in range(0, len(code_to_return), 10):
--- a/backend/poetry.lock
+++ b/backend/poetry.lock
@ -200,6 +200,18 @@ files = [
    {file = "idna-3.6.tar.gz", hash = "sha256:9ecdbbd083b06798ae1e86adcbfe8ab1479cf864e4ee30fe4e46a003d12491ca"},
 ]
 [[package]]
 name = "iniconfig"
 version = "2.0.0"
 description = "brain-dead simple config-ini parsing"
 category = "dev"
 optional = false
 python-versions = ">=3.7"
 files = [
    {file = "iniconfig-2.0.0-py3-none-any.whl", hash = "sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374"},
    {file = "iniconfig-2.0.0.tar.gz", hash = "sha256:2d91e135bf72d31a410b17c16da610a82cb55f6b0477d1a902134b24a455b8b3"},
 ]
 [[package]]
 name = "openai"
 version = "1.3.7"
@ -224,6 +236,34 @@ typing-extensions = ">=4.5,<5"
 [package.extras]
 datalib = ["numpy (>=1)", "pandas (>=1.2.3)", "pandas-stubs (>=1.1.0.11)"]
 [[package]]
 name = "packaging"
 version = "23.2"
 description = "Core utilities for Python packages"
 category = "dev"
 optional = false
 python-versions = ">=3.7"
 files = [
    {file = "packaging-23.2-py3-none-any.whl", hash = "sha256:8c491190033a9af7e1d931d0b5dacc2ef47509b34dd0de67ed209b5203fc88c7"},
    {file = "packaging-23.2.tar.gz", hash = "sha256:048fb0e9405036518eaaf48a55953c750c11e1a1b68e0dd1a9d62ed0c092cfc5"},
 ]
 [[package]]
 name = "pluggy"
 version = "1.3.0"
 description = "plugin and hook calling mechanisms for python"
 category = "dev"
 optional = false
 python-versions = ">=3.8"
 files = [
    {file = "pluggy-1.3.0-py3-none-any.whl", hash = "sha256:d89c696a773f8bd377d18e5ecda92b7a3793cbe66c87060a6fb58c7b6e1061f7"},
    {file = "pluggy-1.3.0.tar.gz", hash = "sha256:cf61ae8f126ac6f7c451172cf30e3e43d3ca77615509771b3a984a0730651e12"},
 ]
 [package.extras]
 dev = ["pre-commit", "tox"]
 testing = ["pytest", "pytest-benchmark"]
 [[package]]
 name = "pydantic"
 version = "1.10.13"
@ -277,6 +317,29 @@ typing-extensions = ">=4.2.0"
 dotenv = ["python-dotenv (>=0.10.4)"]
 email = ["email-validator (>=1.0.3)"]
 [[package]]
 name = "pytest"
 version = "7.4.3"
 description = "pytest: simple powerful testing with Python"
 category = "dev"
 optional = false
 python-versions = ">=3.7"
 files = [
    {file = "pytest-7.4.3-py3-none-any.whl", hash = "sha256:0d009c083ea859a71b76adf7c1d502e4bc170b80a8ef002da5806527b9591fac"},
    {file = "pytest-7.4.3.tar.gz", hash = "sha256:d989d136982de4e3b29dabcc838ad581c64e8ed52c11fbe86ddebd9da0818cd5"},
 ]
 [package.dependencies]
 colorama = {version = "*", markers = "sys_platform == \"win32\""}
 exceptiongroup = {version = ">=1.0.0rc8", markers = "python_version < \"3.11\""}
 iniconfig = "*"
 packaging = "*"
 pluggy = ">=0.12,<2.0"
 tomli = {version = ">=1.0.0", markers = "python_version < \"3.11\""}
 [package.extras]
 testing = ["argcomplete", "attrs (>=19.2.0)", "hypothesis (>=3.56)", "mock", "nose", "pygments (>=2.7.2)", "requests", "setuptools", "xmlschema"]
 [[package]]
 name = "python-dotenv"
 version = "1.0.0"
@ -334,6 +397,18 @@ anyio = ">=3.4.0,<5"
 [package.extras]
 full = ["httpx (>=0.22.0)", "itsdangerous", "jinja2", "python-multipart", "pyyaml"]
 [[package]]
 name = "tomli"
 version = "2.0.1"
 description = "A lil' TOML parser"
 category = "dev"
 optional = false
 python-versions = ">=3.7"
 files = [
    {file = "tomli-2.0.1-py3-none-any.whl", hash = "sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc"},
    {file = "tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f"},
 ]
 [[package]]
 name = "tqdm"
 version = "4.66.1"
@ -472,4 +547,4 @@ files = [
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.10"
-content-hash = "b8d248a44a5eea9638a7726096de77d7a9aa8c00673da806534da2c228ffabb4"
+content-hash = "c31ed2a1ce006749d6f34d8d6aebcbc58d306b9f8925b40cc35972a74979e5c7"
--- a/backend/prompts.py
+++ b/backend/prompts.py
@ -1,124 +1,59 @@
-TAILWIND_SYSTEM_PROMPT = """
+from typing import List, Union
 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(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.
+from openai.types.chat import ChatCompletionMessageParam, ChatCompletionContentPartParam
 - 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.
-In terms of libraries,
+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,
 )
 - Use this script to include Tailwind: <script src="https://cdn.tailwindcss.com"></script>
 - You can use Google Fonts
 - Font Awesome for icons: <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css"></link>
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 BOOTSTRAP_SYSTEM_PROMPT = """
 You are an expert Bootstrap developer
 You take screenshots of a reference web page from the user, and then build single page apps 
 using Bootstrap, HTML and JS.
 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, 
 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.
 In terms of libraries,
 - Use this script to include Bootstrap: <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.2/dist/css/bootstrap.min.css" rel="stylesheet" integrity="sha384-T3c6CoIi6uLrA9TneNEoa7RxnatzjcDSCmG1MXxSR1GAsXEV/Dwwykc2MPK8M2HN" crossorigin="anonymous">
 - You can use Google Fonts
 - Font Awesome for icons: <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css"></link>
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 REACT_TAILWIND_SYSTEM_PROMPT = """
 You are an expert React/Tailwind developer
 You take screenshots of a reference web page from the user, and then build single page apps 
 using React and Tailwind CSS.
 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, 
 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.
 In terms of libraries,
 - Use these script to include React so that it can run on a standalone page:
    <script src="https://unpkg.com/react/umd/react.development.js"></script>
    <script src="https://unpkg.com/react-dom/umd/react-dom.development.js"></script>
    <script src="https://unpkg.com/@babel/standalone/babel.js"></script>
 - Use this script to include Tailwind: <script src="https://cdn.tailwindcss.com"></script>
 - You can use Google Fonts
 - Font Awesome for icons: <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css"></link>
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 IONIC_TAILWIND_SYSTEM_PROMPT = """
 You are an expert Ionic/Tailwind developer
 You take screenshots of a reference web page from the user, and then build single page apps 
 using Ionic and Tailwind CSS.
 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, 
 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.
 In terms of libraries,
 - Use these script to include Ionic so that it can run on a standalone page:
    <script type="module" src="https://cdn.jsdelivr.net/npm/@ionic/core/dist/ionic/ionic.esm.js"></script>
    <script nomodule src="https://cdn.jsdelivr.net/npm/@ionic/core/dist/ionic/ionic.js"></script>
    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/@ionic/core/css/ionic.bundle.css" />
 - Use this script to include Tailwind: <script src="https://cdn.tailwindcss.com"></script>
 - You can use Google Fonts
 - ionicons for icons, add the following <script > tags near the end of the page, right before the closing </body> tag:
    <script type="module">
        import ionicons from 'https://cdn.jsdelivr.net/npm/ionicons/+esm'
    </script>
    <script nomodule src="https://cdn.jsdelivr.net/npm/ionicons/dist/esm/ionicons.min.js"></script>
    <link href="https://cdn.jsdelivr.net/npm/ionicons/dist/collection/components/icon/icon.min.css" rel="stylesheet">
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 USER_PROMPT = """
 Generate code for a web page that looks exactly like this.
 """
 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(
-    image_data_url, generated_code_config: str, result_image_data_url=None
+    image_data_url: 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":
@ -132,7 +67,7 @@ def assemble_prompt(
    else:
        raise Exception("Code config is not one of available options")
-    user_content = [
+    user_content: List[ChatCompletionContentPartParam] = [
        {
            "type": "image_url",
            "image_url": {"url": image_data_url, "detail": "high"},
--- a/backend/pyproject.toml
+++ b/backend/pyproject.toml
@ -15,6 +15,9 @@ python-dotenv = "^1.0.0"
 beautifulsoup4 = "^4.12.2"
 httpx = "^0.25.1"
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.3"
 [build-system]
 requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
--- a/backend/routes/generate_code.py
+++ b/backend/routes/generate_code.py
@ -0,0 +1,261 @@
 import os
 import traceback
 from fastapi import APIRouter, WebSocket
 import openai
 from config import IS_PROD, SHOULD_MOCK_AI_RESPONSE
 from llm import stream_openai_response
 from openai.types.chat import ChatCompletionMessageParam
 from mock_llm import mock_completion
 from typing import Dict, List
 from image_generation import create_alt_url_mapping, generate_images
 from prompts import assemble_imported_code_prompt, assemble_prompt
 from access_token import validate_access_token
 from datetime import datetime
 import json
 from utils import pprint_prompt  # type: ignore
 router = APIRouter()
 def write_logs(prompt_messages: List[ChatCompletionMessageParam], completion: str):
    # Get the logs path from environment, default to the current working directory
    logs_path = os.environ.get("LOGS_PATH", os.getcwd())
    # Create run_logs directory if it doesn't exist within the specified logs path
    logs_directory = os.path.join(logs_path, "run_logs")
    if not os.path.exists(logs_directory):
        os.makedirs(logs_directory)
    print("Writing to logs directory:", logs_directory)
    # Generate a unique filename using the current timestamp within the logs directory
    filename = datetime.now().strftime(f"{logs_directory}/messages_%Y%m%d_%H%M%S.json")
    # Write the messages dict into a new file for each run
    with open(filename, "w") as f:
        f.write(json.dumps({"prompt": prompt_messages, "completion": completion}))
@router.websocket("/generate-code")
 async def stream_code(websocket: WebSocket):
    await websocket.accept()
    print("Incoming websocket connection...")
    async def throw_error(
        message: str,
    ):
        await websocket.send_json({"type": "error", "value": message})
        await websocket.close()
    # TODO: Are the values always strings?
    params: Dict[str, str] = await websocket.receive_json()
    print("Received params")
    # Read the code config settings from the request. Fall back to default if not provided.
    generated_code_config = ""
    if "generatedCodeConfig" in params and params["generatedCodeConfig"]:
        generated_code_config = params["generatedCodeConfig"]
    print(f"Generating {generated_code_config} code")
    # Get the OpenAI API key from the request. Fall back to environment variable if not provided.
    # If neither is provided, we throw an error.
    openai_api_key = None
    if "accessCode" in params and params["accessCode"]:
        print("Access code - using platform API key")
        res = await validate_access_token(params["accessCode"])
        if res["success"]:
            openai_api_key = os.environ.get("PLATFORM_OPENAI_API_KEY")
        else:
            await websocket.send_json(
                {
                    "type": "error",
                    "value": res["failure_reason"],
                }
            )
            return
    else:
        if params["openAiApiKey"]:
            openai_api_key = params["openAiApiKey"]
            print("Using OpenAI API key from client-side settings dialog")
        else:
            openai_api_key = os.environ.get("OPENAI_API_KEY")
            if openai_api_key:
                print("Using OpenAI API key from environment variable")
    if not openai_api_key:
        print("OpenAI API key not found")
        await websocket.send_json(
            {
                "type": "error",
                "value": "No OpenAI API key found. Please add your API key in the settings dialog or add it to backend/.env file.",
            }
        )
        return
    # Get the OpenAI Base URL from the request. Fall back to environment variable if not provided.
    openai_base_url = None
    # Disable user-specified OpenAI Base URL in prod
    if not os.environ.get("IS_PROD"):
        if "openAiBaseURL" in params and params["openAiBaseURL"]:
            openai_base_url = params["openAiBaseURL"]
            print("Using OpenAI Base URL from client-side settings dialog")
        else:
            openai_base_url = os.environ.get("OPENAI_BASE_URL")
            if openai_base_url:
                print("Using OpenAI Base URL from environment variable")
    if not openai_base_url:
        print("Using official OpenAI URL")
    # Get the image generation flag from the request. Fall back to True if not provided.
    should_generate_images = (
        params["isImageGenerationEnabled"]
        if "isImageGenerationEnabled" in params
        else True
    )
    print("generating code...")
    await websocket.send_json({"type": "status", "value": "Generating code..."})
    async def process_chunk(content: str):
        await websocket.send_json({"type": "chunk", "value": content})
    # Image cache for updates so that we don't have to regenerate images
    image_cache: Dict[str, str] = {}
    # If this generation started off with imported code, we need to assemble the prompt differently
    if params.get("isImportedFromCode") and params["isImportedFromCode"]:
        original_imported_code = params["history"][0]
        prompt_messages = assemble_imported_code_prompt(
            original_imported_code, generated_code_config
        )
        for index, text in enumerate(params["history"][1:]):
            if index % 2 == 0:
                message: ChatCompletionMessageParam = {
                    "role": "user",
                    "content": text,
                }
            else:
                message: ChatCompletionMessageParam = {
                    "role": "assistant",
                    "content": text,
                }
            prompt_messages.append(message)
    else:
        # Assemble the prompt
        try:
            if params.get("resultImage") and params["resultImage"]:
                prompt_messages = assemble_prompt(
                    params["image"], generated_code_config, params["resultImage"]
                )
            else:
                prompt_messages = assemble_prompt(
                    params["image"], generated_code_config
                )
        except:
            await websocket.send_json(
                {
                    "type": "error",
                    "value": "Error assembling prompt. Contact support at support@picoapps.xyz",
                }
            )
            await websocket.close()
            return
        if params["generationType"] == "update":
            # Transform the history tree into message format
            # TODO: Move this to frontend
            for index, text in enumerate(params["history"]):
                if index % 2 == 0:
                    message: ChatCompletionMessageParam = {
                        "role": "assistant",
                        "content": text,
                    }
                else:
                    message: ChatCompletionMessageParam = {
                        "role": "user",
                        "content": text,
                    }
                prompt_messages.append(message)
            image_cache = create_alt_url_mapping(params["history"][-2])
    pprint_prompt(prompt_messages)
    if SHOULD_MOCK_AI_RESPONSE:
        completion = await mock_completion(process_chunk)
    else:
        try:
            completion = await stream_openai_response(
                prompt_messages,
                api_key=openai_api_key,
                base_url=openai_base_url,
                callback=lambda x: process_chunk(x),
            )
        except openai.AuthenticationError as e:
            print("[GENERATE_CODE] Authentication failed", e)
            error_message = (
                "Incorrect OpenAI key. Please make sure your OpenAI API key is correct, or create a new OpenAI API key on your OpenAI dashboard."
                + (
                    " Alternatively, you can purchase code generation credits directly on this website."
                    if IS_PROD
                    else ""
                )
            )
            return await throw_error(error_message)
        except openai.NotFoundError as e:
            print("[GENERATE_CODE] Model not found", e)
            error_message = (
                e.message
                + ". Please make sure you have followed the instructions correctly to obtain an OpenAI key with GPT vision access: https://github.com/abi/screenshot-to-code/blob/main/Troubleshooting.md"
                + (
                    " Alternatively, you can purchase code generation credits directly on this website."
                    if IS_PROD
                    else ""
                )
            )
            return await throw_error(error_message)
        except openai.RateLimitError as e:
            print("[GENERATE_CODE] Rate limit exceeded", e)
            error_message = (
                "OpenAI error - 'You exceeded your current quota, please check your plan and billing details.'"
                + (
                    " Alternatively, you can purchase code generation credits directly on this website."
                    if IS_PROD
                    else ""
                )
            )
            return await throw_error(error_message)
    # Write the messages dict into a log so that we can debug later
    write_logs(prompt_messages, completion)
    try:
        if should_generate_images:
            await websocket.send_json(
                {"type": "status", "value": "Generating images..."}
            )
            updated_html = await generate_images(
                completion,
                api_key=openai_api_key,
                base_url=openai_base_url,
                image_cache=image_cache,
            )
        else:
            updated_html = completion
        await websocket.send_json({"type": "setCode", "value": updated_html})
        await websocket.send_json(
            {"type": "status", "value": "Code generation complete."}
        )
    except Exception as e:
        traceback.print_exc()
        print("Image generation failed", e)
        await websocket.send_json(
            {"type": "status", "value": "Image generation failed but code is complete."}
        )
    await websocket.close()
--- a/backend/routes/home.py
+++ b/backend/routes/home.py
@ -0,0 +1,12 @@
 from fastapi import APIRouter
 from fastapi.responses import HTMLResponse
 router = APIRouter()
@router.get("/")
 async def get_status():
    return HTMLResponse(
        content="<h3>Your backend is running correctly. Please open the front-end URL (default is http://localhost:5173) to use screenshot-to-code.</h3>"
    )
--- a/backend/routes/screenshot.py
+++ b/backend/routes/screenshot.py
@ -11,7 +11,9 @@ def bytes_to_data_url(image_bytes: bytes, mime_type: str) -> str:
    return f"data:{mime_type};base64,{base64_image}"
-async def capture_screenshot(target_url, api_key, device="desktop") -> bytes:
+async def capture_screenshot(
    target_url: str, api_key: str, device: str = "desktop"
 ) -> bytes:
    api_base_url = "https://api.screenshotone.com/take"
    params = {
--- a/backend/screenshot_system_prompts.py
+++ b/backend/screenshot_system_prompts.py
@ -0,0 +1,112 @@
 TAILWIND_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(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, 
 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.
 In terms of libraries,
 - Use this script to include Tailwind: <script src="https://cdn.tailwindcss.com"></script>
 - You can use Google Fonts
 - Font Awesome for icons: <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css"></link>
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 BOOTSTRAP_SYSTEM_PROMPT = """
 You are an expert Bootstrap developer
 You take screenshots of a reference web page from the user, and then build single page apps 
 using Bootstrap, HTML and JS.
 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, 
 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.
 In terms of libraries,
 - Use this script to include Bootstrap: <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.2/dist/css/bootstrap.min.css" rel="stylesheet" integrity="sha384-T3c6CoIi6uLrA9TneNEoa7RxnatzjcDSCmG1MXxSR1GAsXEV/Dwwykc2MPK8M2HN" crossorigin="anonymous">
 - You can use Google Fonts
 - Font Awesome for icons: <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css"></link>
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 REACT_TAILWIND_SYSTEM_PROMPT = """
 You are an expert React/Tailwind developer
 You take screenshots of a reference web page from the user, and then build single page apps 
 using React and Tailwind CSS.
 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, 
 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.
 In terms of libraries,
 - Use these script to include React so that it can run on a standalone page:
    <script src="https://unpkg.com/react/umd/react.development.js"></script>
    <script src="https://unpkg.com/react-dom/umd/react-dom.development.js"></script>
    <script src="https://unpkg.com/@babel/standalone/babel.js"></script>
 - Use this script to include Tailwind: <script src="https://cdn.tailwindcss.com"></script>
 - You can use Google Fonts
 - Font Awesome for icons: <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css"></link>
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 IONIC_TAILWIND_SYSTEM_PROMPT = """
 You are an expert Ionic/Tailwind developer
 You take screenshots of a reference web page from the user, and then build single page apps 
 using Ionic and Tailwind CSS.
 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, 
 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.
 In terms of libraries,
 - Use these script to include Ionic so that it can run on a standalone page:
    <script type="module" src="https://cdn.jsdelivr.net/npm/@ionic/core/dist/ionic/ionic.esm.js"></script>
    <script nomodule src="https://cdn.jsdelivr.net/npm/@ionic/core/dist/ionic/ionic.js"></script>
    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/@ionic/core/css/ionic.bundle.css" />
 - Use this script to include Tailwind: <script src="https://cdn.tailwindcss.com"></script>
 - You can use Google Fonts
 - ionicons for icons, add the following <script > tags near the end of the page, right before the closing </body> tag:
    <script type="module">
        import ionicons from 'https://cdn.jsdelivr.net/npm/ionicons/+esm'
    </script>
    <script nomodule src="https://cdn.jsdelivr.net/npm/ionicons/dist/esm/ionicons.min.js"></script>
    <link href="https://cdn.jsdelivr.net/npm/ionicons/dist/collection/components/icon/icon.min.css" rel="stylesheet">
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
--- a/backend/test_prompts.py
+++ b/backend/test_prompts.py
@ -1,4 +1,4 @@
-from prompts import assemble_prompt
+from prompts import assemble_imported_code_prompt, assemble_prompt
 TAILWIND_SYSTEM_PROMPT = """
 You are an expert Tailwind developer
@ -113,6 +113,87 @@ Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 IMPORTED_CODE_TAILWIND_SYSTEM_PROMPT = """
 You are an expert Tailwind developer.
 - 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. 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.
 In terms of libraries,
 - Use this script to include Tailwind: <script src="https://cdn.tailwindcss.com"></script>
 - You can use Google Fonts
 - Font Awesome for icons: <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css"></link>
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 IMPORTED_CODE_REACT_TAILWIND_SYSTEM_PROMPT = """
 You are an expert React/Tailwind developer
 - 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. 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.
 In terms of libraries,
 - Use these script to include React so that it can run on a standalone page:
    <script src="https://unpkg.com/react/umd/react.development.js"></script>
    <script src="https://unpkg.com/react-dom/umd/react-dom.development.js"></script>
    <script src="https://unpkg.com/@babel/standalone/babel.js"></script>
 - Use this script to include Tailwind: <script src="https://cdn.tailwindcss.com"></script>
 - You can use Google Fonts
 - Font Awesome for icons: <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css"></link>
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 IMPORTED_CODE_BOOTSTRAP_SYSTEM_PROMPT = """
 You are an expert Bootstrap developer.
 - 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. 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.
 In terms of libraries,
 - Use this script to include Bootstrap: <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.2/dist/css/bootstrap.min.css" rel="stylesheet" integrity="sha384-T3c6CoIi6uLrA9TneNEoa7RxnatzjcDSCmG1MXxSR1GAsXEV/Dwwykc2MPK8M2HN" crossorigin="anonymous">
 - You can use Google Fonts
 - Font Awesome for icons: <link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/5.15.3/css/all.min.css"></link>
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 IMPORTED_CODE_IONIC_TAILWIND_SYSTEM_PROMPT = """
 You are an expert Ionic/Tailwind developer.
 - 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. 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.
 In terms of libraries,
 - Use these script to include Ionic so that it can run on a standalone page:
    <script type="module" src="https://cdn.jsdelivr.net/npm/@ionic/core/dist/ionic/ionic.esm.js"></script>
    <script nomodule src="https://cdn.jsdelivr.net/npm/@ionic/core/dist/ionic/ionic.js"></script>
    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/@ionic/core/css/ionic.bundle.css" />
 - Use this script to include Tailwind: <script src="https://cdn.tailwindcss.com"></script>
 - You can use Google Fonts
 - ionicons for icons, add the following <script > tags near the end of the page, right before the closing </body> tag:
    <script type="module">
        import ionicons from 'https://cdn.jsdelivr.net/npm/ionicons/+esm'
    </script>
    <script nomodule src="https://cdn.jsdelivr.net/npm/ionicons/dist/esm/ionicons.min.js"></script>
    <link href="https://cdn.jsdelivr.net/npm/ionicons/dist/collection/components/icon/icon.min.css" rel="stylesheet">
 Return only the full code in <html></html> tags.
 Do not include markdown "```" or "```html" at the start or end.
 """
 def test_prompts():
    tailwind_prompt = assemble_prompt(
@ -134,3 +215,41 @@ def test_prompts():
        "image_data_url", "ionic_tailwind", "result_image_data_url"
    )
    assert ionic_tailwind[0]["content"] == IONIC_TAILWIND_SYSTEM_PROMPT
 def test_imported_code_prompts():
    tailwind_prompt = assemble_imported_code_prompt(
        "code", "html_tailwind", "result_image_data_url"
    )
    expected_tailwind_prompt = [
        {"role": "system", "content": IMPORTED_CODE_TAILWIND_SYSTEM_PROMPT},
        {"role": "user", "content": "Here is the code of the app: code"},
    ]
    assert tailwind_prompt == expected_tailwind_prompt
    react_tailwind_prompt = assemble_imported_code_prompt(
        "code", "react_tailwind", "result_image_data_url"
    )
    expected_react_tailwind_prompt = [
        {"role": "system", "content": IMPORTED_CODE_REACT_TAILWIND_SYSTEM_PROMPT},
        {"role": "user", "content": "Here is the code of the app: code"},
    ]
    assert react_tailwind_prompt == expected_react_tailwind_prompt
    bootstrap_prompt = assemble_imported_code_prompt(
        "code", "bootstrap", "result_image_data_url"
    )
    expected_bootstrap_prompt = [
        {"role": "system", "content": IMPORTED_CODE_BOOTSTRAP_SYSTEM_PROMPT},
        {"role": "user", "content": "Here is the code of the app: code"},
    ]
    assert bootstrap_prompt == expected_bootstrap_prompt
    ionic_tailwind = assemble_imported_code_prompt(
        "code", "ionic_tailwind", "result_image_data_url"
    )
    expected_ionic_tailwind = [
        {"role": "system", "content": IMPORTED_CODE_IONIC_TAILWIND_SYSTEM_PROMPT},
        {"role": "user", "content": "Here is the code of the app: code"},
    ]
    assert ionic_tailwind == expected_ionic_tailwind
--- a/backend/utils.py
+++ b/backend/utils.py
@ -1,28 +1,30 @@
 import copy
 import json
 from typing import List
 from openai.types.chat import ChatCompletionMessageParam
-def pprint_prompt(prompt_messages):
+def pprint_prompt(prompt_messages: List[ChatCompletionMessageParam]):
    print(json.dumps(truncate_data_strings(prompt_messages), indent=4))
-def truncate_data_strings(data):
+def truncate_data_strings(data: List[ChatCompletionMessageParam]):  # type: ignore
    # Deep clone the data to avoid modifying the original object
    cloned_data = copy.deepcopy(data)
    if isinstance(cloned_data, dict):
-        for key, value in cloned_data.items():
+        for key, value in cloned_data.items():  # type: ignore
            # Recursively call the function if the value is a dictionary or a list
            if isinstance(value, (dict, list)):
-                cloned_data[key] = truncate_data_strings(value)
+                cloned_data[key] = truncate_data_strings(value)  # type: ignore
            # Truncate the string if it it's long and add ellipsis and length
            elif isinstance(value, str):
-                cloned_data[key] = value[:40]
+                cloned_data[key] = value[:40]  # type: ignore
                if len(value) > 40:
-                    cloned_data[key] += "..." + f" ({len(value)} chars)"
+                    cloned_data[key] += "..." + f" ({len(value)} chars)"  # type: ignore
-    elif isinstance(cloned_data, list):
+    elif isinstance(cloned_data, list):  # type: ignore
        # Process each item in the list
-        cloned_data = [truncate_data_strings(item) for item in cloned_data]
+        cloned_data = [truncate_data_strings(item) for item in cloned_data]  # type: ignore
-    return cloned_data
+    return cloned_data  # type: ignore
--- a/frontend/src/App.tsx
+++ b/frontend/src/App.tsx
@ -33,6 +33,7 @@ import { History } from "./components/history/history_types";
 import HistoryDisplay from "./components/history/HistoryDisplay";
 import { extractHistoryTree } from "./components/history/utils";
 import toast from "react-hot-toast";
 import ImportCodeSection from "./components/ImportCodeSection";
 const IS_OPENAI_DOWN = false;
@ -43,6 +44,7 @@ function App() {
  const [referenceImages, setReferenceImages] = useState<string[]>([]);
  const [executionConsole, setExecutionConsole] = useState<string[]>([]);
  const [updateInstruction, setUpdateInstruction] = useState("");
  const [isImportedFromCode, setIsImportedFromCode] = useState<boolean>(false);
  // Settings
  const [settings, setSettings] = usePersistedState<Settings>(
@ -118,6 +120,8 @@ function App() {
    setReferenceImages([]);
    setExecutionConsole([]);
    setAppHistory([]);
    setCurrentVersion(null);
    setIsImportedFromCode(false);
  };
  const stop = () => {
@ -231,6 +235,7 @@ function App() {
          image: referenceImages[0],
          resultImage: resultImage,
          history: updatedHistory,
          isImportedFromCode,
        },
        currentVersion
      );
@ -240,6 +245,7 @@ function App() {
          generationType: "update",
          image: referenceImages[0],
          history: updatedHistory,
          isImportedFromCode,
        },
        currentVersion
      );
@ -256,6 +262,33 @@ function App() {
    }));
  };
  // TODO: Rename everything to "stack" instead of "config"
  function setStack(stack: GeneratedCodeConfig) {
    setSettings((prev) => ({
      ...prev,
      generatedCodeConfig: stack,
    }));
  }
  function importFromCode(code: string, stack: GeneratedCodeConfig) {
    setIsImportedFromCode(true);
    // Set up this project
    setGeneratedCode(code);
    setStack(stack);
    setAppHistory([
      {
        type: "code_create",
        parentIndex: null,
        code,
        inputs: { code },
      },
    ]);
    setCurrentVersion(0);
    setAppState(AppState.CODE_READY);
  }
  return (
    <div className="mt-2 dark:bg-black dark:text-white">
      {IS_RUNNING_ON_CLOUD && <PicoBadge settings={settings} />}
@ -274,12 +307,7 @@ function App() {
          <OutputSettingsSection
            generatedCodeConfig={settings.generatedCodeConfig}
-            setGeneratedCodeConfig={(config: GeneratedCodeConfig) =>
+            setGeneratedCodeConfig={(config) => setStack(config)}
              setSettings((prev) => ({
                ...prev,
                generatedCodeConfig: config,
              }))
            }
            shouldDisableUpdates={
              appState === AppState.CODING || appState === AppState.CODE_READY
            }
@ -364,22 +392,24 @@ function App() {
              {/* Reference image display */}
              <div className="flex gap-x-2 mt-2">
-                <div className="flex flex-col">
+                {referenceImages.length > 0 && (
-                  <div
+                  <div className="flex flex-col">
-                    className={classNames({
+                    <div
-                      "scanning relative": appState === AppState.CODING,
+                      className={classNames({
-                    })}
+                        "scanning relative": appState === AppState.CODING,
-                  >
+                      })}
-                    <img
+                    >
-                      className="w-[340px] border border-gray-200 rounded-md"
+                      <img
-                      src={referenceImages[0]}
+                        className="w-[340px] border border-gray-200 rounded-md"
-                      alt="Reference"
+                        src={referenceImages[0]}
-                    />
+                        alt="Reference"
                      />
                    </div>
                    <div className="text-gray-400 uppercase text-sm text-center mt-1">
                      Original Screenshot
                    </div>
                  </div>
-                  <div className="text-gray-400 uppercase text-sm text-center mt-1">
+                )}
                    Original Screenshot
                  </div>
                </div>
                <div className="bg-gray-400 px-4 py-2 rounded text-sm hidden">
                  <h2 className="text-lg mb-4 border-b border-gray-800">
                    Console
@ -424,6 +454,7 @@ function App() {
              doCreate={doCreate}
              screenshotOneApiKey={settings.screenshotOneApiKey}
            />
            <ImportCodeSection importFromCode={importFromCode} />
          </div>
        )}
--- a/frontend/src/components/ImportCodeSection.tsx
+++ b/frontend/src/components/ImportCodeSection.tsx
@ -0,0 +1,78 @@
 import { useState } from "react";
 import { Button } from "./ui/button";
 import {
  Dialog,
  DialogContent,
  DialogDescription,
  DialogFooter,
  DialogHeader,
  DialogTitle,
  DialogTrigger,
 } from "./ui/dialog";
 import { Textarea } from "./ui/textarea";
 import OutputSettingsSection from "./OutputSettingsSection";
 import { GeneratedCodeConfig } from "../types";
 import toast from "react-hot-toast";
 interface Props {
  importFromCode: (code: string, stack: GeneratedCodeConfig) => void;
 }
 function ImportCodeSection({ importFromCode }: Props) {
  const [code, setCode] = useState("");
  const [stack, setStack] = useState<GeneratedCodeConfig | undefined>(
    undefined
  );
  const doImport = () => {
    if (code === "") {
      toast.error("Please paste in some code");
      return;
    }
    if (stack === undefined) {
      toast.error("Please select your stack");
      return;
    }
    importFromCode(code, stack);
  };
  return (
    <Dialog>
      <DialogTrigger asChild>
        <Button variant="secondary">Import from Code</Button>
      </DialogTrigger>
      <DialogContent className="sm:max-w-[425px]">
        <DialogHeader>
          <DialogTitle>Paste in your HTML code</DialogTitle>
          <DialogDescription>
            Make sure that the code you're importing is valid HTML.
          </DialogDescription>
        </DialogHeader>
        <Textarea
          value={code}
          onChange={(e) => setCode(e.target.value)}
          className="w-full h-64"
        />
        <OutputSettingsSection
          generatedCodeConfig={stack}
          setGeneratedCodeConfig={(config: GeneratedCodeConfig) =>
            setStack(config)
          }
          label="Stack:"
          shouldDisableUpdates={false}
        />
        <DialogFooter>
          <Button type="submit" onClick={doImport}>
            Import
          </Button>
        </DialogFooter>
      </DialogContent>
    </Dialog>
  );
 }
 export default ImportCodeSection;
--- a/frontend/src/components/OutputSettingsSection.tsx
+++ b/frontend/src/components/OutputSettingsSection.tsx
@ -43,20 +43,22 @@ function generateDisplayComponent(config: GeneratedCodeConfig) {
 }
 interface Props {
-  generatedCodeConfig: GeneratedCodeConfig;
+  generatedCodeConfig: GeneratedCodeConfig | undefined;
  setGeneratedCodeConfig: (config: GeneratedCodeConfig) => void;
  label?: string;
  shouldDisableUpdates?: boolean;
 }
 function OutputSettingsSection({
  generatedCodeConfig,
  setGeneratedCodeConfig,
  label = "Generating:",
  shouldDisableUpdates = false,
 }: Props) {
  return (
    <div className="flex flex-col gap-y-2 justify-between text-sm">
      <div className="grid grid-cols-3 items-center gap-4">
-        <span>Generating:</span>
+        <span>{label}</span>
        <Select
          value={generatedCodeConfig}
          onValueChange={(value: string) =>
@ -65,7 +67,9 @@ function OutputSettingsSection({
          disabled={shouldDisableUpdates}
        >
          <SelectTrigger className="col-span-2" id="output-settings-js">
-            {generateDisplayComponent(generatedCodeConfig)}
+            {generatedCodeConfig
              ? generateDisplayComponent(generatedCodeConfig)
              : "Select a stack"}
          </SelectTrigger>
          <SelectContent>
            <SelectGroup>
--- a/frontend/src/components/history/HistoryDisplay.tsx
+++ b/frontend/src/components/history/HistoryDisplay.tsx
@ -21,6 +21,8 @@ function displayHistoryItemType(itemType: HistoryItemType) {
      return "Create";
    case "ai_edit":
      return "Edit";
    case "code_create":
      return "Imported from code";
    default: {
      const exhaustiveCheck: never = itemType;
      throw new Error(`Unhandled case: ${exhaustiveCheck}`);
@ -62,7 +64,11 @@ export default function HistoryDisplay({
                {" "}
                <div className="flex gap-x-1 truncate">
                  <h2 className="text-sm truncate">
-                    {item.type === "ai_edit" ? item.inputs.prompt : "Create"}
+                    {item.type === "ai_edit"
                      ? item.inputs.prompt
                      : item.type === "ai_create"
                      ? "Create"
                      : "Imported from code"}
                  </h2>
                  {/* <h2 className="text-sm">{displayHistoryItemType(item.type)}</h2> */}
                  {item.parentIndex !== null &&
@ -76,7 +82,11 @@ export default function HistoryDisplay({
              </HoverCardTrigger>
              <HoverCardContent>
                <div>
-                  {item.type === "ai_edit" ? item.inputs.prompt : "Create"}
+                  {item.type === "ai_edit"
                    ? item.inputs.prompt
                    : item.type === "ai_create"
                    ? "Create"
                    : "Imported from code"}
                </div>
                <Badge>{displayHistoryItemType(item.type)}</Badge>
              </HoverCardContent>
--- a/frontend/src/components/history/history_types.ts
+++ b/frontend/src/components/history/history_types.ts
@ -1,4 +1,4 @@
-export type HistoryItemType = "ai_create" | "ai_edit";
+export type HistoryItemType = "ai_create" | "ai_edit" | "code_create";
 type CommonHistoryItem = {
  parentIndex: null | number;
@ -13,6 +13,10 @@ export type HistoryItem =
  | ({
      type: "ai_edit";
      inputs: AiEditInputs;
    } & CommonHistoryItem)
  | ({
      type: "code_create";
      inputs: CodeCreateInputs;
    } & CommonHistoryItem);
 export type AiCreateInputs = {
@ -23,4 +27,8 @@ export type AiEditInputs = {
  prompt: string;
 };
 export type CodeCreateInputs = {
  code: string;
 };
 export type History = HistoryItem[];
--- a/frontend/src/components/history/utils.ts
+++ b/frontend/src/components/history/utils.ts
@ -14,9 +14,11 @@ export function extractHistoryTree(
      if (item.type === "ai_create") {
        // Don't include the image for ai_create
        flatHistory.unshift(item.code);
-      } else {
+      } else if (item.type === "ai_edit") {
        flatHistory.unshift(item.code);
        flatHistory.unshift(item.inputs.prompt);
      } else if (item.type === "code_create") {
        flatHistory.unshift(item.code);
      }
      // Move to the parent of the current item
--- a/frontend/src/generateCode.ts
+++ b/frontend/src/generateCode.ts
@ -12,6 +12,7 @@ export interface CodeGenerationParams {
  image: string;
  resultImage?: string;
  history?: string[];
  isImportedFromCode?: boolean;
  // isImageGenerationEnabled: boolean; // TODO: Merge with Settings type in types.ts
 }
`@ -1,3 +1,3 @@`
	`Run tests`	`# Run tests`

	`pytest test_prompts.py`	`poetry run pytest`