import os import litellm from fastapi import HTTPException import logging from ..utils.system_prompts import get_system_prompt from ..utils.code_postprocessor import post_process_code from ..utils.code_validator import CodeValidator from ..utils.code_fixer import CodeFixer from ..inputs.processor import InputProcessor from ..utils.content_preprocessor import preprocess_long_content, get_script_mode_prompt_for_long_content logger = logging.getLogger(__name__) def _generate_with_reflexion(prompt: str, category: str, job_id: str = "unknown") -> str: """ Generate code using the Reflexion Agent. Args: prompt: Processed user prompt category: Animation category job_id: Job identifier for logging Returns: Generated and refined code """ from ..agents.reflexion_agent import ReflexionAgent agent = ReflexionAgent() code, state = agent.generate_with_reflection( goal=prompt, category=category, job_id=job_id ) logger.info(f"🔄 Reflexion complete: {state.iteration} iterations, {state.total_issues_found} issues found, {state.total_issues_fixed} fixed") return code def _generate_legacy(prompt: str, category: str, max_attempts: int = 3) -> str: """ Legacy code generation without Reflexion. Used as fallback when Reflexion is disabled or fails. Args: prompt: Processed user prompt category: Animation category max_attempts: Maximum generation attempts Returns: Generated code """ validator = CodeValidator() fixer = CodeFixer() primary_model = os.getenv("CODE_GEN_MODEL") fallback_model = os.getenv("CODE_GEN_FALLBACK_MODEL", primary_model) for attempt in range(max_attempts): try: # Use fallback model on retry attempts model = fallback_model if attempt > 0 else primary_model # Get dynamic system prompt based on category system_prompt = get_system_prompt(category) # Detect if input is a ready-made script (long content) vs short prompt word_count = len(prompt.split()) is_script_mode = word_count > 200 # For very long content, preprocess into sections processed_prompt = prompt section_count = 0 if word_count > 1000: processed_prompt, section_count = preprocess_long_content(prompt) if section_count > 0: # Very long content - use sectioned prompt logger.info(f"📝 LONG DOCUMENT MODE (Legacy): {word_count} words -> {section_count} sections") user_content = get_script_mode_prompt_for_long_content(processed_prompt, section_count) elif is_script_mode: logger.info(f"📝 SCRIPT MODE (Legacy): Input has {word_count} words - treating as ready-made script") user_content = f"""# 🎬 SCRIPT MODE - ANIMATE THE USER'S CONTENT ## IMPORTANT: The user has provided their COMPLETE script/content below. This is NOT a topic to research - this IS the exact narration/content they want animated. ## YOUR TASK: 1. **Use the content below AS the voiceover text** - split it into logical sections 2. **Create beautiful animations that MATCH each section** of their content 3. **Do NOT rewrite, summarize, or generate new information** - animate THEIR words 4. **Every paragraph/section should become a voiceover block** with matching visuals ## USER'S SCRIPT TO ANIMATE: --- {prompt} --- NOTE!!!: 1. NO BLANK SCREENS: Keep the screen populated. If a voiceover is playing, show something. 2. NO OVERLAPS: Ensure text and objects do not overlap. Use `next_to` and `arrange`. 3. CLEAN TRANSITIONS: Fade out old content before showing new content, but don't leave the screen empty for long. 4. VARIED ANIMATIONS: Use a mix of Write, FadeIn, GrowFromCenter, etc. 5. STAY ON SCREEN: Ensure all text and objects are within the screen boundaries. Use .scale_to_fit_width(config.frame_width - 1) for large groups.""" else: logger.info(f"📝 GENERATION MODE (Legacy): Input has {word_count} words - LLM will generate content") user_content = f"Create a video about:\n\n{prompt}\n\n NOTE!!!:\n1. NO BLANK SCREENS: Keep the screen populated. If a voiceover is playing, show something.\n2. NO OVERLAPS: Ensure text and objects do not overlap. Use `next_to` and `arrange`.\n3. CLEAN TRANSITIONS: Fade out old content before showing new content, but don't leave the screen empty for long.\n4. VARIED ANIMATIONS: Use a mix of Write, FadeIn, GrowFromCenter, etc.\n5. STAY ON SCREEN: Ensure all text and objects are within the screen boundaries. Use .scale_to_fit_width(config.frame_width - 1) for large groups." messages = [ { "role": "system", "content": system_prompt, }, { "role": "user", "content": user_content, }, ] logger.info(f"Generating code (attempt {attempt + 1}/{max_attempts}) with model {model}") # Only set max_tokens for long documents kwargs = { "model": model, "messages": messages, "num_retries": 2 } if section_count > 0: kwargs["max_tokens"] = 12000 response = litellm.completion(**kwargs) raw_code = response.choices[0].message.content # Extract code if wrapped in markdown (handle various formats) import re # Try different markdown patterns code_patterns = [ r'```python\n(.*?)```', # Standard: ```python ... ``` r'````python\n(.*?)````', # Quad backticks r'```py\n(.*?)```', # ```py r'```\n(.*?)```', # Just backticks without language ] for pattern in code_patterns: match = re.search(pattern, raw_code, re.DOTALL) if match: raw_code = match.group(1).strip() break # If still has backticks, try to clean up if raw_code.startswith('```'): lines = raw_code.split('\n') # Remove first line if it's just ```python or similar if lines[0].strip().startswith('```'): lines = lines[1:] # Remove last line if it's just ``` if lines and lines[-1].strip() == '```': lines = lines[:-1] raw_code = '\n'.join(lines) # Post-process the code to fix common issues processed_code = post_process_code(raw_code) # Validate code is_valid, errors = validator.validate(processed_code) if is_valid: logger.info("Code validation passed") return processed_code # Try to auto-fix logger.warning(f"Code validation failed with {len(errors)} errors, attempting auto-fix") fixed_code, is_fixed, remaining_errors = fixer.fix_and_validate(processed_code, max_attempts=2) if is_fixed: logger.info("Code auto-fixed successfully") return fixed_code # If last attempt, return best code we have if attempt == max_attempts - 1: error_msg = f"Code generation failed after {max_attempts} attempts. Errors: {remaining_errors}" logger.error(error_msg) raise HTTPException( status_code=500, detail=error_msg ) logger.info(f"Retrying code generation (attempt {attempt + 2}/{max_attempts})") except HTTPException: raise except Exception as e: logger.error(f"Error in code generation attempt {attempt + 1}: {str(e)}") if attempt == max_attempts - 1: raise HTTPException( status_code=500, detail=f"Failed to generate animation response after {max_attempts} attempts: {str(e)}" ) # Should not reach here, but just in case raise HTTPException( status_code=500, detail="Failed to generate valid animation code after all attempts" ) def generate_animation_response( input_data: str, input_type: str = "text", category: str = "mathematical", max_attempts: int = 3, job_id: str = "unknown" ) -> str: """Generate Manim animation code from input with validation and auto-fixing. Uses Reflexion Agent when enabled (REFLEXION_ENABLED=true) for improved code quality through self-critique and iterative improvement. Args: input_data (str): User's input (text, URL, or PDF path) input_type (str): Type of input ('text', 'url', 'pdf') category (str): Animation category (tech_system, product_startup, mathematical) max_attempts (int): Maximum generation attempts (for legacy mode) job_id (str): Job identifier for logging Returns: str: Generated Manim animation code (validated and post-processed) Raises: HTTPException: If code generation fails after all attempts """ # Process input to get the actual prompt text try: prompt = InputProcessor.process(input_type, input_data) except Exception as e: logger.error(f"Input processing failed: {e}") raise HTTPException(status_code=400, detail=f"Input processing failed: {str(e)}") # Check if Reflexion is enabled reflexion_enabled = os.getenv("REFLEXION_ENABLED", "true").lower() == "true" if reflexion_enabled: logger.info("🔄 Using Reflexion Agent for code generation") try: code = _generate_with_reflexion(prompt, category, job_id) # Final validation validator = CodeValidator() fixer = CodeFixer() is_valid, errors = validator.validate(code) if not is_valid: logger.warning(f"Reflexion code has {len(errors)} validation errors, attempting auto-fix") code, is_fixed, _ = fixer.fix_and_validate(code, max_attempts=2) return code except Exception as e: logger.error(f"Reflexion failed, falling back to legacy: {e}") # Fall through to legacy generation else: logger.info("📝 Using legacy code generation (Reflexion disabled)") # Legacy generation (fallback or when Reflexion is disabled) return _generate_legacy(prompt, category, max_attempts)