app.py

import os
from typing import Optional, Tuple

import gradio as gr
import pandas as pd
from buster.completers import Completion
from buster.utils import extract_zip

from embed_docs import crawl_and_embed_docs
import cfg
from cfg import setup_buster

# Typehint for chatbot history
ChatHistory = list[list[Optional[str], Optional[str]]]


# Because this is a one-click deploy app, we will be relying on env. variables being set
openai_api_key = os.getenv("OPENAI_API_KEY")  # Mandatory for app to work


if os.path.exists("outputs.zip"):
    print("Found outputs.zip, Skipping crawl and embed.")
    extract_zip("outputs.zip", output_path="outputs")

else:
    readthedocs_url = os.getenv("READTHEDOCS_URL")  # Mandatory for app to work as intended
    readthedocs_version = os.getenv("READTHEDOCS_VERSION")

    if openai_api_key is None:
        print(
            "Warning: No OPENAI_API_KEY detected. Set it with 'export OPENAI_API_KEY=sk-...'."
        )

    if readthedocs_url is None:
        raise ValueError(
            "No READTHEDOCS_URL detected. Set it with e.g. 'export READTHEDOCS_URL=https://orion.readthedocs.io/'"
        )

    if readthedocs_version is None:
        print(
            """
        Warning: No READTHEDOCS_VERSION detected. If multiple versions of the docs exist, they will all be scraped.
        Set it with e.g. 'export READTHEDOCS_VERSION=en/stable'
        """
        )


    # scrape and embed content from readthedocs website
    crawl_and_embed_docs(
        homepage_url=readthedocs_url,
        save_directory="outputs",  # Expected to be in outputs/ by buster cfg
        target_version=readthedocs_version,
    )

# Setup RAG agent
buster = setup_buster(cfg.buster_cfg)


# Setup Gradio app
def add_user_question(
    user_question: str, chat_history: Optional[ChatHistory] = None
) -> ChatHistory:
    """Adds a user's question to the chat history.

    If no history is provided, the first element of the history will be the user conversation.
    """
    if chat_history is None:
        chat_history = []
    chat_history.append([user_question, None])
    return chat_history


def format_sources(matched_documents: pd.DataFrame) -> str:
    if len(matched_documents) == 0:
        return ""

    matched_documents.similarity_to_answer = (
        matched_documents.similarity_to_answer * 100
    )

    # drop duplicate pages (by title), keep highest ranking ones
    matched_documents = matched_documents.sort_values(
        "similarity_to_answer", ascending=False
    ).drop_duplicates("title", keep="first")

    documents_answer_template: str = "📝 Here are the sources I used to answer your question:\n\n{documents}\n\n{footnote}"
    document_template: str = "[🔗 {document.title}]({document.url}), relevance: {document.similarity_to_answer:2.1f} %"

    documents = "\n".join(
        [
            document_template.format(document=document)
            for _, document in matched_documents.iterrows()
        ]
    )
    footnote: str = "I'm a bot 🤖 and not always perfect."

    return documents_answer_template.format(documents=documents, footnote=footnote)


def add_sources(history, completion):
    if completion.answer_relevant:
        formatted_sources = format_sources(completion.matched_documents)
        history.append([None, formatted_sources])

    return history


def chat(chat_history: ChatHistory) -> Tuple[ChatHistory, Completion]:
    """Answer a user's question using retrieval augmented generation."""

    # We assume that the question is the user's last interaction
    user_input = chat_history[-1][0]

    # Do retrieval + augmented generation with buster
    completion = buster.process_input(user_input)

    # Stream tokens one at a time to the user
    chat_history[-1][1] = ""
    for token in completion.answer_generator:
        chat_history[-1][1] += token

        yield chat_history, completion


demo = gr.Blocks()
with demo:
    with gr.Row():
        gr.Markdown("<h1><center>RAGTheDocs - docs.mila.quebec </center></h1>")

    gr.Markdown(
        """
        ## About
        RAGTheDocs allows you to ask questions found on the docs.mila.quebec website.

        Try it out by asking a question below about [mila docs](https://docs.mila.quebec/).

        ## How it works
        This app uses [Buster 🤖](https://github.com/jerpint/buster) and ChatGPT to search the docs for relevant info and
        answer questions.
        View the code on the [project homepage](https://github.com/jerpint/RAGTheDocs)
        """
    )

    chatbot = gr.Chatbot()

    with gr.Row():
        question = gr.Textbox(
            label="What's your question?",
            placeholder="Type your question here...",
            lines=1,
        )
        submit = gr.Button(value="Send", variant="secondary")

    examples = gr.Examples(
        examples=[
            "How can I request a job with multiple GPUs?",
            "Where should I store large datasets?",
            "how can i view my GPU usage?",
        ],
        inputs=question,
    )

    response = gr.State()

    # fmt: off
    gr.on(
        triggers=[submit.click, question.submit],
        fn=add_user_question,
        inputs=[question],
        outputs=[chatbot]
    ).then(
        chat,
        inputs=[chatbot],
        outputs=[chatbot, response]
    ).then(
        add_sources,
        inputs=[chatbot, response],
        outputs=[chatbot]
    )


demo.launch()