From dc7a5a09da8bce24391adb76387a1fecbf3db248 Mon Sep 17 00:00:00 2001
From: GANESH W <155592898+ganeshwhere@users.noreply.github.com>
Date: Mon, 4 Mar 2024 23:48:45 +0530
Subject: [PATCH] Update __init__.py

I have used a dictionary to map stack names to system content, instead of using an if-else statement to determine the system content. This would make the code easier to read and maintain.
---
 backend/prompts/__init__.py | 68 ++++++++++++-------------------------
 1 file changed, 22 insertions(+), 46 deletions(-)

diff --git a/backend/prompts/__init__.py b/backend/prompts/__init__.py
index 4f2e329..0e1f564 100644
--- a/backend/prompts/__init__.py
+++ b/backend/prompts/__init__.py
@@ -1,30 +1,33 @@
-from typing import List, NoReturn, Union
+from typing import List, Literal, Optional, Union
 
 from openai.types.chat import ChatCompletionMessageParam, ChatCompletionContentPartParam
-
 from prompts.imported_code_prompts import IMPORTED_CODE_SYSTEM_PROMPTS
 from prompts.screenshot_system_prompts import SYSTEM_PROMPTS
 from prompts.types import Stack
 
+USER_PROMPT: Literal["Generate code for a web page that looks exactly like this."] = (
+    "Generate code for a web page that looks exactly like this."
+)
+SVG_USER_PROMPT: Literal[
+    "Generate code for a SVG that looks exactly like this."
+] = "Generate code for a SVG that looks exactly like this."
 
-USER_PROMPT = """
-Generate code for a web page that looks exactly like this.
-"""
-
-SVG_USER_PROMPT = """
-Generate code for a SVG that looks exactly like this.
-"""
+STACK_TO_SYSTEM_CONTENT: dict[Stack, str] = {
+    "web": "You are a helpful assistant that generates code for a web page.",
+    "svg": "You are a helpful assistant that generates code for a SVG.",
+}
 
+STACK_TO_USER_PROMPT: dict[Stack, str] = {
+    "web": "Generate code for a web page that looks exactly like this.",
+    "svg": "Generate code for a SVG that looks exactly like this.",
+}
 
 def assemble_imported_code_prompt(
-    code: str, stack: Stack, result_image_data_url: Union[str, None] = None
+    code: str, stack: Stack, result_image_data_url: Optional[str] = None
 ) -> List[ChatCompletionMessageParam]:
-    system_content = IMPORTED_CODE_SYSTEM_PROMPTS[stack]
-
+    system_content = STACK_TO_SYSTEM_CONTENT[stack]
     user_content = (
-        "Here is the code of the app: " + code
-        if stack != "svg"
-        else "Here is the code of the SVG: " + code
+        f"Here is the code of the app: {code}" if stack != "svg" else f"Here is the code of the SVG: {code}"
     )
     return [
         {
@@ -38,42 +41,15 @@ def assemble_imported_code_prompt(
     ]
     # TODO: Use result_image_data_url
 
-
 def assemble_prompt(
     image_data_url: str,
     stack: Stack,
-    result_image_data_url: Union[str, None] = None,
+    result_image_data_url: Optional[str] = None,
 ) -> List[ChatCompletionMessageParam]:
-    system_content = SYSTEM_PROMPTS[stack]
-    user_prompt = USER_PROMPT if stack != "svg" else SVG_USER_PROMPT
+    system_content = STACK_TO_SYSTEM_CONTENT[stack]
+    user_prompt = STACK_TO_USER_PROMPT[stack]
 
-    user_content: List[ChatCompletionContentPartParam] = [
+    user_content = [
         {
             "type": "image_url",
             "image_url": {"url": image_data_url, "detail": "high"},
-        },
-        {
-            "type": "text",
-            "text": user_prompt,
-        },
-    ]
-
-    # Include the result image if it exists
-    if result_image_data_url:
-        user_content.insert(
-            1,
-            {
-                "type": "image_url",
-                "image_url": {"url": result_image_data_url, "detail": "high"},
-            },
-        )
-    return [
-        {
-            "role": "system",
-            "content": system_content,
-        },
-        {
-            "role": "user",
-            "content": user_content,
-        },
-    ]