Update app.py

app.py

@@ -477,63 +477,73 @@ def get_response_from_pdf(query, model, selected_docs, num_calls=3, temperature=
     embed = get_embeddings()
     if os.path.exists("faiss_database"):
         logging.info("Loading FAISS database")
-        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
-        logging.info(f"FAISS database loaded with {len(database.index)} vectors")
+        try:
+            database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+            logging.info(f"FAISS database loaded with {database.index.ntotal} vectors")
+        except Exception as e:
+            logging.error(f"Error loading FAISS database: {str(e)}")
+            yield "Error loading the document database. Please try uploading the documents again."
+            return
     else:
         logging.warning("No FAISS database found")
         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."
-        return
-
-    for i, doc in enumerate(filtered_docs):
-        logging.info(f"Document {i+1} source: {doc.metadata['source']}")
-        logging.info(f"Document {i+1} content preview: {doc.page_content[:100]}...")
-
-    context_str = "\n".join([doc.page_content for doc in filtered_docs])
-    logging.info(f"Total context length: {len(context_str)}")
-
-    if model == "@cf/meta/llama-3.1-8b-instruct":
-        logging.info("Using Cloudflare API")
-        # Use Cloudflare API with the retrieved context
-        for response in get_response_from_cloudflare(prompt="", context=context_str, query=query, num_calls=num_calls, temperature=temperature, search_type="pdf"):
-            yield response
-    else:
-        logging.info("Using Hugging Face API")
-        # Use Hugging Face API
-        prompt = f"""Using the following context from the PDF documents:
-{context_str}
-Write a detailed and complete response that answers the following user question: '{query}'"""
+    try:
+        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)}")
         
-        client = InferenceClient(model, token=huggingface_token)
+        # 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)}")
         
-        response = ""
-        for i in range(num_calls):
-            logging.info(f"API call {i+1}/{num_calls}")
-            for message in client.chat_completion(
-                messages=[{"role": "user", "content": prompt}],
-                max_tokens=10000,
-                temperature=temperature,
-                stream=True,
-            ):
-                if message.choices and message.choices[0].delta and message.choices[0].delta.content:
-                    chunk = message.choices[0].delta.content
-                    response += chunk
-                    yield response  # Yield partial response
-        
-        logging.info("Finished generating response")
+        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."
+            return
+
+        for i, doc in enumerate(filtered_docs):
+            logging.info(f"Document {i+1} source: {doc.metadata['source']}")
+            logging.info(f"Document {i+1} content preview: {doc.page_content[:100]}...")
+
+        context_str = "\n".join([doc.page_content for doc in filtered_docs])
+        logging.info(f"Total context length: {len(context_str)}")
+
+        if model == "@cf/meta/llama-3.1-8b-instruct":
+            logging.info("Using Cloudflare API")
+            # Use Cloudflare API with the retrieved context
+            for response in get_response_from_cloudflare(prompt="", context=context_str, query=query, num_calls=num_calls, temperature=temperature, search_type="pdf"):
+                yield response
+        else:
+            logging.info("Using Hugging Face API")
+            # Use Hugging Face API
+            prompt = f"""Using the following context from the PDF documents:
+{context_str}
+Write a detailed and complete response that answers the following user question: '{query}'"""
+            
+            client = InferenceClient(model, token=huggingface_token)
+            
+            response = ""
+            for i in range(num_calls):
+                logging.info(f"API call {i+1}/{num_calls}")
+                for message in client.chat_completion(
+                    messages=[{"role": "user", "content": prompt}],
+                    max_tokens=10000,
+                    temperature=temperature,
+                    stream=True,
+                ):
+                    if message.choices and message.choices[0].delta and message.choices[0].delta.content:
+                        chunk = message.choices[0].delta.content
+                        response += chunk
+                        yield response  # Yield partial response
+            
+            logging.info("Finished generating response")
+
+    except Exception as e:
+        logging.error(f"Error in get_response_from_pdf: {str(e)}")
+        yield f"An error occurred while processing your query: {str(e)}. Please try again or contact support."
 
 def vote(data: gr.LikeData):
     if data.liked: