Github/screenshot-to-code
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update __init__.py

Browse Source 
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:
GANESH W 2024-03-04 23:48:45 +05:30 committed by GitHub
parent 072b286b6d
commit dc7a5a09da
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 22 additions and 46 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

68
backend/prompts/__init__.py
View File

@ -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 openai.types.chat import ChatCompletionMessageParam, ChatCompletionContentPartParam
from prompts.imported_code_prompts import IMPORTED_CODE_SYSTEM_PROMPTS from prompts.imported_code_prompts import IMPORTED_CODE_SYSTEM_PROMPTS
from prompts.screenshot_system_prompts import SYSTEM_PROMPTS from prompts.screenshot_system_prompts import SYSTEM_PROMPTS
from prompts.types import Stack 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 = """ STACK_TO_SYSTEM_CONTENT: dict[Stack, str] = {
Generate code for a web page that looks exactly like this. "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.",
}
SVG_USER_PROMPT = """
Generate code for a SVG that looks exactly like this.
"""
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( 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]: ) -> List[ChatCompletionMessageParam]:
system_content = IMPORTED_CODE_SYSTEM_PROMPTS[stack] system_content = STACK_TO_SYSTEM_CONTENT[stack]
user_content = ( user_content = (
"Here is the code of the app: " + code f"Here is the code of the app: {code}" if stack != "svg" else f"Here is the code of the SVG: {code}"
if stack != "svg"
else "Here is the code of the SVG: " + code
) )
return [ return [
{ {
@ -38,42 +41,15 @@ def assemble_imported_code_prompt(
] ]
# TODO: Use result_image_data_url # TODO: Use result_image_data_url
def assemble_prompt( def assemble_prompt(
image_data_url: str, image_data_url: str,
stack: Stack, stack: Stack,
result_image_data_url: Union[str, None] = None, result_image_data_url: Optional[str] = None,
) -> List[ChatCompletionMessageParam]: ) -> List[ChatCompletionMessageParam]:
system_content = SYSTEM_PROMPTS[stack] system_content = STACK_TO_SYSTEM_CONTENT[stack]
user_prompt = USER_PROMPT if stack != "svg" else SVG_USER_PROMPT user_prompt = STACK_TO_USER_PROMPT[stack]
user_content: List[ChatCompletionContentPartParam] = [ user_content = [
{ {
"type": "image_url", "type": "image_url",
"image_url": {"url": image_data_url, "detail": "high"}, "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,
},
]