app.py

import asyncio
import os.path
import tempfile
import uuid
from typing import List

import gradio
import gradio as gr
import openai
import pandas as pd
from autorag.evaluator import Evaluator

from src.data.chunk import chunk
from src.data.parse import parse_pdf
from src.runner import GradioStreamRunner

root_dir = os.path.dirname(os.path.realpath(__file__))

pseudo_trial_yaml_path = os.path.join(root_dir, "config", "init_project_for_pseudo_trial.yaml")
init_run_yaml = os.path.join(root_dir, "config", "init_project_for_run.yaml")

gradio_runner = None

# Code for Task 1
def file_ingest(input_files: List[str], temp_project_dir, progress=gr.Progress()):
    if os.getenv("OPENAI_API_KEY") is None:
        return "Please submit your OpenAI API key first."
    if not input_files:
        return "Please upload a file first."
    progress(0.05)
    # do parse
    raw_df = parse_pdf(file_lists=input_files)
    progress(0.3)
    # do chunk
    corpus_df = chunk(raw_df, method="recursivecharacter",
                      lang="en", chunk_size=512, chunk_overlap=128)
    progress(0.5)
    asyncio.sleep(0.5)

    # Logic for button click
    empty_qa_df = make_empty_qa(corpus_df=corpus_df)
    with tempfile.TemporaryDirectory() as temp_data_dir:
        empty_qa_df.to_parquet(os.path.join(temp_data_dir, "empty_qa.parquet"))
        corpus_df.to_parquet(os.path.join(temp_data_dir, "corpus.parquet"))

        evaluator = Evaluator(qa_data_path=os.path.join(temp_data_dir, "empty_qa.parquet"),
                              corpus_data_path=os.path.join(temp_data_dir, "corpus.parquet"),
                              project_dir=temp_project_dir)
        evaluator.start_trial(pseudo_trial_yaml_path, skip_validation=True)
        yield "Setting up"
        progress(0.9)
        set_runner(temp_project_dir)
        progress(1.0)
        yield "File uploaded complete. You can use it at chatbot now."


def make_empty_qa(corpus_df: pd.DataFrame):
    doc_id = corpus_df["doc_id"].iloc[0]
    return pd.DataFrame({
        "qid": str(uuid.uuid4()),
        "query": ["Who is Kai Havertz?"],
        "retrieval_gt": [[[doc_id]]],
        "generation_gt": [["Havertz is the greatest footballer."]],
    })


def on_submit_openai_key(openai_key):
    os.environ["OPENAI_API_KEY"] = openai_key
    # Test openai key
    try:
        client = openai.OpenAI()
        response = client.chat.completions.create(
            messages=[
                {"role": "user", "content": "What is the capital of France?"},
            ],
            model="gpt-4o-mini",
            max_tokens=3,
        )
        assert isinstance(response.choices[0].message.content, str)
        gr.Info("OpenAI API key submitted.", duration=3)
        return "Setting complete."
    except openai.AuthenticationError as e:
        gr.Error("OpenAI API key is invalid.", duration=3)
        return "Not Set"
    except AssertionError as e:
        gr.Error("OpenAI server is not working properly.", duration=3)
        return "Not Set"


def set_runner(project_dir):
    runner = GradioStreamRunner.from_yaml(yaml_path=init_run_yaml, project_dir=project_dir)
    global gradio_runner
    gradio_runner = runner


def get_response(message, history):
    global gradio_runner
    if gradio_runner is None:
        gradio.Warning("Please set the AutoRAG server first.")
        return
    if os.getenv("OPENAI_API_KEY", None) is None:
        gradio.Warning("Please submit your OpenAI API key first.")
        return

    for output in gradio_runner.stream_run(message):
        yield output[0]


# interface one
with gr.Blocks(theme="earneleh/paris") as demo:
    with tempfile.TemporaryDirectory() as project_dir:
        # Define components
        with gr.Row():
            with gr.Column(scale=3):
                textbox = gr.Textbox(label="Please input your OpenAI API key and press Enter.", type="password",
                         info="You can get your API key from https://platform.openai.com/account/api-keys\n"
                         "AutoRAG do not store your API key.",
                                     autofocus=True)
                api_key_status_box = gr.Textbox(label="OpenAI API status", value="Not Set", interactive=False)

                gr.Markdown("## Ingest Your Data")

                file_input = gr.File(label="Upload Files", type="filepath", file_count="multiple")
                button = gr.Button("Submit file")
                text_output = gr.Textbox(label="Status update", interactive=False)

                # Define layout and interactions
                textbox.submit(on_submit_openai_key, inputs=[textbox], outputs=api_key_status_box)
                button.click(file_ingest, inputs=[file_input, gr.State(project_dir)], outputs=[text_output])

            with gr.Column(scale=7):
                gr.ChatInterface(
                    get_response, title="This is your Naive RAG Chatbot 🚀", retry_btn=None, undo_btn=None,
                )

        gr.Markdown("## Do you like the result?\n\nIf you don't like it, try to optimize it with AutoRAG. Press below button and go to make evaluation data and optimize it. Both on the Huggingface space so you don't need to install anything.")
        with gr.Row():
            open_data_creation = gr.Button(value="1️⃣ : Data Creation",
                                           link="https://huggingface.co/spaces/AutoRAG/AutoRAG-data-creation")
            open_optimize = gr.Button(value="2️⃣ : Optimize", link="https://www.auto-rag.com/")


demo.launch(share=False, debug=True)