Update app.py

app.py

@@ -24,7 +24,7 @@ from llama_index.core import SimpleDirectoryReader, VectorStoreIndex, StorageCon
 from llama_index.vector_stores.qdrant import QdrantVectorStore
 from llama_index.core.memory import ChatMemoryBuffer
 
-# Global variables to hold persistent objects.
+# Global variables to hold the index and chat engine.
 chat_engine = None
 index = None
 query_engine = None
@@ -33,122 +33,109 @@ client = None
 vector_store = None
 storage_context = None
 
-# Define a global collection name (you can change this as needed)
-collection_name = "paper"
-
+# -------------------------------------------------------
+# Function to process uploaded files and build the index.
+# -------------------------------------------------------
 def process_upload(files):
     """
-    Process newly uploaded files by copying them into a persistent folder,
-    loading their content, and then either building a new index or inserting
-    new documents into the existing index.
+    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"
-    # Create the upload folder if it does not exist.
     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))
     
-    # Copy new files into the upload directory.
-    new_file_paths = []
+    # '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)
-        # Copy the file if it doesn't already exist.
-        if not os.path.exists(dest):
-            shutil.copy(file_path, dest)
-        new_file_paths.append(dest)
+        shutil.copy(file_path, dest)
     
-    # Load only the newly uploaded documents.
-    # (SimpleDirectoryReader can accept a list of file paths via the 'input_files' parameter.)
-    documents = SimpleDirectoryReader(input_files=new_file_paths).load_data()
+    # 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
-
-    # Initialize Qdrant client if not already done.
-    if client is None:
-        client = qdrant_client.QdrantClient(
-            path="./qdrant_db",  
-            prefer_grpc=True
-        )
+    client = qdrant_client.QdrantClient(location=":memory:")
     
-    # Ensure the collection exists.
-    from qdrant_client.http import models
-    existing_collections = {col.name for col in client.get_collections().collections}
-    if collection_name not in existing_collections:
-        client.create_collection(
-            collection_name=collection_name,
-            vectors_config=models.VectorParams(
-                size=1536,  # OpenAI's text-embedding-ada-002 produces 1536-d vectors.
-                distance=models.Distance.COSINE
-            )
-        )
-    
-    # Initialize the vector store if not already done.
-    if vector_store is None:
-        vector_store = QdrantVectorStore(
-            collection_name=collection_name,
-            client=client,
-            enable_hybrid=True,
-            batch_size=20,
-        )
+    vector_store = QdrantVectorStore(
+        collection_name="paper",
+        client=client,
+        enable_hybrid=True,
+        batch_size=20,
+    )
     
-    # Initialize storage context if not already done.
-    if storage_context is None:
-        storage_context = StorageContext.from_defaults(vector_store=vector_store)
+    storage_context = StorageContext.from_defaults(vector_store=vector_store)
     
-    # If no index exists yet, create one from the documents.
-    if index is None:
-        index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
-    else:
-        # Append the new documents to the existing index.
-        index.insert_documents(documents)
+    index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
     
-    # (Optional) Reinitialize the query and chat engines so they reflect the updated index.
     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,"
+        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 updated successfully!"
+    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 = []
+    ref, pages = [], []
     
-    # Extract referenced file names from the response.
+    # 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:
+    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("# AI Assistant")
+        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"  # Returns file paths.
+                type="filepath"  # returns file paths
             )
             upload_status = gr.Textbox(label="Upload Status", interactive=False)
             upload_button = gr.Button("Process Upload")
@@ -156,7 +143,7 @@ def gradio_interface():
             upload_button.click(process_upload, inputs=file_upload, outputs=upload_status)
         
         with gr.Tab("Chat"):
-            chatbot = gr.Chatbot(label="AI Assistant Chat Interface")
+            chatbot = gr.Chatbot(label="LlamaIndex Chatbot")
             user_input = gr.Textbox(
                 placeholder="Ask a question...", label="Enter your question"
             )
@@ -172,4 +159,5 @@ def gradio_interface():
     
     return demo
 
+# Launch the Gradio app.
 gradio_interface().launch(debug=True)