src/rag.py

from pypdf import PdfReader
from .chunker import Chunker
from .amodel import AModel
from .store import Store

CHUNK_CHAR_COUNT = 1000
CHUNK_OVERLAP = 200

class Rag:
    '''
        RAG naïf
        Classe qui s'occupe de toute la chaine du RAG.
        Elle permet :
            d'interroger un llm directement (sans RAG) avec ask_llm()
            d'interroger le RAG lui même avec ask_rag()
            d'ajouter des documents à la base de données du RAG
            de remettre la base à zéro
            de créer des vecteurs
            de charger des pdf
    '''

    # Le prompt qui sera utilisé uniquement avec le RAG
    prompt_template = """
        En vous basant **uniquement** sur les informations fournies dans le contexte 
        ci-dessous, répondez à la question posée.
        Les équations seront écrites en latex.
        Si vous ne trouvez pas la réponse dans le contexte, répondez "Je ne sais pas".
        Contexte : {context}
        Question : {question}
    """

    def __init__(self, llm:AModel, emb:AModel, store_dir:str)->None:
        '''
            Constructeur du Rag
            Args: 
                llm: le model de langage
                emb: le model d'embeddings
                store_dir: le répertoire de persistance de la base de données ou None pour éphémère
            Exception:
                Si le store ne peut pas se créer (répertoire inaccessible par ex.)
        '''
        self.llm:AModel = llm
        self.emb:AModel = emb
        self.store_dir:str = store_dir
        try:
            self.emb_store = Store(store_dir) # persistant
            # self.emb_store = Store(None) # éphémère
        except:
            raise

    def get_llm_name(self):
        return self.llm.get_llm_name()
    
    def get_feature_name(self):
        return self.emb.get_feature_name()
    
    def get_temperature(self):
        return self.llm.get_temperature()
    
    def set_temperature(self, temperature:float):
        self.llm.set_temperature(temperature)

    def reset_store(self):
        self.emb_store.reset()

    def delete_collection(self, name:str)->None:
        self.emb_store.delete_collection(name)

    def create_vectors(self, chunks:list[str])->list[list[float]]:
        '''
            Renvoie les vecteurs correspondant à 'chunks', calculés par 'emb'
            Args:
                chunks: les extraits de texte à calculer
            Return:
                la liste des vecteurs calculés
        '''        
        vectors:list = []
        tokens:int = 0
        try:
            vectors:list[list[float]] = self.emb.create_vectors(chunks) # batch si le model le permet
            return vectors
        except:
            raise

    def load_pdf(self, file_name:str)->str:
        ''' Charge le fichier 'file_name' et renvoie son contenu sous forme de texte. '''
        reader = PdfReader(file_name)
        content = ""
        for page in reader.pages:
            content += page.extract_text() + "\n"
        return content
    
    def get_chunks(self, text:str)->list:
        '''
            Découpe le 'text' en chunks de taille chunk_size avec un recouvrement
            Args:
                text: Le texte à découper
            Return:
                La liste des chunks
        '''
        chunker = Chunker()
        chunks = chunker.split_basic(text=text, char_count=CHUNK_CHAR_COUNT, overlap=CHUNK_OVERLAP)
        return chunks

    def add_pdf_to_store(self, file_name:str, collection_name:str)->None:
        '''
            Ajoute un pdf à la base de données du RAG.
            Args:
                file_name: le chemin vers le fichier à ajouter
                collection_name: Le nom de la collection dans laquelle il faut ajouter les chunks
                    La collection est créée si elle n'existe pas.
        '''
        text:str = self.load_pdf(file_name)
        chunks:list[str] = self.get_chunks(text)
        self.add_chunks_to_store(chunks=chunks, collection_name=collection_name, source=file_name)

    def add_pdf_stream_to_store(self, stream, collection_name:str)->None:
        '''
            Ajoute un stream provenant de file_uploader de streamlit par exemple
        '''
        text:str = self.load_pdf(stream)
        chunks:list[str] = self.get_chunks(text)
        self.add_chunks_to_store(chunks=chunks, collection_name=collection_name, source="stream")

    def add_chunks_to_store(self, chunks:list[str], collection_name:str, source:str)->None:
        '''
            Ajoute des chunks à la base de données du RAG.
            Args:
                chunks: les chunks à ajouter
                collection_name: Le nom de la collection dans laquelle il faut ajouter les chunks
                    La collection est créée si elle n'existe pas.
                source: la source des chunks (nom du fichier, url ...)
        ''' 
        try:       
            vectors = self.create_vectors(chunks=chunks)
        except:
            raise
        self.emb_store.add_to_collection(
            collection_name=collection_name, 
            source=source, 
            vectors=vectors,
            chunks=chunks
        )


    def ask_llm(self, question:str)->str:
        '''
            Pose une question au llm, attend sa réponse et la renvoie.
            Args:
                question: La question qu'on veut lui poser
            Returns:
                La réponse du llm
        '''
        try:
            return self.llm.ask_llm(question=question)
        except:
            return "Error while comminicating with model !"

    def ask_rag(self, question:str, collection_name:str)->tuple[str, str, list[str], list[str]]:
        '''
            Pose une question au RAG, attend sa réponse et la renvoie.
            Args:
                question: La question qu'on veut lui poser
                collection_name: le nom de la collection que l'on veut interroger
            Returns:
                Le prompt effectivement donné au llm
                La réponse du llm
                Les sources du RAG utilisées
                Les ids des documents du RAG
        '''
        if not question:
            return "", "Error: No question !", [], []
        if not collection_name:
            return "", "Error: No collection specified !", [], []
        if not collection_name in self.emb_store.get_collection_names():
            return "", "Error: {name} is no more in the database !".format(name=collection_name), [], []
        try:
            # Transformer la 'question' en vecteur avec emb_model
            query_vector:list[float] = self.emb.create_vector(question) 
            # Récupérer les chunks du store similaires à la question
            chunks, sources, ids = self.emb_store.get_similar_chunks( 
                query_vector=query_vector, 
                count=2, 
                collection_name=collection_name
            )
            # Préparer le prompt final à partir du prompt_template
            prompt:str = self.prompt_template.format( 
                context="\n\n\n".join(chunks),
                question=question
            )
            # demander au llm_model de répondre
            resp:str = self.ask_llm(question=prompt)

            return prompt, resp, sources, ids
        except:
            return "", "Error with communicating with model !", [], []