Spaces:
Running
Running
File size: 9,197 Bytes
0805c5b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 |
"""
Code Generation Tools for Manim MCP Server
This module provides tools for generating and refining Manim animation code.
"""
import logging
from typing import Any, Dict, Optional
from mcp.types import CallToolResult, TextContent
from utils.hf_wrapper import HFInferenceWrapper, ModelConfig
logger = logging.getLogger(__name__)
async def generate_manim_code(
hf_wrapper: HFInferenceWrapper, arguments: Dict[str, Any]
) -> CallToolResult:
"""
Generate Manim Python code for an animation concept.
Uses a code LLM to generate complete, runnable Manim code based on:
- A concept description
- Scene details
- Desired visual elements
- Optional error feedback for retries
Args:
hf_wrapper: HuggingFace inference wrapper instance
arguments: Dictionary containing:
- concept (str): The animation concept
- scene_description (str): Detailed scene description
- visual_elements (list, optional): List of visual elements to include
- model (str, optional): Hugging Face model to use
- previous_code (str, optional): Previous code attempt (for retries)
- error_message (str, optional): Error from previous attempt (for retries)
Returns:
CallToolResult with the generated Manim code
"""
concept = arguments["concept"]
scene_description = arguments["scene_description"]
visual_elements = arguments.get("visual_elements", [])
model = arguments.get("model")
previous_code = arguments.get("previous_code")
error_message = arguments.get("error_message")
try:
model_config = ModelConfig()
selected_model = model or model_config.code_models[0]
# Build prompt based on whether this is a retry
if previous_code and error_message:
prompt = f"""
You are an expert animation engineer using Manim Community Edition (v0.18.0+).
The previous code attempt had an error. Your task is to FIX the code.
PREVIOUS CODE:
```python
{previous_code}
```
ERROR ENCOUNTERED:
{error_message}
TASK: Fix the error in the code above. Pay special attention to:
- Closing all parentheses, brackets, and braces
- Completing all function calls
- Proper indentation
- Valid Python syntax
Concept: {concept}
Scene Description: {scene_description}
Visual Elements: {", ".join(visual_elements)}
STRICT CODE REQUIREMENTS:
1. Header: MUST start with `from manim import *`
2. Class Structure: Define a class inheriting from `MovingCameraScene` (use this instead of `Scene` to enable camera zoom/pan with `self.camera.frame`)
3. Method: All logic must be inside the `def construct(self):` method
4. SYNTAX: Ensure ALL parentheses, brackets, and function calls are properly closed
5. Colors: Use ONLY valid Manim colors (WHITE, BLACK, RED, GREEN, BLUE, YELLOW, ORANGE, PINK, PURPLE, TEAL, GOLD, etc.)
6. Text: Use `Text()` objects for strings
7. Positioning: Use `.next_to()`, `.move_to()`, or `.shift()`
8. Animations: Use Write(), Create(), FadeIn(), FadeOut(), Transform(), Flash(), Indicate() - capitalize properly!
9. Pacing: Include `self.wait(1)` between animations
OUTPUT FORMAT:
Provide ONLY the complete, corrected Python code. No markdown blocks. No explanations.
"""
else:
prompt = f"""
You are an expert animation engineer using Manim Community Edition (v0.18.0+).
Generate a complete, runnable Python script for the following request.
Concept: {concept}
Scene Description: {scene_description}
Visual Elements: {", ".join(visual_elements)}
STRICT CODE REQUIREMENTS:
1. Header: MUST start with `from manim import *`
2. Class Structure: Define a class inheriting from `MovingCameraScene` (e.g., `class GenScene(MovingCameraScene):`) - this enables camera operations like zoom/pan via `self.camera.frame`
3. Method: All logic must be inside the `def construct(self):` method
4. SYNTAX: Ensure ALL parentheses, brackets, and function calls are properly closed
5. Colors: Use ONLY these valid Manim color constants:
- Basic: WHITE, BLACK, GRAY, GREY, LIGHT_GRAY, DARK_GRAY
- Primary: RED, GREEN, BLUE, YELLOW, ORANGE, PINK, PURPLE, TEAL, GOLD, MAROON
- Variants: RED_A, RED_B, RED_C, RED_D, RED_E, GREEN_A, GREEN_B, GREEN_C, GREEN_D, GREEN_E,
BLUE_A, BLUE_B, BLUE_C, BLUE_D, BLUE_E, YELLOW_A, YELLOW_B, YELLOW_C, YELLOW_D, YELLOW_E
- NEVER use: DARK_GREEN, LIGHT_GREEN, DARK_BLUE, LIGHT_BLUE, DARK_RED, LIGHT_RED (these don't exist!)
6. Text: Use `Text()` objects for strings. Avoid `Tex()` or `MathTex()` unless necessary
7. Positioning: Use `.next_to()`, `.move_to()`, or `.shift()` to arrange elements
8. Animations: Use ONLY these valid animations:
- Write(), Create(), FadeIn(), FadeOut(), GrowFromCenter(), ShrinkToCenter()
- Transform(), ReplacementTransform(), MoveToTarget(), ApplyMethod()
- Rotate(), Indicate(), Flash(), ShowCreation() - DO NOT use lowercase like 'flash'
- For custom effects use .animate.method() (e.g., obj.animate.scale(2), obj.animate.shift(UP))
9. Pacing: Include `self.wait(1)` between major animation groups
OUTPUT FORMAT:
Provide ONLY the raw Python code. Do not wrap in markdown blocks (no ```python). Do not include conversational text.
"""
response = await hf_wrapper.text_generation(
model=selected_model,
prompt=prompt,
max_new_tokens=2048,
temperature=0.3,
)
logger.info(f"Successfully generated Manim code for concept: {concept}")
return CallToolResult(
content=[
TextContent(
type="text",
text=f"Generated Manim Code:\n\n```python\n{response}\n```",
)
]
)
except Exception as e:
logger.error(f"Code generation failed: {str(e)}")
return CallToolResult(
content=[
TextContent(type="text", text=f"Code generation failed: {str(e)}")
],
isError=True,
)
async def refine_animation(
hf_wrapper: HFInferenceWrapper, arguments: Dict[str, Any]
) -> CallToolResult:
"""
Refine animation code based on feedback.
Uses a code LLM to improve existing Manim code based on:
- User feedback or error messages
- Specific improvement goals
- Visual or educational quality issues
Args:
hf_wrapper: HuggingFace inference wrapper instance
arguments: Dictionary containing:
- original_code (str): The original Manim code to refine
- feedback (str): Feedback or error message about the code
- improvement_goals (list, optional): List of specific improvement goals
- model (str, optional): Hugging Face model to use
Returns:
CallToolResult with the refined Manim code
"""
original_code = arguments["original_code"]
feedback = arguments["feedback"]
improvement_goals = arguments.get("improvement_goals", [])
model = arguments.get("model")
try:
model_config = ModelConfig()
selected_model = model or model_config.code_models[0]
prompt = f"""
You are a Manim Code Repair Agent. Your task is to rewrite the FULL Python script to fix issues or apply improvements.
Previous Code:
{original_code}
User Feedback/Error:
{feedback}
Improvement Goals:
{", ".join(improvement_goals)}
INSTRUCTIONS:
1. Output the COMPLETE corrected script, including `from manim import *`.
2. Do not output diffs or partial snippets.
3. Ensure the class inherits from `MovingCameraScene` and uses `def construct(self):`.
4. Fix logic errors based on the feedback.
5. Animations: Use ONLY valid animations like Write(), FadeIn(), FadeOut(), Create(), Flash(), Transform() - NEVER lowercase!
6. Colors: Use ONLY these valid Manim color constants:
- Basic: WHITE, BLACK, GRAY, GREY, LIGHT_GRAY, DARK_GRAY
- Primary: RED, GREEN, BLUE, YELLOW, ORANGE, PINK, PURPLE, TEAL, GOLD, MAROON
- Variants: RED_A, RED_B, RED_C, RED_D, RED_E, GREEN_A, GREEN_B, GREEN_C, GREEN_D, GREEN_E,
BLUE_A, BLUE_B, BLUE_C, BLUE_D, BLUE_E, YELLOW_A, YELLOW_B, YELLOW_C, YELLOW_D, YELLOW_E
- NEVER use: DARK_GREEN, LIGHT_GREEN, DARK_BLUE, LIGHT_BLUE, DARK_RED, LIGHT_RED (these don't exist!)
- For darker/lighter variants, use the letter suffixes (e.g., GREEN_E for dark green, GREEN_A for light green).
OUTPUT:
Return ONLY the raw Python code. No markdown backticks. No explanation.
"""
response = await hf_wrapper.text_generation(
model=selected_model,
prompt=prompt,
max_new_tokens=2048,
temperature=0.3,
)
logger.info("Successfully refined animation code")
return CallToolResult(
content=[
TextContent(
type="text",
text=f"Refined Manim Code:\n\n```python\n{response}\n```",
)
]
)
except Exception as e:
logger.error(f"Code refinement failed: {str(e)}")
return CallToolResult(
content=[
TextContent(type="text", text=f"Code refinement failed: {str(e)}")
],
isError=True,
)
|