Update app.py

app.py

@@ -104,6 +104,11 @@ def update_vectors(files, parser):
                 logging.warning(f"No chunks loaded from {file.name}")
                 continue
             logging.info(f"Loaded {len(data)} chunks from {file.name}")
+            
+            # Add source to metadata
+            for chunk in data:
+                chunk.metadata["source"] = file.name
+            
             all_data.extend(data)
             total_chunks += len(data)
             if not any(doc["name"] == file.name for doc in uploaded_documents):
@@ -460,25 +465,34 @@ def get_response_from_pdf(query, model, selected_docs, num_calls=3, temperature=
         yield "No documents available. Please upload PDF documents to answer questions."
         return
 
-    retriever = database.as_retriever()
-    logging.info(f"Retrieving relevant documents for query: {query}")
-    relevant_docs = retriever.get_relevant_documents(query)
-    logging.info(f"Number of relevant documents retrieved: {len(relevant_docs)}")
-    
-    # Filter relevant_docs based on selected documents
-    filtered_docs = [doc for doc in relevant_docs if doc.metadata["source"] in selected_docs]
-    logging.info(f"Number of filtered documents: {len(filtered_docs)}")
-    
-    if not filtered_docs:
-        logging.warning(f"No relevant information found in the selected documents: {selected_docs}")
-        yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
+    # New function to get documents by source
+    def get_documents_by_source(database, sources):
+        return [doc for doc in database.docstore._dict.values() if doc.metadata["source"] in sources]
+
+    # Log all documents in the database
+    logging.info(f"All documents in database: {[doc.metadata['source'] for doc in database.docstore._dict.values()]}")
+
+    # Get only the selected documents
+    selected_docs_content = get_documents_by_source(database, selected_docs)
+    logging.info(f"Number of selected documents: {len(selected_docs_content)}")
+
+    if not selected_docs_content:
+        logging.warning(f"No documents found for the selected sources: {selected_docs}")
+        yield "No documents found for the selected sources. Please check your document selection."
         return
 
-    for doc in filtered_docs:
+    # Perform similarity search on selected documents
+    relevant_docs = sorted(selected_docs_content, 
+                           key=lambda doc: database.similarity_search_with_score(query, k=1, filter={"source": doc.metadata["source"]})[0][1], 
+                           reverse=True)
+
+    logging.info(f"Number of relevant documents: {len(relevant_docs)}")
+
+    for doc in relevant_docs:
         logging.info(f"Document source: {doc.metadata['source']}")
         logging.info(f"Document content preview: {doc.page_content[:100]}...")  # Log first 100 characters of each document
 
-    context_str = "\n".join([doc.page_content for doc in filtered_docs])
+    context_str = "\n".join([doc.page_content for doc in relevant_docs])
     logging.info(f"Total context length: {len(context_str)}")
 
     if model == "@cf/meta/llama-3.1-8b-instruct":