Spaces:
Sleeping
Sleeping
| """ | |
| 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) | |