Spaces:
Running
Running
| """ | |
| Utilities for cleaning and validating Manim code generated by LLMs. | |
| """ | |
| import re | |
| import logging | |
| import json | |
| logger = logging.getLogger(__name__) | |
| def clean_manim_code(raw_code): | |
| """ | |
| Clean Manim code from LLM responses by removing markdown formatting | |
| and ensuring proper structure. | |
| Args: | |
| raw_code (str): The raw code from the LLM response | |
| Returns: | |
| str: Cleaned, executable Python code | |
| """ | |
| # Start with the raw code | |
| code = raw_code | |
| # Extract code from markdown code blocks if present | |
| if "```python" in code: | |
| parts = code.split("```python") | |
| if len(parts) > 1: | |
| code = parts[1] | |
| if "```" in code: | |
| code = code.split("```")[0] | |
| elif "```" in code: | |
| parts = code.split("```") | |
| if len(parts) > 1: | |
| code = parts[1] | |
| if "```" in parts[1]: | |
| code = code.split("```")[0] | |
| # Remove any remaining backticks | |
| code = code.replace('```', '') | |
| # Ensure code begins with the necessary import | |
| if not code.strip().startswith('from manim import'): | |
| code = 'from manim import *\n\n' + code | |
| # Verify the code contains a Scene class | |
| if 'class' not in code or 'Scene' not in code: | |
| logger.warning("Generated code does not contain a proper Scene class") | |
| # Add a basic scene structure if missing | |
| if 'class ManimScene(Scene):' not in code: | |
| code = 'from manim import *\n\nclass ManimScene(Scene):\n def construct(self):\n ' + code | |
| # Verify the code has a construct method | |
| if 'def construct(self)' not in code: | |
| logger.warning("Generated code does not contain a construct method") | |
| # Try to find where the class is defined and add construct method | |
| class_match = re.search(r'class\s+\w+\s*\(\s*Scene\s*\)\s*:', code) | |
| if class_match: | |
| insert_pos = class_match.end() | |
| code = code[:insert_pos] + '\n def construct(self):\n pass\n' + code[insert_pos:] | |
| # Ensure there's a wait at the end if not present | |
| if 'self.wait(' not in code.split('def construct')[-1]: | |
| # Find the end of the construct method to add wait | |
| construct_body_match = re.search(r'def\s+construct\s*\(\s*self\s*\)\s*:', code) | |
| if construct_body_match: | |
| # Check if the method has content | |
| method_content = code[construct_body_match.end():] | |
| indentation = ' ' # Default indentation | |
| # Try to determine indentation from code | |
| indent_match = re.search(r'\n(\s+)', method_content) | |
| if indent_match: | |
| indentation = indent_match.group(1) | |
| # Find a good place to insert the wait | |
| if '}' in method_content.splitlines()[-1]: # If last line closes something | |
| code = code.rstrip() + f'\n{indentation}self.wait(1)\n' | |
| else: | |
| code = code.rstrip() + f'\n{indentation}self.wait(1)\n' | |
| return code.strip() | |
| def parse_scenario_from_llm_response(content): | |
| """ | |
| Extract structured scenario information from an LLM response. | |
| Args: | |
| content (str): The LLM response text | |
| Returns: | |
| dict: Extracted scenario dictionary | |
| """ | |
| try: | |
| # Try to find and extract a JSON object | |
| json_match = re.search(r'\{.*\}', content, re.DOTALL) | |
| if json_match: | |
| json_str = json_match.group(0) | |
| scenario_dict = json.loads(json_str) | |
| return scenario_dict | |
| except Exception as e: | |
| logger.error(f"Error parsing scenario JSON: {e}") | |
| # Manual parsing fallback | |
| scenario = { | |
| "title": "", | |
| "objects": [], | |
| "transformations": [], | |
| "equations": [] | |
| } | |
| # Simple pattern matching to extract information | |
| title_match = re.search(r'title["\s:]+([^"]+)', content, re.IGNORECASE) | |
| if title_match: | |
| scenario["title"] = title_match.group(1).strip() | |
| # Extract lists with various possible formats | |
| objects_pattern = r'objects[":\s\[]+([^\]]+)' | |
| objects_match = re.search(objects_pattern, content, re.IGNORECASE | re.DOTALL) | |
| if objects_match: | |
| objects_text = objects_match.group(1) | |
| # Handle both comma-separated and quote-wrapped items | |
| objects = re.findall(r'"([^"]+)"', objects_text) | |
| if not objects: | |
| objects = [item.strip() for item in objects_text.split(',')] | |
| scenario["objects"] = objects | |
| # Similar extraction for transformations | |
| trans_pattern = r'transformations[":\s\[]+([^\]]+)' | |
| trans_match = re.search(trans_pattern, content, re.IGNORECASE | re.DOTALL) | |
| if trans_match: | |
| trans_text = trans_match.group(1) | |
| transformations = re.findall(r'"([^"]+)"', trans_text) | |
| if not transformations: | |
| transformations = [item.strip() for item in trans_text.split(',')] | |
| scenario["transformations"] = transformations | |
| # Extract equations if present | |
| equations_pattern = r'equations[":\s\[]+([^\]]+)' | |
| equations_match = re.search(equations_pattern, content, re.IGNORECASE | re.DOTALL) | |
| if equations_match: | |
| equations_text = equations_match.group(1) | |
| if equations_text.lower().strip() in ['null', 'none']: | |
| scenario["equations"] = None | |
| else: | |
| equations = re.findall(r'"([^"]+)"', equations_text) | |
| if not equations: | |
| equations = [item.strip() for item in equations_text.split(',')] | |
| scenario["equations"] = equations | |
| return scenario | |
| def validate_manim_code(code): | |
| """ | |
| Perform basic validation on Manim code to catch common issues. | |
| Args: | |
| code (str): The Manim code to validate | |
| Returns: | |
| tuple: (is_valid, error_message) | |
| """ | |
| # Check for basic Python syntax errors | |
| try: | |
| compile(code, '<string>', 'exec') | |
| except SyntaxError as e: | |
| return False, f"Syntax error: {str(e)}" | |
| # Check for necessary components | |
| if 'from manim import' not in code: | |
| return False, "Missing Manim import" | |
| if 'class' not in code or 'Scene' not in code: | |
| return False, "No Scene class defined" | |
| if 'def construct(self)' not in code: | |
| return False, "No construct method defined" | |
| # Check for common Manim issues | |
| if 'self.play(' not in code and 'self.add(' not in code: | |
| return False, "No objects added to scene (missing self.play or self.add calls)" | |
| # All checks passed | |
| return True, "Code appears valid" | |