Update app.py

app.py

@@ -316,35 +316,47 @@ def respond(message, history, use_web_search, model, temperature, num_calls, sel
 
     try:
         if use_web_search:
-            search_results = duckduckgo_search(message)
-            context = "\n".join(f"{result['title']}\n{result['body']}\nSource: {result['href']}\n" 
-                                for result in search_results if 'body' in result)
-            prompt = f"""Using the following context:
-            {context}
-            Write a detailed and complete research document that fulfills the following user request: '{message}'
-            After writing the document, please provide a list of sources used in your response."""
+            for main_content, sources in get_response_with_search(message, model, num_calls=num_calls, temperature=temperature):
+                response = f"{main_content}\n\n{sources}"
+                yield response
         else:
             embed = get_embeddings()
             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)
-                filtered_docs = [doc for doc in relevant_docs if doc.metadata["source"] in selected_docs]
-                if not filtered_docs:
-                    return "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
-                context_str = "\n".join([doc.page_content for doc in filtered_docs])
+                
+                # 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:
+                yield "No documents available. Please upload PDF documents to answer questions."
+                return
+            
+            if model == "@cf/meta/llama-3.1-8b-instruct":
+                # Use Cloudflare 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: '{message}'"""
+                for partial_response in generate_chunked_response(prompt, model, num_calls=num_calls, temperature=temperature):
+                    yield partial_response
             else:
-                return "No documents available. Please upload PDF documents to answer questions."
-
-        for response in generate_chunked_response(prompt, model, num_calls=num_calls, temperature=temperature):
-            yield response
-
+                # Use Hugging Face API
+                for partial_response in get_response_from_pdf(message, model, selected_docs, num_calls=num_calls, temperature=temperature):
+                    yield partial_response
     except Exception as e:
         logging.error(f"Error with {model}: {str(e)}")
-        yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
+        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, use_web_search, fallback_model, temperature, num_calls, selected_docs)
+        else:
+            yield f"An error occurred with the {model} model: {str(e)}. Please try again or select a different model."
 
 logging.basicConfig(level=logging.DEBUG)