Spaces:
Sleeping
Sleeping
File size: 3,899 Bytes
6fc3143 |
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 |
"""
3D Animation Code Generation
Generates Manim code for 3D animations using LLM.
"""
import os
import re
from typing import Optional
from fastapi import HTTPException
import litellm
from .prompts_3d import get_3d_system_prompt, get_3d_examples, SYSTEM_PROMPT_3D
from manimator.utils.code_postprocessor import post_process_code
def generate_3d_animation_response(
user_prompt: str,
model: Optional[str] = None,
include_examples: bool = True
) -> str:
"""
Generate 3D Manim animation code from a text prompt.
Args:
user_prompt: User's request for a 3D animation
model: LiteLLM model to use (defaults to CODE_GEN_MODEL env var)
include_examples: Whether to include example code in the system prompt
Returns:
Generated 3D Manim animation code (post-processed)
Raises:
HTTPException: If code generation fails
"""
try:
# Build system prompt
system_prompt = get_3d_system_prompt()
if include_examples:
examples = get_3d_examples()
examples_text = "\n\n".join([
f"## {name.upper()} EXAMPLE\n{code}"
for name, code in examples.items()
])
system_prompt += f"\n\nHERE ARE SOME EXAMPLES:\n\n{examples_text}"
# Build messages
messages = [
{"role": "system", "content": system_prompt},
{"role": "user", "content": f"""Create a complete, runnable Manim 3D animation for the following request:
{user_prompt}
Requirements:
1. Use ThreeDScene as base class
2. Include voiceover narration
3. Set appropriate camera angles
4. Make it educational and visually appealing
5. Include mathematical equations if relevant
6. Use smooth animations and camera movements
7. Aim for 2-5 minutes duration (or as specified)
Generate ONLY the Python code, nothing else."""}
]
# Use specified model or default
if model is None:
model = os.getenv("CODE_GEN_MODEL")
# Generate code
response = litellm.completion(
model=model,
messages=messages,
num_retries=2
)
raw_code = response.choices[0].message.content
# Extract code if wrapped in markdown
pattern = r'```python\n(.*?)```'
match = re.search(pattern, raw_code, re.DOTALL)
if match:
code = match.group(1)
else:
code = raw_code
# Post-process the code to fix common issues
processed_code = post_process_code(code)
return processed_code
except Exception as e:
raise HTTPException(
status_code=500,
detail=f"Failed to generate 3D animation: {str(e)}"
)
def generate_3d_animation_with_category(
user_prompt: str,
category: str,
model: Optional[str] = None
) -> str:
"""
Generate 3D animation code with a specific STEM category focus.
Args:
user_prompt: User's description
category: One of 'mathematical', 'scientific', 'geometric', 'data'
model: LLM model to use
Returns:
Generated Python code
"""
category_prompts = {
"mathematical": "Focus on mathematical accuracy and clear visualization of mathematical concepts.",
"scientific": "Focus on scientific accuracy, realistic models, and clear explanations.",
"geometric": "Focus on geometric properties, transformations, and spatial relationships.",
"data": "Focus on clear data representation, appropriate chart types, and data-driven insights."
}
enhanced_prompt = f"""{user_prompt}
CATEGORY: {category.upper()}
{category_prompts.get(category, '')}
"""
return generate_3d_animation_response(enhanced_prompt, model=model)
|