diff --git a/src/RAG_pipeline.py b/src/RAG_pipeline.py new file mode 100644 index 0000000..5bee608 --- /dev/null +++ b/src/RAG_pipeline.py @@ -0,0 +1,49 @@ +from typing import List, Optional, Dict +from dataclasses import dataclass +from .interface.base_datastore import BaseDatastore +from .interface.base_evaluator import BaseEvaluator +from .interface.base_indexer import BaseIndexer +from .interface.base_response_generator import BaseResponseGenerator +from .interface.base_retriever import BaseRetriever + +@dataclass +class RAGpipeline: + + indexer: BaseIndexer + datastore: BaseDatastore + retriever: BaseRetriever + response_generator: BaseResponseGenerator + evaluator: Optional[BaseEvaluator] = None + + def reset(self) -> None: + self.datastore.reset_table() + + def add_documents(self, documents: List[str]) -> None: + items = self.indexer.index(documents) + self.datastore.add_items(items= items) + return + + def process_query(self, query: str, top_k: int = 5, source: bool = False) -> str: + """Génère la réponse à la requête à partir du contexte récupéré. + Affiche éventuellement les sources. + Args: + query (str): Requête + top_k (int, optional): Nombre de documents à renvoyer. Defaults to 5. + source (bool, optional): Afficher ou non les sources. Default to false + Returns: + str: Réponse générée + """ + context = self.retriever.search_retriever(query, top_k= top_k) + + if( source): + for i, doc in enumerate(context): + print(f"Based on the document {i+1}: {doc} \n") + + response = self.response_generator.generate_response(query= query, context= context) + return response + + def evaluate(self) -> None: + """A compléter lors de l'ajout de l'évaluateur + """ + print("Function not completed.") + pass \ No newline at end of file