jerpint/RAGTheDocs-mila-qc

1import os2from typing import Optional, Tuple3 4import gradio as gr5import pandas as pd6from buster.completers import Completion7from buster.utils import extract_zip8 9from embed_docs import crawl_and_embed_docs10import cfg11from cfg import setup_buster12 13# Typehint for chatbot history14ChatHistory = list[list[Optional[str], Optional[str]]]15 16 17# Because this is a one-click deploy app, we will be relying on env. variables being set18openai_api_key = os.getenv("OPENAI_API_KEY")  # Mandatory for app to work19 20 21if os.path.exists("outputs.zip"):22    print("Found outputs.zip, Skipping crawl and embed.")23    extract_zip("outputs.zip", output_path="outputs")24 25else:26    readthedocs_url = os.getenv("READTHEDOCS_URL")  # Mandatory for app to work as intended27    readthedocs_version = os.getenv("READTHEDOCS_VERSION")28 29    if openai_api_key is None:30        print(31            "Warning: No OPENAI_API_KEY detected. Set it with 'export OPENAI_API_KEY=sk-...'."32        )33 34    if readthedocs_url is None:35        raise ValueError(36            "No READTHEDOCS_URL detected. Set it with e.g. 'export READTHEDOCS_URL=https://orion.readthedocs.io/'"37        )38 39    if readthedocs_version is None:40        print(41            """42        Warning: No READTHEDOCS_VERSION detected. If multiple versions of the docs exist, they will all be scraped.43        Set it with e.g. 'export READTHEDOCS_VERSION=en/stable'44        """45        )46 47 48    # scrape and embed content from readthedocs website49    crawl_and_embed_docs(50        homepage_url=readthedocs_url,51        save_directory="outputs",  # Expected to be in outputs/ by buster cfg52        target_version=readthedocs_version,53    )54 55# Setup RAG agent56buster = setup_buster(cfg.buster_cfg)57 58 59# Setup Gradio app60def add_user_question(61    user_question: str, chat_history: Optional[ChatHistory] = None62) -> ChatHistory:63    """Adds a user's question to the chat history.64 65    If no history is provided, the first element of the history will be the user conversation.66    """67    if chat_history is None:68        chat_history = []69    chat_history.append([user_question, None])70    return chat_history71 72 73def format_sources(matched_documents: pd.DataFrame) -> str:74    if len(matched_documents) == 0:75        return ""76 77    matched_documents.similarity_to_answer = (78        matched_documents.similarity_to_answer * 10079    )80 81    # drop duplicate pages (by title), keep highest ranking ones82    matched_documents = matched_documents.sort_values(83        "similarity_to_answer", ascending=False84    ).drop_duplicates("title", keep="first")85 86    documents_answer_template: str = "📝 Here are the sources I used to answer your question:\n\n{documents}\n\n{footnote}"87    document_template: str = "[🔗 {document.title}]({document.url}), relevance: {document.similarity_to_answer:2.1f} %"88 89    documents = "\n".join(90        [91            document_template.format(document=document)92            for _, document in matched_documents.iterrows()93        ]94    )95    footnote: str = "I'm a bot 🤖 and not always perfect."96 97    return documents_answer_template.format(documents=documents, footnote=footnote)98 99 100def add_sources(history, completion):101    if completion.answer_relevant:102        formatted_sources = format_sources(completion.matched_documents)103        history.append([None, formatted_sources])104 105    return history106 107 108def chat(chat_history: ChatHistory) -> Tuple[ChatHistory, Completion]:109    """Answer a user's question using retrieval augmented generation."""110 111    # We assume that the question is the user's last interaction112    user_input = chat_history[-1][0]113 114    # Do retrieval + augmented generation with buster115    completion = buster.process_input(user_input)116 117    # Stream tokens one at a time to the user118    chat_history[-1][1] = ""119    for token in completion.answer_generator:120        chat_history[-1][1] += token121 122        yield chat_history, completion123 124 125demo = gr.Blocks()126with demo:127    with gr.Row():128        gr.Markdown("<h1><center>RAGTheDocs - docs.mila.quebec </center></h1>")129 130    gr.Markdown(131        """132        ## About133        RAGTheDocs allows you to ask questions found on the docs.mila.quebec website.134 135        Try it out by asking a question below about [mila docs](https://docs.mila.quebec/).136 137        ## How it works138        This app uses [Buster 🤖](https://github.com/jerpint/buster) and ChatGPT to search the docs for relevant info and139        answer questions.140        View the code on the [project homepage](https://github.com/jerpint/RAGTheDocs)141        """142    )143 144    chatbot = gr.Chatbot()145 146    with gr.Row():147        question = gr.Textbox(148            label="What's your question?",149            placeholder="Type your question here...",150            lines=1,151        )152        submit = gr.Button(value="Send", variant="secondary")153 154    examples = gr.Examples(155        examples=[156            "How can I request a job with multiple GPUs?",157            "Where should I store large datasets?",158            "how can i view my GPU usage?",159        ],160        inputs=question,161    )162 163    response = gr.State()164 165    # fmt: off166    gr.on(167        triggers=[submit.click, question.submit],168        fn=add_user_question,169        inputs=[question],170        outputs=[chatbot]171    ).then(172        chat,173        inputs=[chatbot],174        outputs=[chatbot, response]175    ).then(176        add_sources,177        inputs=[chatbot, response],178        outputs=[chatbot]179    )180 181 182demo.launch()183