A class for managing vector stores for RAG (Retrieval Augmented Generation). This class handles creation, loading and querying of vector stores for both Manim core and plugin documentation. Args: chroma_db_path (str): Path to ChromaDB storage directory manim_docs_path (
| 19 | from task_generator import get_prompt_detect_plugins |
| 20 | |
| 21 | class RAGVectorStore: |
| 22 | """A class for managing vector stores for RAG (Retrieval Augmented Generation). |
| 23 | |
| 24 | This class handles creation, loading and querying of vector stores for both Manim core |
| 25 | and plugin documentation. |
| 26 | |
| 27 | Args: |
| 28 | chroma_db_path (str): Path to ChromaDB storage directory |
| 29 | manim_docs_path (str): Path to Manim documentation files |
| 30 | embedding_model (str): Name of the embedding model to use |
| 31 | trace_id (str, optional): Trace identifier for logging. Defaults to None |
| 32 | session_id (str, optional): Session identifier. Defaults to None |
| 33 | use_langfuse (bool, optional): Whether to use Langfuse logging. Defaults to True |
| 34 | helper_model: Helper model for processing. Defaults to None |
| 35 | """ |
| 36 | |
| 37 | def __init__(self, |
| 38 | chroma_db_path: str = "chroma_db", |
| 39 | manim_docs_path: str = "rag/manim_docs", |
| 40 | embedding_model: str = "text-embedding-ada-002", |
| 41 | trace_id: str = None, |
| 42 | session_id: str = None, |
| 43 | use_langfuse: bool = True, |
| 44 | helper_model = None): |
| 45 | self.chroma_db_path = chroma_db_path |
| 46 | self.manim_docs_path = manim_docs_path |
| 47 | self.embedding_model = embedding_model |
| 48 | self.trace_id = trace_id |
| 49 | self.session_id = session_id |
| 50 | self.use_langfuse = use_langfuse |
| 51 | self.helper_model = helper_model |
| 52 | self.enc = tiktoken.encoding_for_model("gpt-4") |
| 53 | self.plugin_stores = {} |
| 54 | self.vector_store = self._load_or_create_vector_store() |
| 55 | |
| 56 | def _load_or_create_vector_store(self): |
| 57 | """Loads existing or creates new ChromaDB vector stores. |
| 58 | |
| 59 | Creates/loads vector stores for both Manim core documentation and any available plugins. |
| 60 | Stores are persisted to disk for future reuse. |
| 61 | |
| 62 | Returns: |
| 63 | Chroma: The core Manim vector store instance |
| 64 | """ |
| 65 | print("Entering _load_or_create_vector_store with trace_id:", self.trace_id) |
| 66 | core_path = os.path.join(self.chroma_db_path, "manim_core") |
| 67 | |
| 68 | # Load or create core vector store |
| 69 | if os.path.exists(core_path): |
| 70 | print("Loading existing core ChromaDB...") |
| 71 | self.core_vector_store = Chroma( |
| 72 | collection_name="manim_core", |
| 73 | persist_directory=core_path, |
| 74 | embedding_function=self._get_embedding_function() |
| 75 | ) |
| 76 | else: |
| 77 | print("Creating new core ChromaDB...") |
| 78 | self.core_vector_store = self._create_core_store() |