feat!: Add advanced mode support

Zhou-Shilin · Zhou-Shilin · commit 8de8816ead38 · 2024-05-05T20:58:19.000+08:00
diff --git a/console.py b/console.py
@@ -37,6 +37,24 @@ def generate_plugin(description):
     
     return schem
 
+def get_schematic_advanced(description):
+    print("(Advanced Mode) Generating programme...")
+    programme = core.askgpt(config.BTR_DESC_SYS_GEN, config.BTR_DESC_USR_GEN.replace("%DESCRIPTION%", description), config.GENERATE_MODEL, disable_json_mode=True)
+
+    print("(Advanced Mode) Generating image tag...")
+    image_tag = core.askgpt(config.IMG_TAG_SYS_GEN, config.IMG_TAG_USR_GEN.replace("%PROGRAMME%", programme), config.GENERATE_MODEL, disable_json_mode=True)
+
+    print("(Advanced Mode) Generating image...")
+    tag = image_tag + ", minecraft)"
+    image_url = core.ask_dall_e(tag)
+
+    print("(Advanced Mode) Generating schematic...")
+    response = core.askgpt(config.SYS_GEN_ADV, config.USR_GEN_ADV.replace("%DESCRIPTION%", description), config.VISION_MODEL, image_url=image_url)
+
+    schem = core.text_to_schem(response)
+
+    return schem
+
 if __name__ == "__main__":
     core.initialize()
 
@@ -54,7 +72,11 @@ def generate_plugin(description):
 
     print("Generating...")
 
-    schem = generate_plugin(description)
+    if config.ADVANCED_MODE:
+        print("Advanced mode is enabled. Generating a schematic with advanced features.")
+        schem = get_schematic_advanced(description)
+    else:
+        schem = generate_plugin(description)
 
     logger(f"console: Saving {name}.schem to generated/ folder.")
     version_tag = core.input_version_to_mcs_tag(version)
diff --git a/core.py b/core.py
@@ -2,6 +2,9 @@
 import mcschematic
 import sys
 import json
+import requests
+import base64
+import uuid
 
 from log_writer import logger
 import config
@@ -31,24 +34,46 @@ def askgpt(system_prompt: str, user_prompt: str, model_name: str):
     Returns:
         str: The response from ChatGPT.
     """
-    client = OpenAI(api_key=config.API_KEY, base_url=config.BASE_URL)
+    if image_url is not None and config.USE_DIFFERENT_APIKEY_FOR_VISION_MODEL:
+        logger("Using different API key for vision model.")
+        client = OpenAI(api_key=config.VISION_API_KEY, base_url=config.VISION_BASE_URL)
+    else:
+        client = OpenAI(api_key=config.API_KEY, base_url=config.BASE_URL)
+
     logger("Initialized the OpenAI client.")
 
     # Define the messages for the conversation
-    messages = [
-        {"role": "system", "content": system_prompt},
-        {"role": "user", "content": user_prompt}
-    ]
+    if image_url is not None:
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": [
+                {"type": "text", "text": user_prompt},
+                {"type": "image_url", "image_url": {"url": image_url}}
+                ]
+            }
+        ]
+    else:
+        messages = [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt}
+        ]
+
 
     logger(f"askgpt: system {system_prompt}")
     logger(f"askgpt: user {user_prompt}")
 
     # Create a chat completion
-    response = client.chat.completions.create(
-        model=model_name,
-        response_format={"type": "json_object"},
-        messages=messages
-    )
+    if disable_json_mode:
+        response = client.chat.completions.create(
+            model=model_name,
+            messages=messages
+        )
+    else:
+        response = client.chat.completions.create(
+            model=model_name,
+            response_format={"type": "json_object"},
+            messages=messages
+        )
 
     logger(f"askgpt: response {response}")
 
@@ -57,6 +82,37 @@ def askgpt(system_prompt: str, user_prompt: str, model_name: str):
     logger(f"askgpt: extracted reply {assistant_reply}")
     return assistant_reply
 
+def ask_dall_e(description: str):
+    """
+    Generates a design image using the DALL-E API.
+
+    Args:
+        description (str): The prompt or description for generating the image.
+
+    Returns:
+        str: The URL of the generated image.
+    """
+    if config.USE_DIFFERENT_APIKEY_FOR_DALLE_MODEL:
+        client = OpenAI(api_key=config.DALLE_API_KEY, base_url=config.DALLE_BASE_URL)
+    else:
+        client = OpenAI(api_key=config.API_KEY, base_url=config.BASE_URL)
+
+    logger("ask_dall_e: Generating design image using DALL-E API.")
+
+    response = client.images.generate(
+        model=config.IMAGE_GENERATION_MODEL,
+        prompt=description,
+        size=config.IMAGE_SIZE,
+        quality="standard",
+        n=1,
+    )
+
+    image_url = response.data[0].url
+
+    logger(f"ask_dall_e: Generated image URL {image_url}")
+
+    return image_url
+
 def text_to_schem(text: str):
     """
     Converts a JSON string to a Minecraft schematic.
diff --git a/ui.py b/ui.py
@@ -29,6 +29,24 @@ def get_schematic(description):
     
     return schem
 
+def get_schematic_advanced(description):
+    print("(Advanced Mode) Generating programme...")
+    programme = core.askgpt(config.BTR_DESC_SYS_GEN, config.BTR_DESC_USR_GEN.replace("%DESCRIPTION%", description), config.GENERATE_MODEL, disable_json_mode=True)
+
+    print("(Advanced Mode) Generating image tag...")
+    image_tag = core.askgpt(config.IMG_TAG_SYS_GEN, config.IMG_TAG_USR_GEN.replace("%PROGRAMME%", programme), config.GENERATE_MODEL, disable_json_mode=True)
+
+    print("(Advanced Mode) Generating image...")
+    tag = image_tag + ", minecraft)"
+    image_url = core.ask_dall_e(tag)
+
+    print("(Advanced Mode) Generating schematic...")
+    response = core.askgpt(config.SYS_GEN_ADV, config.USR_GEN_ADV.replace("%DESCRIPTION%", description), config.VISION_MODEL, image_url=image_url)
+
+    schem = core.text_to_schem(response)
+
+    return schem
+
 def generate_schematic():
     """
     Generates a schematic file based on user input.
@@ -42,6 +60,11 @@ def generate_schematic():
     """
     generate_button.config(state=tk.DISABLED, text="Generating...")
 
+    if config.ADVANCED_MODE:
+        msgbox.showwarning("Warning", "You are using advanced mode. This mode will generate schematic with higher quality, but it may take longer to generate.")
+
+    msgbox.showinfo("Info", "It is expected to take 30 seconds to 5 minutes. The programme may \"not responding\", this is normal, just be patient. DO NOT CLOSE THE PROGRAM. Click the button below to start generating.")
+
     version = version_entry.get()
     name = name_entry.get()
     description = description_entry.get()
@@ -50,7 +73,10 @@ def generate_schematic():
     logger(f"console: input name {name}")
     logger(f"console: input description {description}")
 
-    schem = get_schematic(description)
+    if config.ADVANCED_MODE:
+        schem = get_schematic_advanced(description)
+    else:
+        schem = get_schematic(description)
 
     logger(f"console: Saving {name}.schem to generated/ folder.")
     version_tag = core.input_version_to_mcs_tag(version)
