code_generator.pyā¢16.7 kB
import google.generativeai as genai
import json
import os
from typing import Dict, List, Optional
import logging
import time
class AECodeGenerator:
def __init__(self, api_key: str):
genai.configure(api_key=api_key)
self.model = genai.GenerativeModel('gemini-pro')
self.max_attempts = 5
self.logger = logging.getLogger(__name__)
def generate_code(self,
action: str,
params: Dict,
error_message: Optional[str] = None,
previous_attempts: List[str] = None) -> Dict:
"""
Generate After Effects code with continuous improvement capability.
"""
try:
# Build the prompt with context
prompt = self._build_prompt(action, params, error_message, previous_attempts)
# Generate code with the model
response = self.model.generate_content(prompt)
# Extract and validate the code
code = self._extract_code(response.text)
# Validate the generated code
validation_result = self._validate_code(code, action, params)
if validation_result["is_valid"]:
return {
"status": "success",
"code": code,
"message": "Code generated successfully",
"validation": validation_result
}
else:
# If not valid and we haven't reached max attempts, try again
if len(previous_attempts or []) < self.max_attempts:
return self.generate_code(
action,
params,
validation_result["error"],
(previous_attempts or []) + [code]
)
else:
return {
"status": "error",
"message": f"Failed to generate valid code after {self.max_attempts} attempts",
"last_error": validation_result["error"],
"last_attempt": code
}
except Exception as e:
self.logger.error(f"Error generating code: {str(e)}")
return {
"status": "error",
"message": f"Error generating code: {str(e)}"
}
def _build_prompt(self,
action: str,
params: Dict,
error_message: Optional[str] = None,
previous_attempts: List[str] = None) -> str:
"""
Build a comprehensive prompt for code generation.
"""
prompt_parts = [
"You are an expert After Effects ExtendScript developer. ",
"Generate code that is robust, efficient, and follows best practices. ",
"\nAction: " + action,
"\nParameters: " + json.dumps(params, indent=2)
]
if error_message:
prompt_parts.append(f"\nPrevious error: {error_message}")
prompt_parts.append("\nPlease fix the following issues:")
if previous_attempts:
prompt_parts.append("\nPrevious attempts:")
for i, attempt in enumerate(previous_attempts, 1):
prompt_parts.append(f"\nAttempt {i}:")
prompt_parts.append(attempt)
prompt_parts.extend([
"\nRequirements:",
"1. Code must be valid ExtendScript",
"2. Include proper error handling",
"3. Use app.beginUndoGroup() and app.endUndoGroup()",
"4. Validate all inputs",
"5. Include comments explaining complex logic",
"6. Handle edge cases",
"7. Use efficient algorithms",
"8. Follow After Effects best practices",
"\nGenerate the code:"
])
return "\n".join(prompt_parts)
def _extract_code(self, response: str) -> str:
"""
Extract code from the model's response.
"""
# Look for code blocks
if "```javascript" in response:
code = response.split("```javascript")[1].split("```")[0].strip()
elif "```jsx" in response:
code = response.split("```jsx")[1].split("```")[0].strip()
else:
code = response.strip()
return code
def _validate_code(self, code: str, action: str, params: Dict) -> Dict:
"""
Validate the generated code.
"""
validation = {
"is_valid": True,
"error": None,
"warnings": []
}
# Check for basic syntax
if not code.strip():
validation["is_valid"] = False
validation["error"] = "Empty code generated"
return validation
# Check for required elements
required_elements = [
"app.beginUndoGroup",
"app.endUndoGroup",
"try",
"catch"
]
for element in required_elements:
if element not in code:
validation["warnings"].append(f"Missing {element}")
# Check for parameter validation
if "params" in code and not any(param in code for param in params.keys()):
validation["warnings"].append("Not all parameters are being used")
# Check for proper error handling
if "throw new Error" not in code:
validation["warnings"].append("No explicit error throwing")
# Check for proper function structure
if not any(keyword in code for keyword in ["function", "=>"]):
validation["is_valid"] = False
validation["error"] = "No function definition found"
return validation
return validation
def improve_code(self,
code: str,
feedback: str,
context: Dict) -> Dict:
"""
Improve existing code based on feedback.
"""
prompt = f"""
Improve the following After Effects ExtendScript code based on the feedback.
Original Code:
```javascript
{code}
```
Feedback:
{feedback}
Context:
{json.dumps(context, indent=2)}
Requirements:
1. Fix any issues mentioned in the feedback
2. Maintain or improve error handling
3. Keep the core functionality intact
4. Add any missing features
5. Optimize performance if possible
Generate the improved code:
"""
try:
response = self.model.generate_content(prompt)
improved_code = self._extract_code(response.text)
return {
"status": "success",
"code": improved_code,
"message": "Code improved successfully"
}
except Exception as e:
return {
"status": "error",
"message": f"Error improving code: {str(e)}"
}
def generate_animation_code(self,
animation_type: str,
params: Dict) -> Dict:
"""
Generate specific animation code.
"""
prompt = f"""
Generate After Effects ExtendScript code for {animation_type} animation.
Parameters:
{json.dumps(params, indent=2)}
Requirements:
1. Create smooth, professional animations
2. Include proper easing
3. Handle all edge cases
4. Support customization
5. Include error handling
6. Use expressions where appropriate
7. Optimize for performance
Generate the animation code:
"""
try:
response = self.model.generate_content(prompt)
code = self._extract_code(response.text)
return {
"status": "success",
"code": code,
"message": f"{animation_type} animation code generated successfully"
}
except Exception as e:
return {
"status": "error",
"message": f"Error generating animation code: {str(e)}"
}
def generate_custom_code_from_prompt(self, prompt: str, error_context: Optional[Dict] = None, previous_attempts: List[str] = None) -> Dict:
"""
Generate custom After Effects code from a natural language prompt.
Includes self-healing and memory capabilities.
Args:
prompt: Natural language description of what the code should do
error_context: Optional error information from previous attempts
previous_attempts: List of previously generated code that failed
Returns:
Dict containing status, code, and message
"""
try:
# Build a comprehensive system prompt
system_prompt = """You are an expert After Effects ExtendScript developer with deep knowledge of the After Effects API.
Your task is to generate robust, production-ready ExtendScript code based on natural language descriptions.
IMPORTANT GUIDELINES:
1. Generate ONLY ExtendScript code that can run directly in After Effects
2. Always use proper error handling with try/catch blocks
3. Always use app.beginUndoGroup() and app.endUndoGroup() for proper undo support
4. Validate all inputs and handle edge cases
5. Use clear variable names and add comments for complex logic
6. Return meaningful success/error messages
7. Focus on solving the specific task requested
8. If the task involves layers, compositions, or effects, make sure they exist before manipulating them
9. Use efficient algorithms and follow After Effects best practices
10. If the task is complex, break it down into smaller functions
COMMON AFTER EFFECTS OBJECTS AND PROPERTIES:
- app.project - The current project
- app.project.activeItem - The currently active composition
- app.project.item(index) - Get item at index
- comp.layer(index) - Get layer at index
- layer.property("property_name") - Access layer properties
- layer.effect("effect_name") - Access layer effects
- layer.mask("mask_name") - Access layer masks
IMPORTANT: Your code will be executed directly in After Effects, so it must be valid ExtendScript syntax.
"""
# Determine if this is a retry with error information
if error_context and isinstance(error_context, dict):
error_message = error_context.get("error", "")
error_details = error_context.get("details", "")
# Add error context to the prompt
user_prompt = f"""Task: {prompt}
Previous Error: {error_message}
Error Details: {error_details}
Please fix the code to address this error and complete the task successfully.
"""
else:
# First attempt
user_prompt = f"Task: {prompt}\n\nGenerate ExtendScript code to accomplish this task in After Effects."
# Add previous attempts if available
if previous_attempts and len(previous_attempts) > 0:
user_prompt += "\n\nPrevious attempts that failed:"
for i, attempt in enumerate(previous_attempts, 1):
user_prompt += f"\n\nAttempt {i}:\n```javascript\n{attempt}\n```"
user_prompt += "\n\nPlease generate a new solution that addresses the issues in the previous attempts."
# Generate code with the model
response = self.model.generate_content(
contents=[
{"role": "system", "parts": [{"text": system_prompt}]},
{"role": "user", "parts": [{"text": user_prompt}]}
],
generation_config={
"temperature": 0.2,
"top_p": 0.8,
"top_k": 40,
"max_output_tokens": 8192,
}
)
# Extract the code
code = self._extract_code(response.text)
# Basic validation
if not code or len(code.strip()) < 10:
return {
"status": "error",
"message": "Generated code is too short or empty",
"code": code
}
# Check for required elements
validation = self._validate_code(code, "custom_prompt", {"prompt": prompt})
# Log the generated code
self.logger.info(f"Generated code for prompt: {prompt[:50]}...")
# Store in memory for future reference
self._store_in_memory(prompt, code, error_context)
return {
"status": "success",
"code": code,
"message": "Custom code generated successfully",
"validation": validation
}
except Exception as e:
self.logger.error(f"Error generating custom code: {str(e)}")
return {
"status": "error",
"message": f"Error generating custom code: {str(e)}",
"code": ""
}
def _store_in_memory(self, prompt: str, code: str, error_context: Optional[Dict] = None):
"""Store generated code in memory for learning and future reference"""
try:
# Create memory directory if it doesn't exist
memory_dir = os.path.join(os.path.dirname(__file__), "memory")
os.makedirs(memory_dir, exist_ok=True)
# Create a memory file with sanitized prompt as filename
sanitized = "".join(c if c.isalnum() else "_" for c in prompt[:30])
memory_file = os.path.join(memory_dir, f"{sanitized}_{int(time.time())}.json")
# Store the memory
with open(memory_file, "w") as f:
json.dump({
"prompt": prompt,
"code": code,
"timestamp": time.time(),
"error_context": error_context
}, f, indent=2)
except Exception as e:
self.logger.error(f"Error storing memory: {str(e)}")
def fix_code_from_error(self, code: str, error_message: str, task_description: str) -> Dict:
"""
Fix code based on error message and task description.
Args:
code: The original code that produced an error
error_message: The error message from After Effects
task_description: Description of what the code should do
Returns:
Dict containing status, fixed code, and message
"""
try:
# Build the prompt for fixing the code
prompt = f"""Fix the following After Effects ExtendScript code that produced an error.
Original Code:
```javascript
{code}
```
Error Message:
{error_message}
Task Description:
{task_description}
Requirements:
1. Fix the specific issue that caused the error
2. Maintain the original functionality
3. Ensure proper error handling
4. Use app.beginUndoGroup() and app.endUndoGroup()
5. Return meaningful success/error messages
Generate the fixed code:
"""
# Generate fixed code
response = self.model.generate_content(prompt)
fixed_code = self._extract_code(response.text)
return {
"status": "success",
"code": fixed_code,
"message": "Code fixed successfully"
}
except Exception as e:
self.logger.error(f"Error fixing code: {str(e)}")
return {
"status": "error",
"message": f"Error fixing code: {str(e)}",
"code": code # Return original code if fixing failed
}
# Example usage:
if __name__ == "__main__":
# Initialize the code generator
generator = AECodeGenerator(api_key="your_gemini_api_key")
# Example: Generate color animation code
result = generator.generate_animation_code(
"color",
{
"layerName": "text1",
"duration": 5,
"colors": ["#FF0000", "#00FF00", "#0000FF"]
}
)
print(json.dumps(result, indent=2))