app.py

from pathlib import Path
# import base64
import gradio as gr
from src.rag import Rag
from src.amodel import ModelType

STORE_DIR = "./files/rag_app" # Le répertoire de la base
MAX_DOCS = 4 # Le nombre max de documents dans la base

def main():
    # Création du rag
    rag:Rag = Rag(ModelType.MTHUGGINGFACE, store_dir=STORE_DIR) 
    # Reset de la base à chaque démarrage du serveur
    rag.reset_store()

    # UI
    with gr.Blocks() as demo:

        def upload_file(file_path):
            # Récupérer la liste des collections du store
            names = list(rag.emb_store.get_collection_names())
            # Supprimer la première s'il y en a déjà 4 (le max)
            if len(names) == MAX_DOCS:
                rag.delete_collection(names[0])
            # Ajouter le pdf
            name:str = Path(file_path).name
            rag.add_pdf_to_store(file_name=file_path, collection_name=name)
            # Mettre à jour la liste des noms de collections
            names = list(rag.emb_store.get_collection_names())
            combo.choices = names
            return gr.update(choices=names, value=name, interactive=True), gr.update(value="")

        def ask_rag(question:str, choice:str):
            col_name:str = choice
            if col_name == None:
                return "Aucun pdf actif, veuillez en uploader un !"
            prompt, resp, sources, ids = rag.ask_rag(question, col_name)
            # print("choice:" + col_name, "RAG:" + resp)
            return resp
            
        def on_temperature_change(temp):
            rag.set_temperature(temp)
        
        
        with gr.Tab("RAG"):
            # with gr.Row():
            gr.Image("./files/drane.jpg", height=100, show_download_button=False, show_fullscreen_button=False, show_label=False)
            # Le combo qui permet le choix du pdf
            names:list[str] = rag.emb_store.get_collection_names()
            combo = gr.Dropdown(names, label="PDFs", multiselect=False)
            # Le button qui permet d'uploader un pdf
            upload_button = gr.UploadButton("Clique pour uploader un pdf", file_types=[".pdf"], file_count="single")
                    
            # La zone où on pose une question au RAG
            ask_input = gr.Text(label="Pose une question à ton pdf")
            # La réponse du RAG (Markdown pour afficher les formules .tex)
            rag_output = gr.Markdown(label="Réponse")
        
        with gr.Tab("Réglages"):
            gr.Markdown("## Modèles:")
            gr.Markdown("- " + rag.get_llm_name())
            gr.Markdown("- " + rag.get_feature_name())
            temperature_slider = gr.Slider(minimum=0,
                        maximum=1.0,
                        value=0.5,
                        step=0.1,
                        label="Température")

        # Réponses aux évènements
        upload_button.upload(fn=upload_file, inputs=upload_button, outputs=[combo, rag_output], show_progress="full")
        ask_input.submit(fn=ask_rag, inputs=[ask_input, combo], outputs=rag_output)
        temperature_slider.change(fn=on_temperature_change, inputs=temperature_slider)

    # demo.launch(allowed_paths=["./file/"], share=True)
    demo.launch(allowed_paths=["./file/"])

if __name__ == "__main__":
    main()