Update app.py

app.py

@@ -45,6 +45,8 @@ llama_parser = LlamaParse(
     language="en",
 )
 
+uploaded_documents = []
+
 def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[Document]:
     """Loads and splits the document into pages."""
     if parser == "pypdf":
@@ -66,6 +68,7 @@ def get_embeddings():
     return HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
 
 def update_vectors(files, parser):
+    global uploaded_documents
     if not files:
         return "Please upload at least one PDF file."
     
@@ -77,6 +80,7 @@ def update_vectors(files, parser):
         data = load_document(file, parser)
         all_data.extend(data)
         total_chunks += len(data)
+        uploaded_documents.append({"name": file.name, "selected": True})
     
     if os.path.exists("faiss_database"):
         database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
@@ -214,10 +218,11 @@ def retry_last_response(history, use_web_search, model, temperature, num_calls):
     
     return chatbot_interface(last_user_msg, history, use_web_search, model, temperature, num_calls)
 
-def respond(message, history, model, temperature, num_calls, use_web_search):
+def respond(message, history, model, temperature, num_calls, use_web_search, selected_docs):
     logging.info(f"User Query: {message}")
     logging.info(f"Model Used: {model}")
     logging.info(f"Search Type: {'Web Search' if use_web_search else 'PDF Search'}")
+    logging.info(f"Selected Documents: {selected_docs}")
 
     try:
         if use_web_search:
@@ -231,10 +236,20 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
             if os.path.exists("faiss_database"):
                 database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
                 retriever = database.as_retriever()
-                relevant_docs = retriever.get_relevant_documents(message)
+                
+                # Filter relevant documents based on user selection
+                all_relevant_docs = retriever.get_relevant_documents(message)
+                relevant_docs = [doc for doc in all_relevant_docs if doc.metadata["source"] in selected_docs]
+                
+                if not relevant_docs:
+                    yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
+                    return
+
                 context_str = "\n".join([doc.page_content for doc in relevant_docs])
             else:
                 context_str = "No documents available."
+                yield "No documents available. Please upload PDF documents to answer questions."
+                return
             
             if model == "@cf/meta/llama-3.1-8b-instruct":
                 # Use Cloudflare API
@@ -244,7 +259,7 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
                     yield partial_response
             else:
                 # Use Hugging Face API
-                for partial_response in get_response_from_pdf(message, model, num_calls=num_calls, temperature=temperature):
+                for partial_response in get_response_from_pdf(message, model, selected_docs, num_calls=num_calls, temperature=temperature):
                     first_line = partial_response.split('\n')[0] if partial_response else ''
                     logging.info(f"Generated Response (first line): {first_line}")
                     yield partial_response
@@ -253,7 +268,7 @@ def respond(message, history, model, temperature, num_calls, use_web_search):
         if "microsoft/Phi-3-mini-4k-instruct" in model:
             logging.info("Falling back to Mistral model due to Phi-3 error")
             fallback_model = "mistralai/Mistral-7B-Instruct-v0.3"
-            yield from respond(message, history, fallback_model, temperature, num_calls, use_web_search)
+            yield from respond(message, history, fallback_model, temperature, num_calls, use_web_search, selected_docs)
         else:
             yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
 
@@ -344,7 +359,7 @@ After writing the document, please provide a list of sources used in your respon
                     main_content += chunk
                     yield main_content, ""  # Yield partial main content without sources
 
-def get_response_from_pdf(query, model, num_calls=3, temperature=0.2):
+def get_response_from_pdf(query, model, selected_docs, num_calls=3, temperature=0.2):
     embed = get_embeddings()
     if os.path.exists("faiss_database"):
         database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
@@ -354,7 +369,11 @@ def get_response_from_pdf(query, model, num_calls=3, temperature=0.2):
 
     retriever = database.as_retriever()
     relevant_docs = retriever.get_relevant_documents(query)
-    context_str = "\n".join([doc.page_content for doc in relevant_docs])
+    
+    # Filter relevant_docs based on selected documents
+    filtered_docs = [doc for doc in relevant_docs if doc.metadata["source"] in selected_docs]
+    
+    context_str = "\n".join([doc.page_content for doc in filtered_docs])
 
     if model == "@cf/meta/llama-3.1-8b-instruct":
         # Use Cloudflare API with the retrieved context
@@ -392,6 +411,15 @@ css = """
 """
 
 # Define the checkbox outside the demo block
+def display_documents():
+    return gr.CheckboxGroup(
+        choices=[doc["name"] for doc in uploaded_documents],
+        value=[doc["name"] for doc in uploaded_documents if doc["selected"]],
+        label="Select documents to query"
+    )
+
+document_selector = gr.CheckboxGroup(label="Select documents to query")
+
 use_web_search = gr.Checkbox(label="Use Web Search", value=False)
 
 demo = gr.ChatInterface(
@@ -400,7 +428,8 @@ demo = gr.ChatInterface(
         gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[0]),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
         gr.Slider(minimum=1, maximum=5, value=1, step=1, label="Number of API Calls"),
-        use_web_search  # Add this line to include the checkbox
+        use_web_search,
+        document_selector  # Add this line to include the document selector
     ],
     title="AI-powered Web Search and PDF Chat Assistant",
     description="Chat with your PDFs or use web search to answer questions.",