File size: 5,239 Bytes
f0e2e50 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 |
import os
import logging
import tempfile
import shutil
from pathlib import Path
from typing import Optional, Tuple, Dict, Any
from groq import Groq
from config.settings import (
GROQ_API_BASE_URL,
DEFAULT_MODEL,
AVAILABLE_MODELS,
QUALITY_SETTINGS,
APP_SETTINGS,
TEMP_DIR
)
from config.constants import (
ModelType,
VideoQuality,
ErrorType,
FILE_CONSTANTS,
ANIMATION_CONSTANTS
)
from utils.code_validator import CodeValidator
from utils.prompt_manager import PromptManager
from utils.error_handler import ErrorHandler
from utils.metrics_collector import MetricsCollector
logger = logging.getLogger(__name__)
class ManimVideoGenerator:
def __init__(self, groq_api_key: str):
self.client = Groq(api_key=groq_api_key)
self.available_models = AVAILABLE_MODELS
self.temp_dir = None
self.metrics = MetricsCollector()
self.prompt_manager = PromptManager()
self.code_validator = CodeValidator()
def enhance_prompt(self, user_prompt: str, model: str = DEFAULT_MODEL) -> str:
try:
model_config = self.available_models[model]
return self.prompt_manager.enhance_prompt(
user_prompt,
self.client,
model,
temperature=model_config["temperature"],
max_tokens=model_config["max_tokens"]
)
except Exception as e:
logger.error(f"Error enhancing prompt: {e}")
ErrorHandler.handle_error(ErrorType.PROMPT_ERROR, str(e))
return user_prompt
def generate_manim_code(self, enhanced_prompt: str, model: str = DEFAULT_MODEL) -> str:
try:
model_config = self.available_models[model]
return self.prompt_manager.generate_code(
enhanced_prompt,
self.client,
model,
temperature=0.3, # Lower temperature for code generation
max_tokens=model_config["max_tokens"]
)
except Exception as e:
logger.error(f"Error generating Manim code: {e}")
ErrorHandler.handle_error(ErrorType.CODE_GENERATION_ERROR, str(e))
raise
def process_and_validate_code(self, raw_code: str) -> str:
try:
cleaned_code = self.code_validator.fix_common_issues(raw_code)
validation_result = self.code_validator.validate_code(cleaned_code)
if not validation_result["is_valid"]:
raise ValueError(validation_result["error_message"])
return cleaned_code
except Exception as e:
logger.error(f"Error validating code: {e}")
ErrorHandler.handle_error(ErrorType.VALIDATION_ERROR, str(e))
raise
def create_temp_directory(self) -> str:
if self.temp_dir and os.path.exists(self.temp_dir):
shutil.rmtree(self.temp_dir)
self.temp_dir = tempfile.mkdtemp(
prefix=FILE_CONSTANTS["TEMP_DIR_PREFIX"],
dir=TEMP_DIR
)
return self.temp_dir
def save_code_to_file(self, code: str, temp_dir: str) -> str:
code_file = os.path.join(temp_dir, FILE_CONSTANTS["SCENE_FILE_NAME"])
try:
with open(code_file, 'w', encoding='utf-8') as f:
f.write(code)
return code_file
except Exception as e:
logger.error(f"Error saving code to file: {e}")
ErrorHandler.handle_error(ErrorType.SYSTEM_ERROR, str(e))
raise
def render_video(
self,
code_file: str,
temp_dir: str,
quality: str = VideoQuality.MEDIUM.value
) -> Tuple[bool, str, str]:
from utils.renderer import VideoRenderer
try:
renderer = VideoRenderer(
code_file=code_file,
temp_dir=temp_dir,
quality=quality,
metrics=self.metrics
)
return renderer.render()
except Exception as e:
logger.error(f"Error rendering video: {e}")
ErrorHandler.handle_error(ErrorType.RENDERING_ERROR, str(e))
return False, str(e), ""
def cleanup(self) -> None:
if self.temp_dir and os.path.exists(self.temp_dir):
try:
shutil.rmtree(self.temp_dir)
logger.info(f"Cleaned up temporary directory: {self.temp_dir}")
except Exception as e:
logger.error(f"Error cleaning up: {e}")
ErrorHandler.handle_error(ErrorType.SYSTEM_ERROR, str(e))
def get_metrics_summary(self) -> Dict[str, Any]:
return self.metrics.get_summary()
@property
def model_options(self) -> Dict[str, str]:
return {
model_id: f"{config['name']} ({config['description']})"
for model_id, config in self.available_models.items()
}
@property
def quality_options(self) -> Dict[str, str]:
return {
quality_id: config["description"]
for quality_id, config in QUALITY_SETTINGS.items()
} |