110 lines
2.7 KiB
Python
110 lines
2.7 KiB
Python
from typing import List, NoReturn, Union
|
|
|
|
from openai.types.chat import ChatCompletionMessageParam, ChatCompletionContentPartParam
|
|
from llm import Llm
|
|
|
|
from prompts.imported_code_prompts import IMPORTED_CODE_SYSTEM_PROMPTS
|
|
from prompts.screenshot_system_prompts import SYSTEM_PROMPTS
|
|
from prompts.text_prompts import SYSTEM_PROMPTS as TEXT_SYSTEM_PROMPTS
|
|
from prompts.types import Stack
|
|
|
|
|
|
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.
|
|
"""
|
|
|
|
|
|
def assemble_imported_code_prompt(
|
|
code: str, stack: Stack, model: Llm
|
|
) -> List[ChatCompletionMessageParam]:
|
|
system_content = IMPORTED_CODE_SYSTEM_PROMPTS[stack]
|
|
|
|
user_content = (
|
|
"Here is the code of the app: " + code
|
|
if stack != "svg"
|
|
else "Here is the code of the SVG: " + code
|
|
)
|
|
|
|
if model == Llm.CLAUDE_3_5_SONNET_2024_06_20:
|
|
return [
|
|
{
|
|
"role": "system",
|
|
"content": system_content + "\n " + user_content,
|
|
}
|
|
]
|
|
else:
|
|
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: Union[str, None] = None,
|
|
) -> List[ChatCompletionMessageParam]:
|
|
system_content = SYSTEM_PROMPTS[stack]
|
|
user_prompt = USER_PROMPT if stack != "svg" else SVG_USER_PROMPT
|
|
|
|
user_content: List[ChatCompletionContentPartParam] = [
|
|
{
|
|
"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,
|
|
},
|
|
]
|
|
|
|
|
|
def assemble_text_prompt(
|
|
text_prompt: str,
|
|
stack: Stack,
|
|
) -> List[ChatCompletionMessageParam]:
|
|
|
|
system_content = TEXT_SYSTEM_PROMPTS[stack]
|
|
|
|
return [
|
|
{
|
|
"role": "system",
|
|
"content": system_content,
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": "Generate UI for " + text_prompt,
|
|
},
|
|
]
|