dc7a5a09da
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.
56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
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."
|
|
|
|
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: Optional[str] = None
|
|
) -> List[ChatCompletionMessageParam]:
|
|
system_content = STACK_TO_SYSTEM_CONTENT[stack]
|
|
user_content = (
|
|
f"Here is the code of the app: {code}" if stack != "svg" else f"Here is the code of the SVG: {code}"
|
|
)
|
|
return [
|
|
{
|
|
"role": "system",
|
|
"content": system_content,
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": user_content,
|
|
},
|
|
]
|
|
# TODO: Use result_image_data_url
|
|
|
|
def assemble_prompt(
|
|
image_data_url: str,
|
|
stack: Stack,
|
|
result_image_data_url: Optional[str] = None,
|
|
) -> List[ChatCompletionMessageParam]:
|
|
system_content = STACK_TO_SYSTEM_CONTENT[stack]
|
|
user_prompt = STACK_TO_USER_PROMPT[stack]
|
|
|
|
user_content = [
|
|
{
|
|
"type": "image_url",
|
|
"image_url": {"url": image_data_url, "detail": "high"},
|