""" Simplified example of a Manim animation generator using pydantic-ai. """ import os from typing import List, Optional from dotenv import load_dotenv from pydantic_ai.models.openai import OpenAIModel from pydantic_ai.providers.openai import OpenAIProvider from pydantic_ai import Agent, RunContext from pydantic import BaseModel, Field from datetime import datetime import openai import tempfile import subprocess import shutil import time import logging # Load environment variables load_dotenv() # Configure logging if not already done logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class AnimationPrompt(BaseModel): """User input for animation generation.""" description: str = Field(..., description="Description of the mathematical concept to animate") complexity: str = Field("medium", description="Desired complexity of the animation (simple, medium, complex)") class AnimationOutput(BaseModel): """Output of the animation generation.""" manim_code: str = Field(..., description="Generated Manim code") explanation: str = Field(..., description="Explanation of the animation") # Create the animation agent with basic static system prompt model = OpenAIModel( 'deepseek-ai/DeepSeek-V3', provider=OpenAIProvider( base_url='https://api.together.xyz/v1', api_key=os.environ.get('TOGETHER_API_KEY') ), ) animation_agent = Agent( model, deps_type=AnimationPrompt, system_prompt=( "You are a mathematical animation specialist. Your job is to convert text descriptions " "into Manim code that visualizes mathematical concepts. Provide clear and accurate code." ) ) # Configure OpenAI client to use Together API client = openai.OpenAI( api_key=os.environ.get("TOGETHER_API_KEY"), base_url="https://api.together.xyz/v1", ) # Add dynamic system prompts @animation_agent.system_prompt def add_complexity_guidance(ctx: RunContext[AnimationPrompt]) -> str: """Add guidance based on requested complexity.""" complexity = ctx.deps.complexity if complexity == "simple": return "Generate simple, beginner-friendly Manim code with minimal elements and clear explanations." elif complexity == "complex": return "Generate advanced Manim code with sophisticated animations and detailed mathematical representations." else: # medium return "Generate standard Manim code that balances simplicity and detail to effectively demonstrate the concept." @animation_agent.system_prompt def add_timestamp() -> str: """Add a timestamp to help with freshness of information.""" return f"Current timestamp: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" @animation_agent.tool def generate_manim_code(ctx: RunContext[AnimationPrompt]) -> str: """Generate Manim code based on the user's description.""" prompt = ctx.deps # Use Together API with OpenAI client response = client.chat.completions.create( model="meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo", messages=[ {"role": "system", "content": """ Generate Manim code for mathematical animations. The code MUST: 1. Be fully compilable without errors using Manim Community edition 2. Use only the Scene class with a class name 'ManimScene' exactly 3. Include 'from manim import *' at the top 4. Implement the construct method only 5. Use only standard Manim objects and methods 6. Include proper self.play() and self.wait() calls 7. Use valid LaTeX syntax for any mathematical expressions 8. Avoid experimental or uncommon Manim features 9. Keep the animation clean, concise, and educational 10. Include proper error handling for all mathematical operations 11. DO NOT include any backticks (```) or markdown formatting in your response RESPOND WITH CODE ONLY, NO EXPLANATIONS OUTSIDE OF CODE COMMENTS, NO MARKDOWN FORMATTING. """ }, {"role": "user", "content": f"Create Manim code for a {prompt.complexity} animation of {prompt.description}"} ] ) generated_code = response.choices[0].message.content # Strip markdown formatting if it appears in the response if "```python" in generated_code: generated_code = generated_code.split("```python")[1] if "```" in generated_code: generated_code = generated_code.split("```")[0] return generated_code @animation_agent.tool def explain_animation(ctx: RunContext[AnimationPrompt], code: str) -> str: """Explain the generated animation in plain language.""" prompt = ctx.deps # Use Together API with OpenAI client response = client.chat.completions.create( model="meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo", messages=[ {"role": "system", "content": "Explain mathematical animations in simple terms."}, {"role": "user", "content": f"Explain this Manim animation of {prompt.description} " + f"with complexity {prompt.complexity} in simple terms:\n{code}"} ] ) return response.choices[0].message.content def render_manim_video(code, quality="medium_quality"): try: temp_dir = tempfile.mkdtemp() script_path = os.path.join(temp_dir, "manim_script.py") with open(script_path, "w") as f: f.write(code) class_name = None for line in code.split("\n"): if line.startswith("class ") and "Scene" in line: class_name = line.split("class ")[1].split("(")[0].strip() break if not class_name: return "Error: Could not identify the Scene class in the generated code." if quality == "high_quality": command = ["manim", "-qh", script_path, class_name] quality_dir = "1080p60" elif quality == "low_quality": command = ["manim", "-ql", script_path, class_name] quality_dir = "480p15" else: command = ["manim", "-qm", script_path, class_name] quality_dir = "720p30" logger.info(f"Executing command: {' '.join(command)}") result = subprocess.run(command, cwd=temp_dir, capture_output=True, text=True) logger.info(f"Manim stdout: {result.stdout}") logger.error(f"Manim stderr: {result.stderr}") if result.returncode != 0: logger.error(f"Manim execution failed: {result.stderr}") return f"Error rendering video: {result.stderr}" media_dir = os.path.join(temp_dir, "media") videos_dir = os.path.join(media_dir, "videos") if not os.path.exists(videos_dir): return "Error: No video was generated. Check if Manim is installed correctly." scene_dirs = [d for d in os.listdir(videos_dir) if os.path.isdir(os.path.join(videos_dir, d))] if not scene_dirs: return "Error: No scene directory found in the output." scene_dir = max([os.path.join(videos_dir, d) for d in scene_dirs], key=os.path.getctime) mp4_files = [f for f in os.listdir(os.path.join(scene_dir, quality_dir)) if f.endswith(".mp4")] if not mp4_files: return "Error: No MP4 file was generated." video_file = max([os.path.join(scene_dir, quality_dir, f) for f in mp4_files], key=os.path.getctime) output_dir = os.path.join(os.getcwd(), "generated_videos") os.makedirs(output_dir, exist_ok=True) timestamp = int(time.time()) output_file = os.path.join(output_dir, f"manim_video_{timestamp}.mp4") shutil.copy2(video_file, output_file) logger.info(f"Video generated: {output_file}") return output_file except Exception as e: logger.error(f"Error rendering video: {e}") return f"Error rendering video: {str(e)}" finally: if 'temp_dir' in locals(): try: shutil.rmtree(temp_dir) except Exception as e: logger.error(f"Error cleaning up temporary directory: {e}") def run_animation_agent(description: str, complexity: str = "medium", quality: str = "medium_quality") -> AnimationOutput: """Run the animation agent to generate code and explanation.""" prompt = AnimationPrompt(description=description, complexity=complexity) # Use the agent to process the request result = animation_agent.run_sync( "Generate Manim code for this animation and explain what it does", deps=prompt ) # Generate code and explanation code = None explanation = None # As a fallback, provide a direct implementation specific to the Pythagorean theorem if "pythagorean theorem" in description.lower(): code = f""" from manim import * class ManimScene(Scene): def construct(self): # Animation for: {prompt.description} # Complexity level: {prompt.complexity} # Create a right triangle triangle = Polygon( ORIGIN, RIGHT * 3, UP * 4, color=WHITE ) # Labels for sides a_label = MathTex("a").next_to(triangle, DOWN) b_label = MathTex("b").next_to(triangle, RIGHT) c_label = MathTex("c").next_to(triangle.get_center(), UP + LEFT) # The equation equation = MathTex("a^2 + b^2 = c^2").to_edge(DOWN) # Display the triangle and labels self.play(Create(triangle)) self.play(Write(a_label), Write(b_label), Write(c_label)) self.wait() # Show the equation self.play(Write(equation)) self.wait() """ explanation = ( f"This animation visualizes {prompt.description} with a {prompt.complexity} " f"complexity level. It creates a right triangle and labels its sides a, b, and c. " f"It then displays the Pythagorean theorem equation a² + b² = c²." ) else: # Generic fallback code = f""" from manim import * class ManimScene(Scene): def construct(self): # Animation for: {prompt.description} # Complexity level: {prompt.complexity} # Title title = Text("{description}") self.play(Write(title)) self.wait() self.play(title.animate.to_edge(UP)) # Main content based on complexity if "{complexity}" == "simple": # Simple visualization circle = Circle() self.play(Create(circle)) self.wait() else: # More complex visualization axes = Axes( x_range=[-3, 3], y_range=[-3, 3], axis_config={"color": BLUE} ) self.play(Create(axes)) # Add a function graph graph = axes.plot(lambda x: x**2, color=YELLOW) self.play(Create(graph)) self.wait() """ explanation = ( f"This animation visualizes {prompt.description} with a {prompt.complexity} " f"complexity level. It displays a title and creates a visualization that matches " f"the requested complexity." ) # Try to render the video if code: video_path = render_manim_video(code, quality) if video_path and not video_path.startswith("Error"): print(f"Video rendered successfully at: {video_path}") return AnimationOutput(manim_code=code, explanation=explanation) if __name__ == "__main__": # Example usage result = run_animation_agent( "the Pythagorean theorem showing how a² + b² = c²", complexity="simple" ) print("=== Generated Manim Code ===") print(result.manim_code) print("\n=== Explanation ===") print(result.explanation)