Update app.py

app.py

@@ -66,29 +66,6 @@ def load_document(file: NamedTemporaryFile, parser: str = "llamaparse") -> List[
 def get_embeddings():
     return HuggingFaceEmbeddings(model_name="sentence-transformers/stsb-roberta-large")
 
-def scan_faiss_database():
-    global uploaded_documents
-    uploaded_documents = []
-    
-    if os.path.exists("faiss_database"):
-        embed = get_embeddings()
-        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
-        
-        for doc in database.docstore._dict.values():
-            file_name = os.path.basename(doc.metadata['source'])
-            if not any(d['name'] == file_name for d in uploaded_documents):
-                uploaded_documents.append({"name": file_name, "selected": True})
-    
-    return uploaded_documents
-
-# Call this function when the application starts
-uploaded_documents = scan_faiss_database()
-
-def initialize_session():
-    global uploaded_documents
-    uploaded_documents = scan_faiss_database()
-    return uploaded_documents
-
 def update_vectors(files, parser):
     global uploaded_documents
     logging.info(f"Entering update_vectors with {len(files)} files and parser: {parser}")
@@ -112,12 +89,12 @@ def update_vectors(files, parser):
             logging.info(f"Loaded {len(data)} chunks from {file.name}")
             all_data.extend(data)
             total_chunks += len(data)
-            file_name = os.path.basename(file.name)
-            if not any(doc["name"] == file_name for doc in uploaded_documents):
-                uploaded_documents.append({"name": file_name, "selected": True})
-                logging.info(f"Added new document to uploaded_documents: {file_name}")
+            # Append new documents instead of replacing
+            if not any(doc["name"] == file.name for doc in uploaded_documents):
+                uploaded_documents.append({"name": file.name, "selected": True})
+                logging.info(f"Added new document to uploaded_documents: {file.name}")
             else:
-                logging.info(f"Document already exists in uploaded_documents: {file_name}")
+                logging.info(f"Document already exists in uploaded_documents: {file.name}")
         except Exception as e:
             logging.error(f"Error processing file {file.name}: {str(e)}")
     
@@ -133,6 +110,7 @@ def update_vectors(files, parser):
     
     database.save_local("faiss_database")
     logging.info("FAISS database saved")
+
     
     return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files using {parser}.", gr.CheckboxGroup(
         choices=[doc["name"] for doc in uploaded_documents],
@@ -240,14 +218,14 @@ class CitingSources(BaseModel):
         ...,
         description="List of sources to cite. Should be an URL of the source."
     )
-def chatbot_interface(message, history, use_web_search, model, temperature, num_calls, selected_docs):
+def chatbot_interface(message, history, use_web_search, model, temperature, num_calls):
     if not message.strip():
         return "", history
 
     history = history + [(message, "")]
 
     try:
-        for response in respond(message, history, model, temperature, num_calls, use_web_search, selected_docs):
+        for response in respond(message, history, model, temperature, num_calls, use_web_search):
             history[-1] = (message, response)
             yield history
     except gr.CancelledError:
@@ -270,12 +248,15 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
     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:
             for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
                 response = f"{main_content}\n\n{sources}"
+                first_line = response.split('\n')[0] if response else ''
+#                logging.info(f"Generated Response (first line): {first_line}")
                 yield response
         else:
             embed = get_embeddings()
@@ -283,8 +264,9 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
                 database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
                 retriever = database.as_retriever()
                 
+                # 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 os.path.basename(doc.metadata["source"]) in selected_docs]
+                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."
@@ -299,10 +281,14 @@ def respond(message, history, model, temperature, num_calls, use_web_search, sel
             if model == "@cf/meta/llama-3.1-8b-instruct":
                 # Use Cloudflare API
                 for partial_response in get_response_from_cloudflare(prompt="", context=context_str, query=message, num_calls=num_calls, temperature=temperature, search_type="pdf"):
+                    first_line = partial_response.split('\n')[0] if partial_response else ''
+                    logging.info(f"Generated Response (first line): {first_line}")
                     yield partial_response
             else:
                 # Use Hugging Face API
                 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
     except Exception as e:
         logging.error(f"Error with {model}: {str(e)}")
@@ -555,9 +541,7 @@ demo = gr.ChatInterface(
 )
 
 # Add file upload functionality
-with gr.Blocks() as demo:
-    session_documents = gr.State(initialize_session)
-    
+with demo:
     gr.Markdown("## Upload PDF Documents")
 
     with gr.Row():
@@ -566,29 +550,11 @@ with gr.Blocks() as demo:
         update_button = gr.Button("Upload Document")
     
     update_output = gr.Textbox(label="Update Status")
-    document_selector = 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"
-    )
     
     # Update both the output text and the document selector
-    update_button.click(
-        update_vectors, 
-        inputs=[file_input, parser_dropdown], 
-        outputs=[update_output, document_selector]
-    )
-
-    # Add a refresh button to update the document selector
-    refresh_button = gr.Button("Refresh Document List")
-    refresh_button.click(
-        lambda: 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"
-        ),
-        outputs=[document_selector]
-    )  
+    update_button.click(update_vectors, 
+                        inputs=[file_input, parser_dropdown], 
+                        outputs=[update_output, document_selector])
 
     gr.Markdown(
     """