File size: 6,364 Bytes
d970572 |
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 |
import os
import re
from google import genai
from google.genai import types as genai_types
import logging
from .gemini import SYSTEM_PROMPT, base_prompt_instructions
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
def fix_manim_code(faulty_code: str, error_message: str, original_context: str):
api_key = os.getenv("GEMINI_API_KEY")
if not api_key:
logging.error("GEMINI_API_KEY not found in environment variables for fallback.")
return None, None
client = genai.Client(api_key=api_key)
fix_prompt_text = (
f"The following Manim code, intended to '{original_context}', failed with an error.\n\n"
"### FAULTY CODE:\n"
f"```python\n{faulty_code}\n```\n\n"
"### ERROR MESSAGE:\n"
f"```\n{error_message}\n```\n\n"
"### INSTRUCTIONS:\n"
"1. Analyze the error message and the faulty code.\n"
"2. Correct the code to fix the specific error reported.\n"
"3. Ensure the corrected code still fulfills the original request and adheres strictly to *all* the requirements listed below.\n"
"4. Pay close attention to vector dimensions, matrix operations, allowed Manim methods, and total duration (30 seconds).\n"
"5. If the code logic changes significantly, update the narration accordingly.\n"
"6. Return *only* the corrected code and narration using the '### MANIM CODE:' and '### NARRATION:' delimiters, just like the original request.\n\n"
"### REQUIREMENTS (Apply these to the corrected code):\n"
f"{base_prompt_instructions}"
)
contents = [fix_prompt_text]
logging.info("Attempting to fix Manim code via fallback...")
try:
generation_config = genai_types.GenerateContentConfig(
system_instruction=SYSTEM_PROMPT
)
response = client.models.generate_content(
model="gemini-2.0-flash",
contents=contents,
config=generation_config
)
if response:
try:
content = response.text
logging.info("Received response from fallback attempt.")
if "### NARRATION:" in content:
manim_code, narration = content.split("### NARRATION:", 1)
manim_code = re.sub(r"```python", "", manim_code).replace("```", "").strip()
narration = narration.strip()
if "from manim import *" not in manim_code:
logging.warning("Adding missing 'from manim import *' (fallback fix).")
manim_code = "from manim import *\nimport numpy as np\n" + manim_code
elif "import numpy as np" not in manim_code:
logging.warning("Adding missing 'import numpy as np' (fallback fix).")
lines = manim_code.splitlines()
for i, line in enumerate(lines):
if "from manim import *" in line:
lines.insert(i + 1, "import numpy as np")
manim_code = "\n".join(lines)
break
logging.info("Successfully parsed fixed code and narration from fallback.")
return {"manim_code": manim_code, "output_file": "output.mp4"}, narration
else:
logging.warning("Delimiter '### NARRATION:' not found in fallback response. Attempting fallback extraction.")
code_match = re.search(r'```python(.*?)```', content, re.DOTALL)
if code_match:
manim_code = code_match.group(1).strip()
narration_part = content.split('```', 2)[-1].strip()
narration = narration_part if len(narration_part) > 20 else ""
if not narration:
logging.warning("Fallback narration extraction resulted in empty or very short text (fallback fix).")
else:
logging.info("Successfully parsed code and narration using fallback regex (fallback fix).")
if "from manim import *" not in manim_code:
logging.warning("Adding missing 'from manim import *' (fallback fix, regex path).")
manim_code = "from manim import *\nimport numpy as np\n" + manim_code
elif "import numpy as np" not in manim_code:
logging.warning("Adding missing 'import numpy as np' (fallback fix, regex path).")
lines = manim_code.splitlines()
for i, line in enumerate(lines):
if "from manim import *" in line:
lines.insert(i + 1, "import numpy as np")
manim_code = "\n".join(lines)
break
logging.info("Successfully parsed fixed code using fallback extraction.")
return {"manim_code": manim_code, "output_file": "output.mp4"}, narration
else:
logging.error("Fallback extraction failed: No Python code block found in fallback response.")
logging.debug(f"Fallback content without code block:\n{content}")
return None, None
except ValueError:
logging.error("Could not extract text from the fallback response.")
if response.prompt_feedback and response.prompt_feedback.block_reason:
logging.error(f"Fallback content generation blocked. Reason: {response.prompt_feedback.block_reason.name}")
return None, None
except Exception as e:
logging.exception(f"Error processing fallback response: {e}")
return None, None
else:
logging.error("No response received from Gemini during fallback attempt.")
return None, None
except Exception as e:
logging.exception(f"Error calling Gemini API during fallback: {e}")
return None, None
|