Use manim-trainer prompt templates

Co-authored-by: Codex <codex@openai.com>

app.py

@@ -214,17 +214,6 @@ footer,
   box-shadow: inset 0 0 22px rgba(0, 255, 204, 0.08);
 }
 
-.viewport-frame::before {
-  content: "";
-  position: absolute;
-  top: 10px;
-  left: 14px;
-  z-index: 5;
-  color: var(--cyber-cyan);
-  font-size: 12px;
-  text-shadow: 0 0 8px var(--cyber-cyan);
-}
-
 .viewport-frame::after {
   content: "";
   position: absolute;
@@ -482,7 +471,7 @@ with gr.Blocks(css=custom_css, title="SciVisual-Agent") as demo:
             with gr.Tabs():
                 with gr.Tab("Rendered Video"):
                     video = gr.Video(label=None, elem_id="render-video", height=560)
-                with gr.Tab("Generated Source"):
+                with gr.Tab("Generated Code"):
                     source = gr.Textbox(
                         label="Clean Python Source",
                         lines=24,

generator.py

@@ -15,6 +15,7 @@ from pathlib import Path
 from typing import Dict, List, Optional, Tuple
 
 from executor import render_manim_scene
+from prompt_templates import build_feedback_messages, build_initial_messages
 
 
 DEFAULT_MODEL_PATH = "vankhieu/Seed_Coder_8B_Instruct_unsloth_bnb_4bit_lora_r8_sft_grpo_rw_mean_text_visual"
@@ -22,25 +23,6 @@ DEFAULT_SELECTED_MODEL = "unsloth/Seed-Coder-8B-Instruct-unsloth-bnb-4bit"
 THINKING_TOKEN_ID = 151668
 THINKING_TOKEN = "</think>"
 
-SYSTEM_PROMPT = """You are SciVisual-Agent, a Physics and Mathematics visualization genius using ManimCE.
-Return only complete executable Python code for a Manim scene.
-
-Rules:
-- Put the code inside exactly one <CODE>...</CODE> block or exactly one ```python ... ``` block.
-- Use `from manim import *`.
-- Define a renderable Scene class, preferably `MainScene`.
-- Always include `construct(self)`.
-- Avoid external assets, file reads, network calls, audio, and interactive input.
-- Keep the animation concise enough for low-quality fast rendering.
-- Use readable labels and high-contrast colors on a dark background.
-"""
-
-MATH_PROMPT = """Generate ManimCE code for this mathematics visualization.
-Emphasize intuition, labeled diagrams, clean equations, and a clear progression."""
-
-PHYSICS_PROMPT = """Generate ManimCE code for this physics visualization.
-Emphasize motion, forces, fields, trajectories, waves, and cause/effect relationships."""
-
 
 @dataclass
 class LocalModelConfig:
@@ -371,18 +353,8 @@ class ManimVisualAgent:
         return completion
 
     def _initial_messages(self, user_prompt: str, domain: str) -> List[Dict[str, str]]:
-        domain_prompt = PHYSICS_PROMPT if domain.lower().startswith("phys") else MATH_PROMPT
-        return [
-            {"role": "system", "content": SYSTEM_PROMPT},
-            {
-                "role": "user",
-                "content": (
-                    f"{domain_prompt}\n\n"
-                    f"Text script:\n{user_prompt.strip()}\n\n"
-                    "Now generate the ManimCE Python code."
-                ),
-            },
-        ]
+        del domain
+        return build_initial_messages(user_prompt)
 
     def _feedback_messages(
         self,
@@ -391,21 +363,8 @@ class ManimVisualAgent:
         initial_code: str,
         render_errors: str,
     ) -> List[Dict[str, str]]:
-        domain_prompt = PHYSICS_PROMPT if domain.lower().startswith("phys") else MATH_PROMPT
-        return [
-            {"role": "system", "content": SYSTEM_PROMPT},
-            {
-                "role": "user",
-                "content": (
-                    f"{domain_prompt}\n\n"
-                    f"Text script:\n{user_prompt.strip()}\n\n"
-                    "The previous ManimCE code failed to render. Fix it completely.\n\n"
-                    f"<INITIAL_CODE>\n{initial_code}\n</INITIAL_CODE>\n\n"
-                    f"<RENDER_ERRORS>\n{render_errors}\n</RENDER_ERRORS>\n\n"
-                    "Return the corrected complete code in one <CODE>...</CODE> block."
-                ),
-            },
-        ]
+        del domain
+        return build_feedback_messages(user_prompt, initial_code, render_errors)
 
     def _fallback_scene(self, prompt: str) -> str:
         safe_message = (self.model_error or "Local model is not ready.").replace("\\", "\\\\").replace('"', '\\"')

prompt_templates.py

@@ -0,0 +1,154 @@
+"""Prompt templates aligned with the manim-trainer fine-tuning format."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List
+
+
+@dataclass(frozen=True)
+class PromptChatTemplate:
+    system_prompt_template: str
+    user_prompt_template: str
+    completion_prompt_template: str = ""
+
+
+MANIM_VIDEO_GEN_CHAT_TEMPLATE = PromptChatTemplate(
+    system_prompt_template="""You are an expert Manim Community Edition (ManimCE) educator and Python developer.
+
+Please follow these steps precisely:
+1. Read the `<TEXT_SCRIPT>` block which contains a description of a Manim animation that user want to create.
+2. DO NOT provide any explanations.
+3. Generate only executable Python code for Manim, wrapped between `<CODE>` and `</CODE>`.
+4. Do not include any explanations, comments, or instructions on how to run the code. Only include the code wrapped between `<CODE>` and `</CODE>`.
+
+
+Example format:
+
+<TEXT_SCRIPT>
+Display a red square centered on screen, then transform it into a circle.
+</TEXT_SCRIPT>
+
+<CODE>
+```python
+from manim import *
+class RedSquareToCircle(Scene):
+    def construct(self):
+        square = Square(color=RED)
+        self.play(Create(square))
+        circle = Circle(color=RED)
+        self.play(Transform(square, circle))
+        self.wait()
+``` 
+</CODE>""",
+    user_prompt_template="""Generate code for the following text script:
+<TEXT_SCRIPT>
+{reviewed_description}
+</TEXT_SCRIPT>
+
+""",
+    completion_prompt_template="""<CODE>
+```python
+{code}
+```
+</CODE>""",
+)
+
+
+MANIM_VIDEO_GEN_CHAT_FB_ONLY_TEMPLATE = PromptChatTemplate(
+    system_prompt_template="""You are an expert Manim Community Edition (ManimCE) educator and Python developer.
+
+Please follow these steps precisely:
+1. Read the `<TEXT_SCRIPT>` block which contains a description of a Manim animation that I want to create.
+2. Analyze the initial code provided in the `<INITIAL_CODE>` block.
+3. Use the provided render errors (if any) in `<RENDER_ERRORS>` block to revise and fix the initial code.
+4. Generate only executable Python code for Manim, wrapped between `<CODE>` and `</CODE>`.
+5. DO NOT include any explanations, comments, or instructions on how to run the code. Only include the code wrapped between `<CODE>` and `</CODE>`.
+
+
+Example format:
+
+<TEXT_SCRIPT>
+Display a red square centered on screen, then transform it into a circle.
+</TEXT_SCRIPT>
+
+<INITIAL_CODE>
+```python
+from manim import *
+
+class RedSquareToCircle(Scene):
+    def construct(self):
+        square = Square(bg_color=RED)
+        self.play(Create(square))
+        circle = Circle(bg_color=RED)
+        self.play(Transform(square, circle))
+        self.wait()
+``` 
+</INITIAL_CODE>
+
+<RENDER_ERRORS>
+```text
+TypeError: Mobject.__init__() got an unexpected keyword argument 'bg_color'
+```
+</RENDER_ERRORS>
+
+<CODE>
+```python
+from manim import *
+
+class RedSquareToCircle(Scene):
+    def construct(self):
+        square = Square(color=RED)
+        self.play(Create(square))
+        circle = Circle(color=RED)
+        self.play(Transform(square, circle))
+        self.wait()
+```
+</CODE>""",
+    user_prompt_template="""Generate code for the following text script adapting initial code based on the render errors occured from the initial code:
+
+<TEXT_SCRIPT>
+{reviewed_description}
+</TEXT_SCRIPT>
+
+<INITIAL_CODE>
+```python
+{initial_code}
+```
+</INITIAL_CODE>
+
+<RENDER_ERRORS>
+```text
+{render_errors}
+```
+</RENDER_ERRORS>
+
+""",
+)
+
+
+def build_initial_messages(text_script: str) -> List[Dict[str, str]]:
+    return [
+        {"role": "system", "content": MANIM_VIDEO_GEN_CHAT_TEMPLATE.system_prompt_template},
+        {
+            "role": "user",
+            "content": MANIM_VIDEO_GEN_CHAT_TEMPLATE.user_prompt_template.format(
+                reviewed_description=text_script.strip()
+            ),
+        },
+    ]
+
+
+def build_feedback_messages(text_script: str, initial_code: str, render_errors: str) -> List[Dict[str, str]]:
+    return [
+        {"role": "system", "content": MANIM_VIDEO_GEN_CHAT_FB_ONLY_TEMPLATE.system_prompt_template},
+        {
+            "role": "user",
+            "content": MANIM_VIDEO_GEN_CHAT_FB_ONLY_TEMPLATE.user_prompt_template.format(
+                reviewed_description=text_script.strip(),
+                initial_code=initial_code.strip(),
+                render_errors=(render_errors or "NONE").strip() or "NONE",
+            ),
+        },
+    ]
+