app.py

import os
import gradio as gr
from gradio.components import Textbox, Button, Slider, Checkbox
from AinaTheme import theme
from urllib.error import HTTPError

from rag import RAG
from utils import setup

MAX_NEW_TOKENS = 700
SHOW_MODEL_PARAMETERS_IN_UI = os.environ.get("SHOW_MODEL_PARAMETERS_IN_UI", default="True") == "True"

setup()


rag = RAG(embeddings_model=os.getenv("EMBEDDINGS"))


def eadop_rag(prompt, num_chunks):
    model_parameters = {"NUM_CHUNKS": num_chunks}
    try:
        retrun rag.get_context(prompt, model_parameters)
    except HTTPError as err:
        if err.code == 400:
            gr.Warning(
                "The inference endpoint is only available Monday through Friday, from 08:00 to 20:00 CET."
            )
        return None, None, None
    except:
        gr.Warning(
            "Inference endpoint is not available right now. Please try again later."
        )
        return None, None, None


def clear():
    return (
        None, 
        None,
        None,
        None,
        gr.Slider(value=2.0),
    )


def gradio_app():
    with gr.Blocks(theme=theme) as demo:
        with gr.Row(equal_height=True):
            with gr.Column(variant="panel"):
                input_ = Textbox(
                    lines=11,
                    label="Input",
                    placeholder="Quina és la finalitat del Servei Meteorològic de Catalunya?",
                )
                with gr.Row(variant="panel"):
                    clear_btn = Button(
                        "Clear",
                    )
                    submit_btn = Button("Submit", variant="primary", interactive=False)

                with gr.Row(variant="panel"):
                    with gr.Accordion("Model parameters", open=False, visible=SHOW_MODEL_PARAMETERS_IN_UI):
                        num_chunks = Slider(
                            minimum=1,
                            maximum=6,
                            step=1,
                            value=2,
                            label="Number of chunks"
                        )

            with gr.Column(variant="panel"):
                output = Textbox(
                    lines=10, 
                    label="Context", 
                    interactive=False, 
                    show_copy_button=True
                )
                with gr.Accordion("Sources and context:", open=False):
                    source_context = gr.Markdown(
                        label="Sources",
                        show_label=False,
                    )
                    with gr.Accordion("See full context evaluation:", open=False):
                        context_evaluation = gr.Markdown(
                            label="Full context",
                            show_label=False,
                            # interactive=False, 
                            # autoscroll=False,
                            # show_copy_button=True
                        )
                

        input_.change(
            fn=change_interactive,
            inputs=[input_],
            outputs=[clear_btn, submit_btn],
            api_name=False,
        )

        input_.change(
            fn=None,
            inputs=[input_],
            api_name=False,
            js="""(i, m) => {
            document.getElementById('inputlength').textContent = i.length + '  '
            document.getElementById('inputlength').style.color =  (i.length > m) ? "#ef4444" : "";
        }""",
        )

        clear_btn.click(
            fn=clear, 
            inputs=[], 
            outputs=[input_, output, source_context, context_evaluation, num_chunks],
              queue=False, 
              api_name=False
        )
        
        submit_btn.click(
            fn=submit_input, 
            inputs=[input_, num_chunks], 
            outputs=[output, source_context, context_evaluation],
            api_name="get-eadop-rag"
        )

        demo.launch(show_api=True)


if __name__ == "__main__":
    gradio_app()