From 3442ce3ae6a468334e94dcb3e77035152b5eb4ea Mon Sep 17 00:00:00 2001
From: Piyush Mishra <piyush_mishra31@outlook.com>
Date: Thu, 23 Nov 2023 11:00:40 +0530
Subject: [PATCH] prompt update for multi image support rev.1

---
 backend/prompts.py | 98 ++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 98 insertions(+)
 create mode 100644 backend/prompts.py

diff --git a/backend/prompts.py b/backend/prompts.py
new file mode 100644
index 0000000..b56a952
--- /dev/null
+++ b/backend/prompts.py
@@ -0,0 +1,98 @@
+from typing import List, Union
+
+from openai.types.chat import ChatCompletionMessageParam, ChatCompletionContentPartParam
+
+from imported_code_prompts import (
+    IMPORTED_CODE_BOOTSTRAP_SYSTEM_PROMPT,
+    IMPORTED_CODE_IONIC_TAILWIND_SYSTEM_PROMPT,
+    IMPORTED_CODE_REACT_TAILWIND_SYSTEM_PROMPT,
+    IMPORTED_CODE_TAILWIND_SYSTEM_PROMPT,
+)
+from screenshot_system_prompts import (
+    BOOTSTRAP_SYSTEM_PROMPT,
+    IONIC_TAILWIND_SYSTEM_PROMPT,
+    REACT_TAILWIND_SYSTEM_PROMPT,
+    TAILWIND_SYSTEM_PROMPT,
+)
+
+
+USER_PROMPT = """
+Generate code for a web page that looks exactly like 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(
+    images: List[str],
+    generated_code_config: str,
+    result_image_data_url: Union[str, None] = None,
+) -> List[ChatCompletionMessageParam]:
+    # Set the system prompt based on the output settings
+    system_content = TAILWIND_SYSTEM_PROMPT
+    if generated_code_config == "html_tailwind":
+        system_content = TAILWIND_SYSTEM_PROMPT
+    elif generated_code_config == "react_tailwind":
+        system_content = REACT_TAILWIND_SYSTEM_PROMPT
+    elif generated_code_config == "bootstrap":
+        system_content = BOOTSTRAP_SYSTEM_PROMPT
+    elif generated_code_config == "ionic_tailwind":
+        system_content = IONIC_TAILWIND_SYSTEM_PROMPT
+    else:
+        raise Exception("Code config is not one of available options")
+
+    user_content: List[ChatCompletionContentPartParam]= []
+    for image_data_url in images:
+        user_content.append(
+            {
+                "type": "image_url",
+                "image_url": {"url": image_data_url, "detail": "high"},
+            },
+        )
+
+    if result_image_data_url:
+        user_content.insert(1, {
+            "type": "image_url",
+            "image_url": {"url": result_image_data_url, "detail": "high"},
+        })
+
+    user_content.append(
+        {
+            "type": "text",
+            "text": USER_PROMPT,
+        }
+    )
+
+    return [
+        {"role": "system", "content": system_content},
+        {
+            "role": "user",
+            "content": user_content,
+        },
+    ]