""" NeuroAnim Orchestrator This script coordinates the entire STEM animation generation pipeline: 1. Concept Planning 2. Code Generation 3. Rendering 4. Vision-based Analysis 5. Audio Generation 6. Final Merging It uses the MCP servers (renderer and creative) to accomplish these tasks. """ import ast import asyncio import json import logging import os import tempfile from pathlib import Path from typing import Any, Dict, List, Optional import aiofiles from dotenv import load_dotenv from mcp import ClientSession, StdioServerParameters from mcp.client.stdio import stdio_client from utils.tts import TTSGenerator load_dotenv() # Set up logging logging.basicConfig( level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s" ) logger = logging.getLogger(__name__) class NeuroAnimOrchestrator: """Main orchestrator for NeuroAnim pipeline.""" def __init__( self, hf_api_key: Optional[str] = None, elevenlabs_api_key: Optional[str] = None ): self.hf_api_key = hf_api_key or os.getenv("HUGGINGFACE_API_KEY") self.elevenlabs_api_key = elevenlabs_api_key or os.getenv("ELEVENLABS_API_KEY") self.renderer_session: Optional[ClientSession] = None self.creative_session: Optional[ClientSession] = None # Initialize TTS generator self.tts_generator = TTSGenerator( elevenlabs_api_key=self.elevenlabs_api_key, hf_api_key=self.hf_api_key, fallback_enabled=True, ) # Context managers for MCP client connections self._renderer_cm = None self._creative_cm = None self._renderer_streams = None self._creative_streams = None # Working directories self.work_dir: Optional[Path] = None self.output_dir: Optional[Path] = None async def initialize(self): """Initialize MCP server connections.""" # Set up working directories self.work_dir = Path(tempfile.mkdtemp(prefix="neuroanim_work_")) self.output_dir = Path("outputs") self.output_dir.mkdir(exist_ok=True) logger.info(f"Working directory: {self.work_dir}") logger.info(f"Output directory: {self.output_dir}") # Initialize renderer server # stdio_client is an async context manager, must use async with renderer_params = StdioServerParameters( command="python", args=["mcp_servers/renderer.py"] ) self._renderer_cm = stdio_client(renderer_params) self._renderer_streams = await self._renderer_cm.__aenter__() read_stream, write_stream = self._renderer_streams self.renderer_session = ClientSession(read_stream, write_stream) # Start background receive loop for the client session await self.renderer_session.__aenter__() await self.renderer_session.initialize() logger.info("Renderer MCP server connected") # Initialize creative server creative_params = StdioServerParameters( command="python", args=["mcp_servers/creative.py"], env={"HUGGINGFACE_API_KEY": self.hf_api_key} if self.hf_api_key else None, ) self._creative_cm = stdio_client(creative_params) self._creative_streams = await self._creative_cm.__aenter__() read_stream, write_stream = self._creative_streams self.creative_session = ClientSession(read_stream, write_stream) # Start background receive loop for the client session await self.creative_session.__aenter__() await self.creative_session.initialize() logger.info("Creative MCP server connected") async def cleanup(self): """Clean up resources.""" import shutil # Close sessions first if self.renderer_session: try: await self.renderer_session.__aexit__(None, None, None) except (Exception, asyncio.CancelledError) as e: logger.debug(f"Error closing renderer session: {e}") if self.creative_session: try: await self.creative_session.__aexit__(None, None, None) except (Exception, asyncio.CancelledError) as e: logger.debug(f"Error closing creative session: {e}") # Then close the stdio_client context managers with timeout if self._renderer_cm: try: async with asyncio.timeout(2): # 2 second timeout await self._renderer_cm.__aexit__(None, None, None) except (Exception, asyncio.CancelledError, TimeoutError) as e: logger.debug(f"Error closing renderer context manager: {e}") if self._creative_cm: try: async with asyncio.timeout(2): # 2 second timeout await self._creative_cm.__aexit__(None, None, None) except (Exception, asyncio.CancelledError, TimeoutError) as e: logger.debug(f"Error closing creative context manager: {e}") # Clean up working directory if self.work_dir and self.work_dir.exists(): try: shutil.rmtree(self.work_dir) logger.info(f"Cleaned up working directory: {self.work_dir}") except Exception as e: logger.warning(f"Failed to clean up working directory: {e}") async def call_tool( self, session: ClientSession, tool_name: str, arguments: Dict[str, Any] ) -> Dict[str, Any]: """Call a tool on an MCP server.""" result = await session.call_tool(tool_name, arguments) if hasattr(result, "content") and result.content: content = result.content[0] if hasattr(content, "text"): return { "text": content.text, "isError": getattr(result, "isError", False), } return {"text": str(result), "isError": False} async def generate_animation( self, topic: str, target_audience: str = "general", animation_length_minutes: float = 2.0, output_filename: str = "animation.mp4", quality: str = "medium", progress_callback: Optional[callable] = None, ) -> Dict[str, Any]: """Complete animation generation pipeline.""" try: logger.info(f"Starting animation generation for: {topic}") # Step 1: Concept Planning logger.info("Step 1: Planning concept...") if progress_callback: progress_callback("Planning concept", 0.1) concept_result = await self.call_tool( self.creative_session, "plan_concept", { "topic": topic, "target_audience": target_audience, "animation_length_minutes": animation_length_minutes, }, ) if concept_result["isError"]: raise Exception(f"Concept planning failed: {concept_result['text']}") concept_plan = concept_result["text"] logger.info("Concept planning completed") # Step 2: Generate Narration logger.info("Step 2: Generating narration...") if progress_callback: progress_callback("Generating narration script", 0.25) narration_result = await self.call_tool( self.creative_session, "generate_narration", { "concept": topic, "scene_description": concept_plan, "target_audience": target_audience, "duration_seconds": int(animation_length_minutes * 60), }, ) if narration_result["isError"]: raise Exception( f"Narration generation failed: {narration_result['text']}" ) # Clean narration text - remove title/prefix before TTS narration_text = self._clean_narration_text(narration_result["text"]) logger.info("Narration generation completed") logger.info(f"Narration preview: {narration_text[:100]}...") # Step 3: Generate Manim Code with retry logic logger.info("Step 3: Generating Manim code...") if progress_callback: progress_callback("Creating Manim animation code", 0.40) target_duration_seconds = int(animation_length_minutes * 60) manim_code = await self._generate_and_validate_code( topic=topic, concept_plan=concept_plan, duration_seconds=target_duration_seconds, max_retries=3, ) logger.info("Manim code generation completed and validated") # Step 4: Write Manim File logger.info("Step 4: Writing Manim file...") manim_file = self.work_dir / "animation.py" write_result = await self.call_tool( self.renderer_session, "write_manim_file", {"filepath": str(manim_file), "code": manim_code}, ) if write_result["isError"]: raise Exception(f"File writing failed: {write_result['text']}") # Extract scene name from code scene_name = self._extract_scene_name(manim_code) logger.info(f"Scene name detected: {scene_name}") # Step 5: Render Animation with retry on runtime errors logger.info("Step 5: Rendering animation...") if progress_callback: progress_callback("Rendering animation video", 0.55) max_render_retries = 5 video_file = None for render_attempt in range(max_render_retries): render_result = await self.call_tool( self.renderer_session, "render_manim_animation", { "scene_name": scene_name, "file_path": str(manim_file), "output_dir": str(self.work_dir), "quality": quality, # Use the quality parameter "format": "mp4", "frame_rate": 30, }, ) if not render_result["isError"]: # Success! Find the rendered file video_file = self._find_output_file(self.work_dir, scene_name, "mp4") if video_file: # Check video duration try: actual_duration = self._get_video_duration(video_file) logger.info(f"Rendered video duration: {actual_duration:.2f}s (Target: {target_duration_seconds}s)") if actual_duration < target_duration_seconds * 0.5: logger.warning(f"Video is too short ({actual_duration:.2f}s < {target_duration_seconds * 0.5}s). Forcing retry...") error_text = ( f"The generated animation was TOO SHORT ({actual_duration:.1f}s). " f"The target duration is {target_duration_seconds}s. " "You MUST make the animation longer by adding more `self.wait()` calls " "and ensuring animations play slower (use run_time parameter)." ) # Fall through to error handling logic below else: break except Exception as e: logger.warning(f"Could not verify video duration: {e}") break else: logger.warning("Render succeeded but could not find output file") if render_attempt < max_render_retries - 1: continue # Rendering failed - check if it's a runtime error we can fix error_text = render_result["text"] logger.warning(f"Render attempt {render_attempt + 1} failed: {error_text[:200]}...") # Check if this is a Manim runtime error (not a "no scene" error) if render_attempt < max_render_retries - 1 and ( "TypeError" in error_text or "AttributeError" in error_text or "ValueError" in error_text or "KeyError" in error_text ): logger.info(f"Detected runtime error in Manim code. Regenerating code (attempt {render_attempt + 2}/{max_render_retries})...") # Regenerate code with error feedback runtime_error_msg = f"Runtime Error during Manim rendering:\n{error_text}\n\nPlease fix the code to be compatible with Manim version 0.19.0." manim_code = await self._generate_and_validate_code( topic=topic, concept_plan=concept_plan, duration_seconds=target_duration_seconds, max_retries=3, # Allow retries for syntax errors during fix previous_error=runtime_error_msg, previous_code=manim_code, ) # Write the new code write_result = await self.call_tool( self.renderer_session, "write_manim_file", {"filepath": str(manim_file), "code": manim_code}, ) if write_result["isError"]: raise Exception(f"File writing failed: {write_result['text']}") # Extract scene name from new code scene_name = self._extract_scene_name(manim_code) logger.info(f"Regenerated code with scene: {scene_name}") # Loop will retry rendering with new code continue else: # Not a runtime error or out of retries raise Exception(f"Rendering failed: {error_text}") if not video_file: raise Exception("Could not find rendered video file after all attempts") logger.info(f"Animation rendered: {video_file}") # Step 6: Generate Speech Audio logger.info("Step 6: Generating speech audio...") if progress_callback: progress_callback("Generating audio narration", 0.75) audio_file = self.work_dir / "narration.mp3" # Use TTS generator with automatic fallback try: tts_result = await self.tts_generator.generate_speech( text=narration_text, output_path=audio_file, voice="rachel" ) logger.info( f"Audio generated with {tts_result['provider']}: {audio_file}" ) # Validate audio file validation = self.tts_generator.validate_audio_file(audio_file) if not validation["valid"]: logger.warning( f"Audio validation warning: {validation.get('error', 'Unknown issue')}" ) logger.info("Audio file may have issues but continuing...") else: logger.info( f"Audio validated: {validation.get('duration', 'N/A')}s, {validation.get('size', 0)} bytes" ) except Exception as e: logger.error(f"TTS generation failed: {e}") raise Exception(f"Speech generation failed: {str(e)}") # Step 7: Merge Video and Audio logger.info("Step 7: Merging video and audio...") if progress_callback: progress_callback("Merging video and audio", 0.90) final_output = self.output_dir / output_filename merge_result = await self.call_tool( self.renderer_session, "merge_video_audio", { "video_file": str(video_file), "audio_file": str(audio_file), "output_file": str(final_output), }, ) if merge_result["isError"]: raise Exception(f"Merging failed: {merge_result['text']}") # Step 8: Generate Quiz logger.info("Step 8: Generating quiz...") if progress_callback: progress_callback("Creating quiz questions", 0.95) quiz_result = await self.call_tool( self.creative_session, "generate_quiz", { "concept": topic, "difficulty": "medium", "num_questions": 3, "question_types": ["multiple_choice"], }, ) quiz_content = ( quiz_result["text"] if not quiz_result["isError"] else "Quiz generation failed" ) # Return results results = { "success": True, "topic": topic, "target_audience": target_audience, "concept_plan": concept_plan, "narration": narration_text, "manim_code": manim_code, "output_file": str(final_output), "quiz": quiz_content, "work_dir": str(self.work_dir), } logger.info(f"Animation generation completed successfully: {final_output}") return results except Exception as e: logger.error(f"Animation generation failed: {str(e)}") return { "success": False, "error": str(e), "work_dir": str(self.work_dir) if self.work_dir else None, } def _clean_narration_text(self, text: str) -> str: """ Clean narration text by removing title prefixes and formatting artifacts. The creative server returns text with prefixes like "Narration Script:\n\n" which should not be sent to TTS. """ # Remove common prefixes prefixes_to_remove = [ "Narration Script:", "Script:", "Narration:", "Text:", ] cleaned = text.strip() # Remove any of the prefixes (case-insensitive) for prefix in prefixes_to_remove: if cleaned.lower().startswith(prefix.lower()): cleaned = cleaned[len(prefix) :].strip() break # Remove leading newlines and whitespace cleaned = cleaned.lstrip("\n").strip() # Remove any markdown code block markers if cleaned.startswith("```"): lines = cleaned.split("\n") # Remove first line (opening ```) if len(lines) > 1: lines = lines[1:] # Remove last line if it's closing ``` if lines and lines[-1].strip() == "```": lines = lines[:-1] cleaned = "\n".join(lines).strip() return cleaned def _extract_python_code(self, text: str) -> str: """Extract Python code from markdown response.""" # Look for code blocks if "```python" in text: start = text.find("```python") + 9 end = text.find("```", start) if end == -1: end = len(text) return text[start:end].strip() elif "```" in text: start = text.find("```") + 3 end = text.find("```", start) if end == -1: end = len(text) return text[start:end].strip() else: return text.strip() async def _generate_and_validate_code( self, topic: str, concept_plan: str, duration_seconds: int = 60, max_retries: int = 3, previous_error: Optional[str] = None, previous_code: Optional[str] = None, ) -> str: """Generate Manim code with retry logic for syntax errors.""" for attempt in range(max_retries): try: logger.info(f"Code generation attempt {attempt + 1}/{max_retries}") # Build arguments for code generation arguments = { "concept": topic, "scene_description": concept_plan, "visual_elements": ["text", "shapes", "animations"], "duration_seconds": duration_seconds, } # If this is a retry, include error feedback if previous_error: if previous_code: arguments["previous_code"] = previous_code arguments["error_message"] = previous_error logger.info( f"Retrying with error feedback: {previous_error[:100]}..." ) # Generate code code_result = await self.call_tool( self.creative_session, "generate_manim_code", arguments ) if code_result["isError"]: if attempt < max_retries - 1: logger.warning( f"Code generation failed, retrying: {code_result['text']}" ) previous_error = code_result["text"] # Keep previous_code if we had it, for better context in retry continue else: raise Exception( f"Code generation failed: {code_result['text']}" ) # Extract Python code from response manim_code = self._extract_python_code(code_result["text"]) # Validate Python syntax syntax_errors = self._validate_python_syntax(manim_code) if syntax_errors: if attempt < max_retries - 1: logger.warning( f"Syntax error detected, retrying: {syntax_errors}" ) previous_error = f"Syntax Error:\n{syntax_errors}" previous_code = manim_code continue else: raise Exception( f"Generated code has syntax errors after {max_retries} attempts:\n{syntax_errors}" ) # Validate that code contains a Scene class has_scene = self._validate_has_scene_class(manim_code) if not has_scene: if attempt < max_retries - 1: logger.warning( "No Scene class found in generated code, retrying..." ) previous_error = ( "Error: The generated code does not contain any Scene class. " "Please ensure you create a class that inherits from manim.Scene, " "manim.MovingCameraScene, or manim.ThreeDScene." ) previous_code = manim_code continue else: raise Exception( f"Generated code does not contain a Scene class after {max_retries} attempts" ) # Success! logger.info(f"Valid code generated on attempt {attempt + 1}") return manim_code except Exception as e: if attempt < max_retries - 1: logger.warning(f"Attempt {attempt + 1} failed: {str(e)}") previous_error = str(e) continue else: raise raise Exception("Failed to generate valid code after all retries") def _validate_python_syntax(self, code: str) -> Optional[str]: """Validate Python code syntax. Returns error message if invalid, None if valid.""" try: ast.parse(code) return None except SyntaxError as e: # Build detailed error message with context error_msg = f"Line {e.lineno}: {e.msg}" # Show surrounding context (3 lines before and after) if e.lineno is not None: code_lines = code.split("\n") start_line = max(0, e.lineno - 4) # 3 lines before end_line = min(len(code_lines), e.lineno + 2) # 2 lines after error_msg += "\n\nContext:" for i in range(start_line, end_line): line_num = i + 1 prefix = ">>> " if line_num == e.lineno else " " error_msg += f"\n{prefix}{line_num:3d} | {code_lines[i]}" # Add pointer for error line if line_num == e.lineno and e.offset: error_msg += f"\n {' ' * 4}{' ' * (e.offset - 1)}^" return error_msg except Exception as e: return f"Unexpected error during syntax validation: {str(e)}" def _validate_has_scene_class(self, code: str) -> bool: """Check if code contains at least one Scene class.""" import re # Check for Scene class inheritance scene_patterns = [ r"class\s+\w+\s*\(\s*Scene\s*\)", r"class\s+\w+\s*\(\s*MovingCameraScene\s*\)", r"class\s+\w+\s*\(\s*ThreeDScene\s*\)", r"class\s+\w+\s*\(\s*\w*Scene\s*\)", ] for pattern in scene_patterns: if re.search(pattern, code): return True # Also check using AST parsing as a backup try: tree = ast.parse(code) for node in ast.walk(tree): if isinstance(node, ast.ClassDef): # Check if any base class contains "Scene" for base in node.bases: if isinstance(base, ast.Name) and "Scene" in base.id: return True except Exception: pass return False def _extract_scene_name(self, code: str) -> str: """Extract scene class name from Manim code.""" import re # Try multiple patterns to find Scene class patterns = [ r"class\s+(\w+)\s*\(\s*Scene\s*\)", # class Name(Scene) r"class\s+(\w+)\s*\(\s*MovingCameraScene\s*\)", # class Name(MovingCameraScene) r"class\s+(\w+)\s*\(\s*ThreeDScene\s*\)", # class Name(ThreeDScene) r"class\s+(\w+)\s*\(\s*\w*Scene\s*\)", # class Name(AnyScene) ] for pattern in patterns: match = re.search(pattern, code) if match: scene_name = match.group(1) logger.info(f"Found scene class: {scene_name}") return scene_name # If no scene found, look for any class definition and warn any_class = re.search(r"class\s+(\w+)\s*\(", code) if any_class: class_name = any_class.group(1) logger.warning( f"Could not find Scene class, using first class found: {class_name}" ) return class_name # Last resort - parse the AST to find classes try: tree = ast.parse(code) for node in ast.walk(tree): if isinstance(node, ast.ClassDef): logger.warning( f"Using first class from AST parsing: {node.name}" ) return node.name except Exception as e: logger.error(f"Failed to parse code AST: {e}") # Absolute fallback logger.error("No scene class found in code! This will likely cause rendering to fail.") return "Scene" # fallback def _find_output_file( self, directory: Path, scene_name: str, extension: str ) -> Optional[Path]: """Find output file with given scene name and extension.""" for file in directory.glob(f"{scene_name}*.{extension}"): return file return None async def main(): """Main function for running the orchestrator.""" import argparse parser = argparse.ArgumentParser(description="NeuroAnim STEM Animation Generator") parser.add_argument("topic", help="STEM topic for the animation") parser.add_argument( "--audience", choices=["elementary", "middle_school", "high_school", "college", "general"], default="general", help="Target audience", ) parser.add_argument( "--duration", type=float, default=2.0, help="Animation duration in minutes" ) parser.add_argument("--output", default="animation.mp4", help="Output filename") parser.add_argument( "--api-key", help="Hugging Face API key (or set HUGGINGFACE_API_KEY env var)" ) parser.add_argument( "--elevenlabs-key", help="ElevenLabs API key (or set ELEVENLABS_API_KEY env var)", ) args = parser.parse_args() # Initialize and run orchestrator orchestrator = NeuroAnimOrchestrator( hf_api_key=args.api_key, elevenlabs_api_key=args.elevenlabs_key ) try: await orchestrator.initialize() results = await orchestrator.generate_animation( topic=args.topic, target_audience=args.audience, animation_length_minutes=args.duration, output_filename=args.output, ) if results["success"]: print("\nšŸŽ‰ Animation Generated Successfully!") print(f"šŸ“¹ Output file: {results['output_file']}") print(f"šŸŽÆ Topic: {results['topic']}") print(f"šŸ‘„ Audience: {results['target_audience']}") print(f"\nšŸ“ Concept Plan:") print( results["concept_plan"][:500] + "..." if len(results["concept_plan"]) > 500 else results["concept_plan"] ) print(f"\nšŸŽ­ Narration:") print( results["narration"][:300] + "..." if len(results["narration"]) > 300 else results["narration"] ) print(f"\nšŸ“š Quiz Questions:") print(results["quiz"]) else: print(f"\nāŒ Animation Generation Failed: {results['error']}") except KeyboardInterrupt: print("\nāš ļø Process interrupted by user") except Exception as e: print(f"\nšŸ’„ Unexpected error: {str(e)}") finally: await orchestrator.cleanup() if __name__ == "__main__": asyncio.run(main())