Github/screenshot-to-code
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

initial implementation of importing from code

Browse Source
This commit is contained in:
Abi Raja 2023-12-09 21:00:18 -05:00
parent 435402bc85
commit 52fee9e49b
9 changed files with 199 additions and 48 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
backend/imported_code_prompts.py Normal file
View File

@ -0,0 +1,16 @@
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.
"""

19
backend/prompts.py
View File

@ -2,6 +2,8 @@ from typing import List, Union
from openai.types.chat import ChatCompletionMessageParam, ChatCompletionContentPartParam
from imported_code_prompts import IMPORTED_CODE_TAILWIND_SYSTEM_PROMPT
TAILWIND_SYSTEM_PROMPT = """
You are an expert Tailwind developer
@ -121,6 +123,23 @@ Generate code for a web page that looks exactly like this.
"""
def assemble_imported_code_prompt(
code: str, result_image_data_url: Union[str, None] = None
) -> List[ChatCompletionMessageParam]:
system_content = IMPORTED_CODE_TAILWIND_SYSTEM_PROMPT
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: str,
generated_code_config: str,

34
backend/routes/generate_code.py
View File

@ -8,11 +8,13 @@ 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_prompt
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()
@ -122,6 +124,26 @@ async def stream_code(websocket: WebSocket):
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)
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"]:
@ -129,7 +151,9 @@ async def stream_code(websocket: WebSocket):
params["image"], generated_code_config, params["resultImage"]
)
else:
prompt_messages = assemble_prompt(params["image"], generated_code_config)
prompt_messages = assemble_prompt(
params["image"], generated_code_config
)
except:
await websocket.send_json(
{
@ -140,11 +164,8 @@ async def stream_code(websocket: WebSocket):
await websocket.close()
return
# Image cache for updates so that we don't have to regenerate images
image_cache: Dict[str, str] = {}
if params["generationType"] == "update":
# Transform into message format
# Transform the history tree into message format
# TODO: Move this to frontend
for index, text in enumerate(params["history"]):
if index % 2 == 0:
@ -158,6 +179,7 @@ async def stream_code(websocket: WebSocket):
"content": text,
}
prompt_messages.append(message)
image_cache = create_alt_url_mapping(params["history"][-2])
if SHOULD_MOCK_AI_RESPONSE:

24
frontend/src/App.tsx
View File

@ -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,21 @@ function App() {
}));
};
function importFromCode(code: string) {
setAppState(AppState.CODE_READY);
setGeneratedCode(code);
setAppHistory([
{
type: "code_create",
parentIndex: null,
code,
inputs: { code },
},
]);
setCurrentVersion(0);
setIsImportedFromCode(true);
}
return (
<div className="mt-2 dark:bg-black dark:text-white">
{IS_RUNNING_ON_CLOUD && <PicoBadge settings={settings} />}
@ -364,6 +385,7 @@ function App() {
{/* Reference image display */}
<div className="flex gap-x-2 mt-2">
{referenceImages.length > 0 && (
<div className="flex flex-col">
<div
className={classNames({
@ -380,6 +402,7 @@ function App() {
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 +447,7 @@ function App() {
doCreate={doCreate}
screenshotOneApiKey={settings.screenshotOneApiKey}
/>
<ImportCodeSection importFromCode={importFromCode} />
</div>
)}

49
frontend/src/components/ImportCodeSection.tsx Normal file
View File

@ -0,0 +1,49 @@
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";
interface Props {
importFromCode: (code: string) => void;
}
function ImportCodeSection({ importFromCode }: Props) {
const [code, setCode] = useState("");
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"
/>
<DialogFooter>
<Button type="submit" onClick={() => importFromCode(code)}>
Import
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
);
}
export default ImportCodeSection;

14
frontend/src/components/history/HistoryDisplay.tsx
View File

@ -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>

10
frontend/src/components/history/history_types.ts
View File

@ -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[];

4
frontend/src/components/history/utils.ts
View File

@ -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

1
frontend/src/generateCode.ts
View File

@ -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
}