Update app.py

app.py

@@ -1,8 +1,9 @@
 import os
 import shutil
-from getpass import getpass
+import streamlit as st
+from io import BytesIO
 
-import gradio as gr
+# 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
@@ -10,12 +11,16 @@ from llama_index.vector_stores.qdrant import QdrantVectorStore
 from llama_index.core.memory import ChatMemoryBuffer
 import qdrant_client
 
-# Set your OpenAI API key from environment variable
-openai_api_key = os.getenv('OPENAI_API_KEY')
+# =============================================================================
+# 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.")
 
-# Define a system prompt as a global constant
+# 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"
@@ -26,7 +31,10 @@ 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"
-documents = SimpleDirectoryReader("new_file").load_data()
+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:")
@@ -38,7 +46,7 @@ vector_store = QdrantVectorStore(
 )
 storage_context = StorageContext.from_defaults(vector_store=vector_store)
 
-# Build the initial index and query/chat engines
+# 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(
@@ -47,35 +55,42 @@ chat_engine = index.as_chat_engine(
     system_prompt=SYSTEM_PROMPT,
 )
 
-def process_uploaded_file(uploaded_file):
+# =============================================================================
+# 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") to load using SimpleDirectoryReader.
-      3. Extend the global documents list and rebuild the index and chat engine.
+      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."
 
-    # 'uploaded_file' is a temporary file path provided by Gradio.
-    file_name = os.path.basename(uploaded_file)
+    # 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)
-    shutil.copy(uploaded_file, dest_path)
+    with open(dest_path, "wb") as f:
+        f.write(uploaded_file.getbuffer())
 
-    # Prepare a temporary directory to read the file
+    # Prepare a temporary directory for processing the file
     temp_dir = "temp_upload"
     os.makedirs(temp_dir, exist_ok=True)
-    # Clear previous files in temp_dir (optional, to avoid mixing files)
-    for f in os.listdir(temp_dir):
-        os.remove(os.path.join(temp_dir, f))
+    # 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 the new document(s) from the temporary folder
+    # Load new document(s) from the temporary folder using SimpleDirectoryReader
     new_docs = SimpleDirectoryReader(temp_dir).load_data()
 
-    # Update the global documents list and rebuild the index and chat engine
+    # 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)
@@ -85,60 +100,75 @@ def process_uploaded_file(uploaded_file):
         system_prompt=SYSTEM_PROMPT,
     )
 
-    return f"File '{file_name}' processed and added to index."
+    return f"File '{file_name}' processed and added to the index."
 
-def chat_with_ai(user_input, chat_history):
+def chat_with_ai(user_input: str) -> str:
     """
-    Send the user input to the chat engine and update the conversation history.
+    Send user input to the chat engine and return the response.
     """
     response = chat_engine.chat(user_input)
-    # Collect reference filenames from the response (if available)
+    # 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"])
-    # Create a complete response string with references if present
     complete_response = str(response)
     if ref:
         complete_response += "\n\nReferences: " + ", ".join(ref)
-    chat_history.append((user_input, complete_response))
-    return chat_history, ""
+    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.")
 
-def clear_history():
-    return [], ""
-
-def gradio_chatbot():
-    """
-    Create a Gradio interface with two tabs:
-      - "Chat" for interacting with the chat engine.
-      - "Upload" for uploading new files to update the index.
-    """
-  with gr.Blocks() as demo:
-    gr.Markdown("# Chat Interface for LlamaIndex with File Upload")
-    
-    with gr.Tab("Chat"):
-        # Chat interface components
-        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("Delete Context")
-        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(fn=clear_history, outputs=[chatbot, user_input])
-    
-    with gr.Tab("Upload"):
-        gr.Markdown("### Upload a file to add its content to the index")
-        file_upload = gr.File(label="Choose a file")
-        upload_button = gr.Button("Upload and Process")
-        upload_status = gr.Textbox(label="Upload Status")
-        upload_button.click(process_uploaded_file, inputs=[file_upload], outputs=[upload_status])
-
-    
-    return demo
-
-if __name__ == "__main__":
-    gradio_chatbot().launch(debug=True)