app.py

import gradio as gr
import pdfplumber
from docx import Document
import rag

READER_MODEL = rag.load_reader_model()

KNOWLEDGE_VECTOR_DATABASE = None
RAW_TEXT = ""

def extract_text(files):
    global RAW_TEXT

    if files is None:
        return ""

    RAW_TEXT = "" # Reset before processing new files
    
    for file in files:
        file_path = file.name
        
        if file_path.endswith(".txt"):
            with open(file_path, "r", encoding="utf-8") as f:
                RAW_TEXT += f.read() + "\n\n"
        elif file_path.endswith(".pdf"):
            with pdfplumber.open(file_path) as pdf:
                RAW_TEXT += "\n".join(page.extract_text() or "" for page in pdf.pages) + "\n\n"
        elif file_path.endswith(".docx"):
            RAW_TEXT += "\n".join([para.text for para in Document(file_path).paragraphs]) + "\n\n"
        else:
            RAW_TEXT = "Unsupported file type. Please upload a .pdf, .txt, or .docx file."
            return ""
    
    embed_text()
    
    return RAW_TEXT.strip()

def embed_text():
    global KNOWLEDGE_VECTOR_DATABASE, RAW_TEXT

    if not RAW_TEXT.strip():
        return "No content to embed."
    
    KNOWLEDGE_VECTOR_DATABASE = rag.build_vector([RAW_TEXT])

    return "Document embedded and indexed."

def respond(message, chat_history, context_text):
    global KNOWLEDGE_VECTOR_DATABASE

    if not message or message.strip() == "":
        return "", chat_history # do nothing

    if not context_text.strip() or KNOWLEDGE_VECTOR_DATABASE is None:
        chat_history.append({"role": "user", "content": message})
        chat_history.append({"role": "assistant", "content": "I’d love to help! Could you share a bit of context or upload some documents first so I can assist better?"})

        return "", chat_history
    
    answer, sources = rag.answer_with_rag(message, KNOWLEDGE_VECTOR_DATABASE, READER_MODEL)

    source_info = "\n\n**Sources:**\n"
    for i, doc in enumerate(sources[:5]):
        source_info += (
            "<details>\n"
            f"<summary>Document {i+1}</summary>\n"
            f"{doc.page_content[:300]}...\n"
            "</details>\n"
        )

    full_response = answer + source_info if (answer or sources) else "Sorry, I don't have enough information to answer that question."
    
    chat_history.append({"role": "user", "content": message})
    chat_history.append({"role": "assistant", "content": full_response})
    
    return "", chat_history

def on_context_change(text):
    global RAW_TEXT

    RAW_TEXT = text

    embed_text()
    
    return None

gr_css = """
    h1 {
        text-align: center; 
        font-size: 35px;
    }

    footer {
        visibility: hidden;
    }

    #submit-button {
        background-color: #fc7703; /* Orange */
    }
    
    #submit-button:hover {
        background-color: #fc5203; /* Darker orange on hover */
    }

    .icon-button-wrapper.top-panel {
        display: none !important;
    }
"""

with gr.Blocks(css=gr_css) as demo:
    gr.HTML("<h1>Question Answering (QA) Chatbot</h1>")

    with gr.Row():
        with gr.Column(scale=2):
            uploaded_files = gr.Files(label="Documents (.txt, .pdf, .docx)", file_types=[".txt", ".pdf", ".docx"], height=150)
        with gr.Column(scale=8):
            with gr.Accordion("Extracted Context (Click to expand and edit)", open=False):
                extracted_context = gr.Textbox(lines=10, interactive=True, show_label=False, container=False)
                extracted_context.change(fn=on_context_change, inputs=extracted_context, outputs=None)

    greet_msg = [
        {"role": "assistant", "content": "Hello! How can I help you today?"}
    ]

    chatbot = gr.Chatbot(value=greet_msg, type="messages", height=480, show_label=False)
    msg = gr.Textbox(placeholder="Ask anything about the uploaded documents", show_label=False)
    with gr.Row():
        clear = gr.ClearButton([msg, chatbot, extracted_context, uploaded_files])
        submit_button = gr.Button(value="Submit", elem_id="submit-button")

    uploaded_files.change(fn=extract_text, inputs=uploaded_files, outputs=extracted_context)
    submit_button.click(fn=respond, inputs=[msg, chatbot, extracted_context], outputs=[msg, chatbot])
    msg.submit(fn=respond, inputs=[msg, chatbot, extracted_context], outputs=[msg, chatbot])

demo.launch()