Update app.py

app.py

@@ -18,9 +18,7 @@ from huggingface_hub import InferenceClient
 import inspect
 import logging
 import shutil
-from langchain.chains import ConversationalRetrievalChain
-from langchain.memory import ConversationBufferMemory
-from langchain.llms import HuggingFaceHub
+
 
 # Set up basic configuration for logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
@@ -283,16 +281,18 @@ 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:
+        yield history
     except Exception as e:
         logging.error(f"Unexpected error in chatbot_interface: {str(e)}")
         history[-1] = (message, f"An unexpected error occurred: {str(e)}")
@@ -307,65 +307,52 @@ def retry_last_response(history, use_web_search, model, temperature, num_calls):
     
     return chatbot_interface(last_user_msg, history, use_web_search, model, temperature, num_calls)
 
-def create_conversation_chain(model_name, vector_store, temperature=0.5):
-    # Create a HuggingFaceHub instance
-    llm = HuggingFaceHub(
-        repo_id=model_name,
-        model_kwargs={
-            "temperature": temperature
-        },
-        huggingfacehub_api_token=huggingface_token  # Ensure this variable is defined in your environment
-    )
-
-    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
-    
-    chain = ConversationalRetrievalChain.from_llm(
-        llm=llm,
-        retriever=vector_store.as_retriever(),
-        memory=memory
-    )
-    
-    return chain
-
 def respond(message, history, model, temperature, num_calls, use_web_search, selected_docs):
     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()
             if os.path.exists("faiss_database"):
                 database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+                retriever = database.as_retriever(search_kwargs={"k": 20})
                 
                 # Filter relevant documents based on user selection
-                filtered_docs = [doc for doc in database.docstore._dict.values() 
-                                 if isinstance(doc, Document) and doc.metadata.get("source") in selected_docs]
+                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 filtered_docs:
+                if not relevant_docs:
                     yield "No relevant information found in the selected documents. Please try selecting different documents or rephrasing your query."
                     return
-                
-                # Create a new FAISS index with only the selected documents
-                filtered_db = FAISS.from_documents(filtered_docs, embed)
-                
-                # Create the conversation chain
-                chain = create_conversation_chain(model, filtered_db, temperature)
-                
-                # Generate response
-                response = chain({"question": message})
-                
-                yield response['answer']
+
+                context_str = "\n".join([doc.page_content for doc in relevant_docs])
             else:
-                logging.warning("No FAISS database found")
+                context_str = "No documents available."
                 yield "No documents available. Please upload PDF documents to answer questions."
                 return
-
+            
+            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)}")
         if "microsoft/Phi-3-mini-4k-instruct" in model:
@@ -583,7 +570,7 @@ use_web_search = gr.Checkbox(label="Use Web Search", value=True)
 custom_placeholder = "Ask a question (Note: You can toggle between Web Search and PDF Chat in Additional Inputs below)"
 
 demo = gr.ChatInterface(
-    chatbot_interface,
+    respond,
     additional_inputs=[
         gr.Dropdown(choices=MODELS, label="Select Model", value=MODELS[3]),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.2, step=0.1, label="Temperature"),
@@ -671,4 +658,4 @@ with demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=True)