Class for integrating RAG (Retrieval Augmented Generation) functionality. This class handles RAG integration including plugin detection, query generation, and document retrieval. Args: helper_model: Model used for generating queries and processing text output_dir (str):
| 15 | from src.rag.vector_store import RAGVectorStore |
| 16 | |
| 17 | class RAGIntegration: |
| 18 | """Class for integrating RAG (Retrieval Augmented Generation) functionality. |
| 19 | |
| 20 | This class handles RAG integration including plugin detection, query generation, |
| 21 | and document retrieval. |
| 22 | |
| 23 | Args: |
| 24 | helper_model: Model used for generating queries and processing text |
| 25 | output_dir (str): Directory for output files |
| 26 | chroma_db_path (str): Path to ChromaDB |
| 27 | manim_docs_path (str): Path to Manim documentation |
| 28 | embedding_model (str): Name of embedding model to use |
| 29 | use_langfuse (bool, optional): Whether to use Langfuse logging. Defaults to True |
| 30 | session_id (str, optional): Session identifier. Defaults to None |
| 31 | """ |
| 32 | |
| 33 | def __init__(self, helper_model, output_dir, chroma_db_path, manim_docs_path, embedding_model, use_langfuse=True, session_id=None): |
| 34 | self.helper_model = helper_model |
| 35 | self.output_dir = output_dir |
| 36 | self.manim_docs_path = manim_docs_path |
| 37 | self.session_id = session_id |
| 38 | self.relevant_plugins = None |
| 39 | |
| 40 | self.vector_store = RAGVectorStore( |
| 41 | chroma_db_path=chroma_db_path, |
| 42 | manim_docs_path=manim_docs_path, |
| 43 | embedding_model=embedding_model, |
| 44 | session_id=self.session_id, |
| 45 | use_langfuse=use_langfuse, |
| 46 | helper_model=helper_model |
| 47 | ) |
| 48 | |
| 49 | def set_relevant_plugins(self, plugins: List[str]) -> None: |
| 50 | """Set the relevant plugins for the current video. |
| 51 | |
| 52 | Args: |
| 53 | plugins (List[str]): List of plugin names to set as relevant |
| 54 | """ |
| 55 | self.relevant_plugins = plugins |
| 56 | |
| 57 | def detect_relevant_plugins(self, topic: str, description: str) -> List[str]: |
| 58 | """Detect which plugins might be relevant based on topic and description. |
| 59 | |
| 60 | Args: |
| 61 | topic (str): Topic of the video |
| 62 | description (str): Description of the video content |
| 63 | |
| 64 | Returns: |
| 65 | List[str]: List of detected relevant plugin names |
| 66 | """ |
| 67 | # Load plugin descriptions |
| 68 | plugins = self._load_plugin_descriptions() |
| 69 | if not plugins: |
| 70 | return [] |
| 71 | |
| 72 | # Get formatted prompt using the task_generator function |
| 73 | prompt = get_prompt_detect_plugins( |
| 74 | topic=topic, |