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.
This commit is contained in:
parent
072b286b6d
commit
dc7a5a09da
@ -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,
|
||||
},
|
||||
]
|
||||
|
||||
Loading…
Reference in New Issue
Block a user