A class for generating and managing Manim code.
| 27 | from src.rag.vector_store import RAGVectorStore # Import RAGVectorStore |
| 28 | |
| 29 | class CodeGenerator: |
| 30 | """A class for generating and managing Manim code.""" |
| 31 | |
| 32 | def __init__(self, scene_model, helper_model, output_dir="output", print_response=False, use_rag=False, use_context_learning=False, context_learning_path="data/context_learning", chroma_db_path="rag/chroma_db", manim_docs_path="rag/manim_docs", embedding_model="azure/text-embedding-3-large", use_visual_fix_code=False, use_langfuse=True, session_id=None): |
| 33 | """Initialize the CodeGenerator. |
| 34 | |
| 35 | Args: |
| 36 | scene_model: The model used for scene generation |
| 37 | helper_model: The model used for helper tasks |
| 38 | output_dir (str, optional): Directory for output files. Defaults to "output". |
| 39 | print_response (bool, optional): Whether to print model responses. Defaults to False. |
| 40 | use_rag (bool, optional): Whether to use RAG. Defaults to False. |
| 41 | use_context_learning (bool, optional): Whether to use context learning. Defaults to False. |
| 42 | context_learning_path (str, optional): Path to context learning examples. Defaults to "data/context_learning". |
| 43 | chroma_db_path (str, optional): Path to ChromaDB. Defaults to "rag/chroma_db". |
| 44 | manim_docs_path (str, optional): Path to Manim docs. Defaults to "rag/manim_docs". |
| 45 | embedding_model (str, optional): Name of embedding model. Defaults to "azure/text-embedding-3-large". |
| 46 | use_visual_fix_code (bool, optional): Whether to use visual code fixing. Defaults to False. |
| 47 | use_langfuse (bool, optional): Whether to use Langfuse logging. Defaults to True. |
| 48 | session_id (str, optional): Session identifier. Defaults to None. |
| 49 | """ |
| 50 | self.scene_model = scene_model |
| 51 | self.helper_model = helper_model |
| 52 | self.output_dir = output_dir |
| 53 | self.print_response = print_response |
| 54 | self.use_rag = use_rag |
| 55 | self.use_context_learning = use_context_learning |
| 56 | self.context_learning_path = context_learning_path |
| 57 | self.context_examples = self._load_context_examples() if use_context_learning else None |
| 58 | self.manim_docs_path = manim_docs_path |
| 59 | |
| 60 | self.use_visual_fix_code = use_visual_fix_code |
| 61 | self.banned_reasonings = get_banned_reasonings() |
| 62 | self.session_id = session_id # Use session_id passed from VideoGenerator |
| 63 | |
| 64 | if use_rag: |
| 65 | self.vector_store = RAGVectorStore( |
| 66 | chroma_db_path=chroma_db_path, |
| 67 | manim_docs_path=manim_docs_path, |
| 68 | embedding_model=embedding_model, |
| 69 | session_id=self.session_id, |
| 70 | use_langfuse=use_langfuse |
| 71 | ) |
| 72 | else: |
| 73 | self.vector_store = None |
| 74 | |
| 75 | def _load_context_examples(self) -> str: |
| 76 | """Load all context learning examples from the specified directory. |
| 77 | |
| 78 | Returns: |
| 79 | str: Formatted context learning examples, or None if no examples found. |
| 80 | """ |
| 81 | examples = [] |
| 82 | for example_file in glob.glob(f"{self.context_learning_path}/**/*.py", recursive=True): |
| 83 | with open(example_file, 'r') as f: |
| 84 | examples.append(f"# Example from {os.path.basename(example_file)}\n{f.read()}\n") |
| 85 | |
| 86 | # Format examples using get_prompt_context_learning_code instead of _prompt_context_learning |