ui.py

import os

import gradio as gr

from pdfparser import ask_question_all, process_file

PASSWORD = os.environ["OPEN_PASSWORD"]


def add_text(history, text):
    history = history + [(text, None)]
    return history, gr.update(value="", interactive=False)


def add_file(history, file):
    history = history + [((file.name,), None)]
    return history


def bot(history):
    if history[0][0] == PASSWORD:
        if len(history) == 1:
            response = "Access granted."
        else:
            response = ask_question_all(history[1:])
    else:
        response = "Wrong password"
    history[-1][1] = response
    return history


def bot_upload(history):
    if history[0][0] == PASSWORD:
        process_file(history[-1][0][0])
        history[-1][1] = "Ready."
    else:
        history[-1][1] = "Wrong password"
    return history


with gr.Blocks() as demo:
    chatbot = gr.Chatbot([], elem_id="chatbot").style(height=450)

    with gr.Row():
        with gr.Column(scale=0.85):
            txt = gr.Textbox(
                show_label=False,
                placeholder="First upload a pdf file, then query it",
            ).style(container=False)
        with gr.Column(scale=0.15, min_width=0):
            btn = gr.UploadButton("📁", file_types=["pdf"])

    txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
        bot, chatbot, chatbot
    )
    txt_msg.then(lambda: gr.update(interactive=True), None, [txt], queue=False)
    file_msg = btn.upload(add_file, [chatbot, btn], [chatbot], queue=False).then(
        bot_upload, chatbot, chatbot
    )

demo.launch()