""" app.py

Question / answer over a collection of PDF documents using late interaction
ColBERT model for retrieval and DSPy+Mistral for answer generation.

:author: Didier Guillevic
:date: 2024-12-22
"""

import gradio as gr

import logging
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

import os
import pdf_utils # utilities for pdf processing
import colbert_utils # utilities for to build a ColBERT retrieval model
import dspy_utils # utilities for building a DSPy based retrieval generation model

from tqdm.notebook import tqdm
import warnings
warnings.filterwarnings('ignore')


def generate_response(question: str) -> list[str, str, str]:
    """Generate a response to a given question using the RAG model.
    
    """
    global dspy_rag_model

    if dspy_rag_model is None:
        return "RAG model not built. Please build the model first."

    # Generate response
    responses, references, snippets = dspy_rag_model.generate_response(
        question=question, k=5, method='chain_of_thought')
    
    return responses, references, snippets


with gr.Blocks() as demo:
    gr.Markdown("""
        # Retrieval (ColBERT) + Generation (DSPy & Mistral)
        Note: building the retrieval model might take a few minutes.
    """)

    # Input files and build status
    with gr.Row():
        upload_files = gr.File(
            label="Upload PDF files to index", file_count="multiple",
            value=["OECD_Engaging_with_HNW_individuals_tax_compliance_(2009).pdf",],
            scale=5)
        build_status = gr.Textbox(label="Build status", placeholder="", scale=2)
    
    # button
    build_button = gr.Button("Build retrieval generation model", variant='primary')
    
    # Question to answer
    question = gr.Textbox(
        label="Question to answer",
        placeholder="How do tax administrations address aggressive tax planning by HNWIs?"
    )
    response = gr.Textbox(
        label="Response",
        placeholder=""
    )
    with gr.Accordion("References & snippets", open=False):
        references = gr.HTML(label="References")
        snippets = gr.HTML(label="Snippets")
    
    # button
    response_button = gr.Button("Submit", variant='primary')
    
    # Example questions given default provided PDF file
    with gr.Accordion("Sample questions", open=False):
        gr.Examples(
            [
                ["What are the tax risks associated with high net worth individuals (HNWIs)?",],
                ["How do tax administrations address aggressive tax planning by HNWIs?",],
                ["How can tax administrations engage with HNWIs to improve tax compliance?",],
                ["What are the benefits of establishing dedicated HNWI units within tax administrations?",],
                ["How can international cooperation help address offshore tax risks associated with HNWIs?",],
            ],
            inputs=[question,],
            outputs=[response, references, snippets],
            fn=generate_response,
            cache_examples=False,
            label="Sample questions"
        )
    
    # Documentation
    with gr.Accordion("Documentation", open=False):
        gr.Markdown("""
            - What
                - Retrieval augmented generation (RAG) model based on ColBERT and DSPy.
                - Retrieval base model:  'antoinelouis/colbert-xm' (multilingual model)
                - Generation framework: DSPy and Mistral.
            - How
                - Upload PDF files to index.
                - Build the retrieval augmented model (might take a few minutes)
                - Ask a question to generate a response.
        """)

    # Click actions
    build_button.click(
        fn=build_rag_model,
        inputs=[upload_files],
        outputs=[build_status]
    )
    response_button.click(
        fn=generate_response,
        inputs=[question],
        outputs=[response, references, snippets]
    )


demo.launch(show_api=False)