app.py

import os
import shutil
import streamlit as st
from io import BytesIO

# Importing LlamaIndex components
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex, StorageContext
from llama_index.vector_stores.qdrant import QdrantVectorStore
from llama_index.core.memory import ChatMemoryBuffer
import qdrant_client

# =============================================================================
# Configuration and Global Initialization
# =============================================================================

# Ensure that the OpenAI API key is available
openai_api_key = os.getenv("OPENAI_API_KEY")
if not openai_api_key:
    raise ValueError("Please set your OPENAI_API_KEY environment variable.")

# System prompt for the chat engine
SYSTEM_PROMPT = (
    "You are an AI assistant who answers the user questions, "
    "use the schema fields to generate appropriate and valid json queries"
)

# Configure the LLM and embedding models
Settings.llm = OpenAI(model="gpt-3.5-turbo", temperature=0.4)
Settings.embed_model = OpenAIEmbedding(model="text-embedding-ada-002")

# Load initial documents from a directory called "new_file"
if os.path.exists("new_file"):
    documents = SimpleDirectoryReader("new_file").load_data()
else:
    documents = []

# Set up the Qdrant vector store (using an in-memory collection for simplicity)
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)

# Build the initial index and chat engine
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
chat_memory = ChatMemoryBuffer.from_defaults(token_limit=3000)
chat_engine = index.as_chat_engine(
    chat_mode="context",
    memory=chat_memory,
    system_prompt=SYSTEM_PROMPT,
)

# =============================================================================
# Helper Functions
# =============================================================================

def process_uploaded_file(uploaded_file: BytesIO) -> str:
    """
    Process the uploaded file:
      1. Save the file to an "uploads" folder.
      2. Copy it to a temporary folder ("temp_upload") for reading.
      3. Update the global documents list and rebuild the index and chat engine.
    """
    if uploaded_file is None:
        return "No file uploaded."

    # Ensure the uploads directory exists
    uploads_dir = "uploads"
    os.makedirs(uploads_dir, exist_ok=True)

    # Save the uploaded file locally
    file_name = uploaded_file.name
    dest_path = os.path.join(uploads_dir, file_name)
    with open(dest_path, "wb") as f:
        f.write(uploaded_file.getbuffer())

    # Prepare a temporary directory for processing the file
    temp_dir = "temp_upload"
    os.makedirs(temp_dir, exist_ok=True)
    # Clear any existing file in temp_upload directory
    for f_name in os.listdir(temp_dir):
        os.remove(os.path.join(temp_dir, f_name))
    shutil.copy(dest_path, temp_dir)

    # Load new document(s) from the temporary folder using SimpleDirectoryReader
    new_docs = SimpleDirectoryReader(temp_dir).load_data()

    # Update global documents and rebuild the index and chat engine
    global documents, index, chat_engine
    documents.extend(new_docs)
    index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
    chat_engine = index.as_chat_engine(
        chat_mode="context",
        memory=chat_memory,
        system_prompt=SYSTEM_PROMPT,
    )

    return f"File '{file_name}' processed and added to the index."

def chat_with_ai(user_input: str) -> str:
    """
    Send user input to the chat engine and return the response.
    """
    response = chat_engine.chat(user_input)
    # Extract references from the response (if any)
    references = response.source_nodes
    ref = []
    for node in references:
        if "file_name" in node.metadata and node.metadata["file_name"] not in ref:
            ref.append(node.metadata["file_name"])
    complete_response = str(response)
    if ref:
        complete_response += "\n\nReferences: " + ", ".join(ref)
    return complete_response

# =============================================================================
# Streamlit App Layout
# =============================================================================

st.set_page_config(page_title="LlamaIndex Chat & File Upload", layout="wide")
st.title("Chat Interface for LlamaIndex with File Upload")

# Use Streamlit tabs for separate Chat and Upload functionalities
tab1, tab2 = st.tabs(["Chat", "Upload"])

# -----------------------------------------------------------------------------
# Chat Tab
# -----------------------------------------------------------------------------
with tab1:
    st.header("Chat with the AI")
    # Initialize chat history in session state if it does not exist
    if "chat_history" not in st.session_state:
        st.session_state["chat_history"] = []

    # Display conversation history
    for chat in st.session_state["chat_history"]:
        st.markdown(f"**User:** {chat[0]}")
        st.markdown(f"**AI:** {chat[1]}")
        st.markdown("---")

    # Input text for user query
    user_input = st.text_input("Enter your question:")

    # When the "Send" button is clicked, process the chat
    if st.button("Send") and user_input:
        with st.spinner("Processing..."):
            response = chat_with_ai(user_input)
        st.session_state["chat_history"].append((user_input, response))
        st.experimental_rerun()  # Refresh the page to show updated history

    # Button to clear the conversation history
    if st.button("Clear History"):
        st.session_state["chat_history"] = []
        st.experimental_rerun()

# -----------------------------------------------------------------------------
# Upload Tab
# -----------------------------------------------------------------------------
with tab2:
    st.header("Upload a File")
    uploaded_file = st.file_uploader("Choose a file to upload", type=["txt", "pdf", "doc", "docx", "csv", "xlsx"])
    if st.button("Upload and Process"):
        if uploaded_file is not None:
            with st.spinner("Uploading and processing file..."):
                status = process_uploaded_file(uploaded_file)
            st.success(status)
        else:
            st.error("No file uploaded.")