app.py

import os

import gradio as gr
import pixeltable as pxt
from pixeltable.functions import openai
from pixeltable.functions.document import document_splitter
from pixeltable.functions.huggingface import sentence_transformer

from udfs import create_messages, create_prompt

# The OpenAI API key is provided via the OPENAI_API_KEY Space secret; never prompt on stdin.
OPENAI_API_KEY_SET = bool(os.environ.get('OPENAI_API_KEY'))


def process_files(pdf_files, chunk_limit, chunk_separator):
    if not OPENAI_API_KEY_SET:
        return 'OPENAI_API_KEY is not configured. Add it under Settings > Secrets in the Space and restart.'

    # Initialize Pixeltable
    pxt.drop_dir('chatbot_demo', force=True)
    pxt.create_dir('chatbot_demo')

    # Create a table to store the uploaded PDF documents
    t = pxt.create_table(
        'chatbot_demo.documents',
        {
            'document': pxt.Document,
            'question': pxt.String
        }
    )

    # Insert the PDF files into the documents table
    if pdf_files:
        t.insert({'document': file} for file in pdf_files if file.endswith('.pdf'))

    # Create a view that splits the documents into smaller chunks
    chunks_t = pxt.create_view(
        'chatbot_demo.chunks',
        t,
        iterator=document_splitter(
            document=t.document,
            separators=chunk_separator,
            limit=chunk_limit if chunk_separator in ["token_limit", "char_limit"] else None,
            metadata='title,heading,sourceline'
        )
    )

    chunks_t.add_embedding_index(
        'text',
        string_embed=sentence_transformer.using(model_id='intfloat/e5-large-v2')
    )

    @pxt.query
    def top_k(query_text: str):
        sim = chunks_t.text.similarity(string=query_text)
        return (
            chunks_t.order_by(sim, asc=False)
                .select(chunks_t.text, sim=sim)
                .limit(5)
        )

    # Add computed columns using keyword argument syntax
    t.add_computed_column(question_context=top_k(t.question))
    t.add_computed_column(prompt=create_prompt(t.question_context, t.question))
    t.add_computed_column(messages=create_messages(t.prompt))

    # Add the response column using the messages computed column
    t.add_computed_column(response=openai.chat_completions(
        model='gpt-4o-mini',
        messages=t.messages,
        model_kwargs={'max_tokens': 300, 'top_p': 0.9, 'temperature': 0.7}
    ))
    t.add_computed_column(gpt4omini=t.response.choices[0].message.content)

    return "Files processed successfully. You can start the discussion."


def get_answer(msg):
    t = pxt.get_table('chatbot_demo.documents')

    # Insert the question into the table
    t.insert([{'question': msg}])

    # The answer will be automatically generated through the chain of computed columns
    answer = t.select(t.gpt4omini).where(t.question == msg).collect()['gpt4omini'][0]
    return answer


def respond(message, chat_history):
    bot_message = get_answer(message)
    chat_history.append({'role': 'user', 'content': message})
    chat_history.append({'role': 'assistant', 'content': bot_message})
    return "", chat_history


# Gradio interface
with gr.Blocks(theme=gr.themes.Base()) as demo:
    gr.Markdown(
        """
        <div>
            <img src="https://raw.githubusercontent.com/pixeltable/pixeltable/main/docs/resources/pixeltable-logo-large.png" alt="Pixeltable" style="max-width: 200px; margin-bottom: 20px;" />
            <h1 style="margin-bottom: 0.5em;">AI Chatbot With Retrieval-Augmented Generation (RAG)</h1>
        </div>
        """
    )
    gr.HTML(
        """
        <p>
            <a href="https://github.com/pixeltable/pixeltable" target="_blank" style="color: #F25022; text-decoration: none; font-weight: bold;">Pixeltable</a> is a declarative interface for working with text, images, embeddings, and even video, enabling you to store, transform, index, and iterate on data.
        </p>

        <div style="background-color: #E5DDD4; border: 1px solid #e9ecef; color: #000000; border-radius: 8px; padding: 15px; margin-bottom: 20px;">
            <strong style="color: #000000">Disclaimer:</strong> This app is best run on your own hardware with a GPU for optimal performance. This Hugging Face Space uses the free tier (2vCPU, 16GB RAM), which results in slower processing times. If you wish to use this app with your own hardware for improved performance, you can <a href="https://huggingface.co/spaces/Pixeltable/AI-Chatbot-With-Retrieval-Augmented-Generation?duplicate=true" target="_blank" style="color: #4D148C; text-decoration: none; font-weight: bold;">duplicate this Hugging Face Space</a>, run it locally, or use Google Colab with the Free limited GPU support.
        </div>
        """
    )

    if not OPENAI_API_KEY_SET:
        gr.Markdown(
            '**`OPENAI_API_KEY` is not set.** Add it as a Space secret (Settings > Secrets) '
            'and restart the Space to enable document processing and chat.'
        )

    with gr.Row():
        with gr.Column():
           with gr.Accordion("What This Demo Does", open = True):
            gr.Markdown("""
            - Upload multiple PDF documents.
            - Process and index the content of these documents.
            - Ask questions about the content and Receive AI-generated answers that are grounded.
         """)
        with gr.Column():
          with gr.Accordion("How does it work?", open = True):
            gr.Markdown("""
            - When a user asks a question, the system searches for the most relevant chunks of text from the uploaded documents.
            - It then uses these relevant chunks as context for a large language model (LLM) to generate an answer.
            - The LLM formulates a response based on the provided context and the user's question.
          """)

    with gr.Row():
        with gr.Column(scale=1):
            pdf_files = gr.File(label="Upload PDF Documents", file_count="multiple", file_types=['.pdf'])
            chunk_limit = gr.Slider(minimum=100, maximum=500, value=300, step=5, label="Chunk Size Limit")
            chunk_separator = gr.Dropdown(
                choices=["token_limit", "char_limit", "sentence", "paragraph", "heading"],
                value="token_limit",
                label="Chunk Separator"
            )
            process_button = gr.Button("Process Files")
            process_output = gr.Textbox(label="Processing Output")

        with gr.Column(scale=2):
            chatbot = gr.Chatbot(label="Chat History", type='messages')
            msg = gr.Textbox(label="Your Question", placeholder="Ask a question about the uploaded documents")
            submit = gr.Button("Submit")

    process_button.click(process_files, inputs=[pdf_files, chunk_limit, chunk_separator], outputs=[process_output])
    submit.click(respond, inputs=[msg, chatbot], outputs=[msg, chatbot])

if __name__ == "__main__":
    demo.launch()