initial implementation of importing from code

2023-12-09 21:00:18 -05:00 · 2023-12-09 21:00:18 -05:00 · 52fee9e49b
commit 52fee9e49b
parent 435402bc85
9 changed files with 199 additions and 48 deletions
--- a/backend/imported_code_prompts.py
+++ b/backend/imported_code_prompts.py
@ -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.
+"""
--- a/backend/prompts.py
+++ b/backend/prompts.py
@ -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,
--- a/backend/routes/generate_code.py
+++ b/backend/routes/generate_code.py
@ -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,43 +124,63 @@ async def stream_code(websocket: WebSocket):
    async def process_chunk(content: str):
        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: Dict[str, str] = {}

-    if params["generationType"] == "update":
-        # Transform into message format
-        # TODO: Move this to frontend
-        for index, text in enumerate(params["history"]):
+    # 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": "assistant",
-                    "content": text,
-                }
-            else:
                message: ChatCompletionMessageParam = {
                    "role": "user",
                    "content": text,
                }
+            else:
+                message: ChatCompletionMessageParam = {
+                    "role": "assistant",
+                    "content": text,
+                }
            prompt_messages.append(message)
-        image_cache = create_alt_url_mapping(params["history"][-2])
+    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])

    if SHOULD_MOCK_AI_RESPONSE:
        completion = await mock_completion(process_chunk)
--- 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,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,22 +385,24 @@ function App() {

              {/* Reference image display */}
              <div className="flex gap-x-2 mt-2">
-                <div className="flex flex-col">
-                  <div
-                    className={classNames({
-                      "scanning relative": appState === AppState.CODING,
-                    })}
-                  >
-                    <img
-                      className="w-[340px] border border-gray-200 rounded-md"
-                      src={referenceImages[0]}
-                      alt="Reference"
-                    />
+                {referenceImages.length > 0 && (
+                  <div className="flex flex-col">
+                    <div
+                      className={classNames({
+                        "scanning relative": appState === AppState.CODING,
+                      })}
+                    >
+                      <img
+                        className="w-[340px] border border-gray-200 rounded-md"
+                        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 +447,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,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;
--- 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
 }