app.py

import os
import shutil
import gradio as gr
import qdrant_client
from getpass import getpass

# Set your OpenAI API key from environment variables.
openai_api_key = os.getenv('OPENAI_API_KEY')

# -------------------------------------------------------
# Configure LlamaIndex with OpenAI LLM and Embeddings
# -------------------------------------------------------
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import Settings

Settings.llm = OpenAI(model="gpt-3.5-turbo", temperature=0.4)
Settings.embed_model = OpenAIEmbedding(model="text-embedding-ada-002")

# -------------------------------------------------------
# Import document readers, index, vector store, memory, etc.
# -------------------------------------------------------
from llama_index.core import SimpleDirectoryReader, VectorStoreIndex, StorageContext
from llama_index.vector_stores.qdrant import QdrantVectorStore
from llama_index.core.memory import ChatMemoryBuffer

# Global variables to hold the index and chat engine.
chat_engine = None
index = None
query_engine = None
memory = None
client = None
vector_store = None
storage_context = None

# -------------------------------------------------------
# Function to process uploaded files and build the index.
# -------------------------------------------------------
def process_upload(files):
    """
    Accepts a list of uploaded file paths, saves them to a local folder,
    loads them as documents, and builds the vector index and chat engine.
    """
    upload_dir = "uploaded_files"
    if not os.path.exists(upload_dir):
        os.makedirs(upload_dir)
    else:
        # Clear any existing files in the folder.
        for f in os.listdir(upload_dir):
            os.remove(os.path.join(upload_dir, f))
    
    # 'files' is a list of file paths (Gradio's File component with type="file")
    for file_path in files:
        file_name = os.path.basename(file_path)
        dest = os.path.join(upload_dir, file_name)
        shutil.copy(file_path, dest)
    
    # Load documents from the saved folder.
    documents = SimpleDirectoryReader(upload_dir).load_data()
    
    # Build the index and chat engine using Qdrant as the vector store.
    global client, vector_store, storage_context, index, query_engine, memory, chat_engine
    client = qdrant_client.QdrantClient(location=":memory:")
    
    vector_store = QdrantVectorStore(
        collection_name="paper",
        client=client,
        enable_hybrid=True,
        batch_size=20,
    )
    
    storage_context = StorageContext.from_defaults(vector_store=vector_store)
    
    index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
    
    query_engine = index.as_query_engine(vector_store_query_mode="hybrid")
    
    memory = ChatMemoryBuffer.from_defaults(token_limit=3000)
    
    chat_engine = index.as_chat_engine(
        chat_mode="context",
        memory=memory,
        system_prompt=(
            "You are an AI assistant who answers the user questions, "
            "use the schema fields to generate appropriate and valid json queries"
        ),
    )
    
    return "Documents uploaded and index built successfully!"

# -------------------------------------------------------
# Chat function that uses the built chat engine.
# -------------------------------------------------------
def chat_with_ai(user_input, chat_history):
    global chat_engine
    # Check if the chat engine is initialized.
    if chat_engine is None:
        return chat_history, "Please upload documents first."
    
    response = chat_engine.chat(user_input)
    references = response.source_nodes
    ref, pages = [], []
    
    # Extract file names from the source nodes (if available)
    for node in references:
        file_name = node.metadata.get('file_name')
        if file_name and file_name not in ref:
            ref.append(file_name)
    
    complete_response = str(response) + "\n\n"
    if ref or pages:
        chat_history.append((user_input, complete_response))
    else:
        chat_history.append((user_input, str(response)))
    return chat_history, ""

# -------------------------------------------------------
# Function to clear the chat history.
# -------------------------------------------------------
def clear_history():
    return [], ""

# -------------------------------------------------------
# Build the Gradio interface.
# -------------------------------------------------------
def gradio_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# Chat Interface for LlamaIndex with File Upload")
        
        # Use Tabs to separate the file upload and chat interfaces.
        with gr.Tab("Upload Documents"):
            gr.Markdown("Upload PDF, Excel, CSV, DOC/DOCX, or TXT files below:")
            # The file upload widget: we specify allowed file types.
            file_upload = gr.File(
                label="Upload Files",
                file_count="multiple",
                file_types=[".pdf", ".csv", ".txt", ".xlsx", ".xls", ".doc", ".docx"],
                type="filepath"  
            )
            upload_status = gr.Textbox(label="Upload Status", interactive=False)
            upload_button = gr.Button("Process Upload")
            
            upload_button.click(process_upload, inputs=file_upload, outputs=upload_status)
        
        with gr.Tab("Chat"):
            chatbot = gr.Chatbot(label="LlamaIndex Chatbot")
            user_input = gr.Textbox(
                placeholder="Ask a question...", label="Enter your question"
            )
            submit_button = gr.Button("Send")
            btn_clear = gr.Button("Clear History")
            
            # A State to hold the chat history.
            chat_history = gr.State([])
            
            submit_button.click(chat_with_ai, inputs=[user_input, chat_history], outputs=[chatbot, user_input])
            user_input.submit(chat_with_ai, inputs=[user_input, chat_history], outputs=[chatbot, user_input])
            btn_clear.click(clear_history, outputs=[chatbot, user_input])
    
    return demo

# Launch the Gradio app.
gradio_interface().launch(debug=True)