Update app.py

app.py

@@ -460,25 +460,29 @@ 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(search_kwargs={"k": 10})
-    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)}")
+    all_filtered_docs = []
+    k_per_doc = max(5, 20 // len(selected_docs))  # Adjust this value as needed
+
+    for doc_name in selected_docs:
+        logging.info(f"Retrieving documents for: {doc_name}")
+        doc_filter = lambda doc: doc.metadata["source"] == doc_name
+        doc_retriever = database.as_retriever(search_kwargs={"k": k_per_doc, "filter": doc_filter})
+        relevant_docs = doc_retriever.get_relevant_documents(query)
+        all_filtered_docs.extend(relevant_docs)
+        logging.info(f"Retrieved {len(relevant_docs)} documents for {doc_name}")
+
+    logging.info(f"Total number of filtered documents: {len(all_filtered_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:
+    if not all_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."
         return
 
-    for doc in filtered_docs:
+    for doc in all_filtered_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 all_filtered_docs])
     logging.info(f"Total context length: {len(context_str)}")
 
     if model == "@cf/meta/llama-3.1-8b-instruct":