|
|
"""Prompt management and code generation utilities.""" |
|
|
|
|
|
import logging |
|
|
from typing import Dict, Any, Optional |
|
|
from groq import Groq |
|
|
|
|
|
from config.constants import ( |
|
|
AnimationCategory, |
|
|
ANIMATION_CONSTANTS, |
|
|
ErrorType |
|
|
) |
|
|
from utils.error_handler import ErrorHandler |
|
|
|
|
|
logger = logging.getLogger(__name__) |
|
|
|
|
|
class PromptManager: |
|
|
"""Manages prompt enhancement and code generation.""" |
|
|
|
|
|
def __init__(self): |
|
|
"""Initialize the prompt manager.""" |
|
|
self.category_contexts = { |
|
|
AnimationCategory.ALGEBRA: """ |
|
|
Focus on algebraic concepts like: |
|
|
- Equations and inequalities |
|
|
- Function transformations |
|
|
- Polynomial operations |
|
|
- Algebraic proofs |
|
|
Use colors to distinguish different terms and operations. |
|
|
""", |
|
|
AnimationCategory.GEOMETRY: """ |
|
|
Focus on geometric concepts like: |
|
|
- Shapes and their properties |
|
|
- Constructions and transformations |
|
|
- Theorems and proofs |
|
|
- Spatial relationships |
|
|
Use clear lines and shapes with appropriate colors. |
|
|
""", |
|
|
AnimationCategory.CALCULUS: """ |
|
|
Focus on calculus concepts like: |
|
|
- Derivatives and integrals |
|
|
- Limits and continuity |
|
|
- Series and sequences |
|
|
- Applications of calculus |
|
|
Use smooth curves and clear visualizations of rates of change. |
|
|
""", |
|
|
AnimationCategory.STATISTICS: """ |
|
|
Focus on statistical concepts like: |
|
|
- Probability distributions |
|
|
- Statistical measures |
|
|
- Data visualization |
|
|
- Statistical tests |
|
|
Use appropriate graphs and charts with clear labels. |
|
|
""", |
|
|
AnimationCategory.LINEAR_ALGEBRA: """ |
|
|
Focus on linear algebra concepts like: |
|
|
- Matrices and vectors |
|
|
- Linear transformations |
|
|
- Eigenvalues and eigenvectors |
|
|
- Vector spaces |
|
|
Use clear geometric interpretations with appropriate colors. |
|
|
""", |
|
|
AnimationCategory.NUMBER_THEORY: """ |
|
|
Focus on number theory concepts like: |
|
|
- Prime numbers and factorization |
|
|
- Modular arithmetic |
|
|
- Number sequences |
|
|
- Mathematical proofs |
|
|
Use clear visualizations of number patterns and relationships. |
|
|
""", |
|
|
AnimationCategory.OTHER: """ |
|
|
Focus on creating clear and educational animations that: |
|
|
- Explain mathematical concepts visually |
|
|
- Use appropriate colors and shapes |
|
|
- Include clear labels and annotations |
|
|
- Maintain a logical flow of information |
|
|
""" |
|
|
} |
|
|
|
|
|
def detect_category(self, prompt: str) -> AnimationCategory: |
|
|
"""Detect the mathematical category of a prompt.""" |
|
|
prompt_lower = prompt.lower() |
|
|
|
|
|
|
|
|
if any(word in prompt_lower for word in ["equation", "function", "polynomial", "algebra"]): |
|
|
return AnimationCategory.ALGEBRA |
|
|
elif any(word in prompt_lower for word in ["shape", "geometry", "triangle", "circle", "square"]): |
|
|
return AnimationCategory.GEOMETRY |
|
|
elif any(word in prompt_lower for word in ["derivative", "integral", "limit", "calculus"]): |
|
|
return AnimationCategory.CALCULUS |
|
|
elif any(word in prompt_lower for word in ["statistic", "probability", "distribution", "data"]): |
|
|
return AnimationCategory.STATISTICS |
|
|
elif any(word in prompt_lower for word in ["matrix", "vector", "linear", "eigen"]): |
|
|
return AnimationCategory.LINEAR_ALGEBRA |
|
|
elif any(word in prompt_lower for word in ["prime", "number", "sequence", "modular"]): |
|
|
return AnimationCategory.NUMBER_THEORY |
|
|
else: |
|
|
return AnimationCategory.OTHER |
|
|
|
|
|
def get_enhancement_context(self, prompt: str) -> str: |
|
|
"""Get category-specific context for prompt enhancement.""" |
|
|
category = self.detect_category(prompt) |
|
|
return self.category_contexts.get(category, self.category_contexts[AnimationCategory.OTHER]) |
|
|
|
|
|
def enhance_prompt( |
|
|
self, |
|
|
user_prompt: str, |
|
|
client: Groq, |
|
|
model: str, |
|
|
temperature: float = 0.7, |
|
|
max_tokens: int = 500 |
|
|
) -> str: |
|
|
"""Enhance the user prompt for better Manim code generation.""" |
|
|
category_context = self.get_enhancement_context(user_prompt) |
|
|
|
|
|
enhancement_prompt = f""" |
|
|
You are an expert at creating detailed prompts for Manim (Mathematical Animation Engine) video generation. |
|
|
|
|
|
User's original prompt: "{user_prompt}" |
|
|
|
|
|
Context: {category_context} |
|
|
|
|
|
Please enhance this prompt by: |
|
|
1. Adding specific mathematical or visual details |
|
|
2. Suggesting appropriate Manim objects and animations |
|
|
3. Specifying colors, positioning, and timing |
|
|
4. Including educational context if applicable |
|
|
5. Making it clear and comprehensive for code generation |
|
|
|
|
|
Enhanced prompt should be detailed but concise, focusing on visual elements that can be animated with Manim. |
|
|
|
|
|
Return only the enhanced prompt, nothing else. |
|
|
""" |
|
|
|
|
|
try: |
|
|
response = client.chat.completions.create( |
|
|
messages=[{"role": "user", "content": enhancement_prompt}], |
|
|
model=model, |
|
|
temperature=temperature, |
|
|
max_tokens=max_tokens |
|
|
) |
|
|
return response.choices[0].message.content.strip() |
|
|
except Exception as e: |
|
|
logger.error(f"Error enhancing prompt: {e}") |
|
|
ErrorHandler.handle_error(ErrorType.PROMPT_ERROR, str(e)) |
|
|
return user_prompt |
|
|
|
|
|
def generate_code( |
|
|
self, |
|
|
enhanced_prompt: str, |
|
|
client: Groq, |
|
|
model: str, |
|
|
temperature: float = 0.3, |
|
|
max_tokens: int = 2000 |
|
|
) -> str: |
|
|
"""Generate Manim code from the enhanced prompt.""" |
|
|
code_generation_prompt = f""" |
|
|
You are an expert Manim developer. Generate a complete, working Manim scene based on this prompt: |
|
|
|
|
|
"{enhanced_prompt}" |
|
|
|
|
|
Requirements: |
|
|
1. Create a class that inherits from Scene |
|
|
2. Use proper Manim syntax and imports |
|
|
3. Include appropriate animations (Create, Transform, Write, etc.) |
|
|
4. Add colors, positioning, and timing |
|
|
5. Make it visually appealing and educational |
|
|
6. Ensure the code is syntactically correct |
|
|
7. Use self.play() for animations and self.wait() for pauses |
|
|
8. Include comments explaining key parts |
|
|
9. Keep animations under {ANIMATION_CONSTANTS['MAX_DURATION']} seconds total duration |
|
|
10. Use appropriate wait times between animations (between {ANIMATION_CONSTANTS['MIN_WAIT_TIME']} and {ANIMATION_CONSTANTS['MAX_WAIT_TIME']} seconds) |
|
|
11. To set background color, use: self.camera.background_color = WHITE |
|
|
12. For SVG paths, use SVGMobject instead of VMobjectFromSVGPath |
|
|
13. Always use proper animation classes (Create, Write, Transform, etc.) with self.play() |
|
|
14. For positioning, use methods like next_to(), to_edge(), to_corner() instead of direct position manipulation |
|
|
15. Do NOT use the font_size argument with get_x_axis_label or get_y_axis_label. Only use supported arguments for these methods. |
|
|
|
|
|
Structure: |
|
|
```python |
|
|
from manim import * |
|
|
|
|
|
class GeneratedScene(Scene): |
|
|
def construct(self): |
|
|
# Set background color |
|
|
self.camera.background_color = WHITE |
|
|
|
|
|
# Create objects |
|
|
obj = Circle(color=BLUE) |
|
|
|
|
|
# Position objects |
|
|
obj.to_edge(UP) |
|
|
|
|
|
# Animate objects |
|
|
self.play(Create(obj)) |
|
|
self.wait(1) |
|
|
|
|
|
# Transform objects |
|
|
new_obj = Square(color=RED) |
|
|
self.play(Transform(obj, new_obj)) |
|
|
self.wait(1) |
|
|
``` |
|
|
|
|
|
Generate ONLY the Python code, no explanations or markdown formatting. |
|
|
""" |
|
|
|
|
|
try: |
|
|
response = client.chat.completions.create( |
|
|
messages=[{"role": "user", "content": code_generation_prompt}], |
|
|
model=model, |
|
|
temperature=temperature, |
|
|
max_tokens=max_tokens |
|
|
) |
|
|
return response.choices[0].message.content.strip() |
|
|
except Exception as e: |
|
|
logger.error(f"Error generating Manim code: {e}") |
|
|
ErrorHandler.handle_error(ErrorType.CODE_GENERATION_ERROR, str(e)) |
|
|
raise |
