Update app.py

app.py

@@ -13,10 +13,12 @@ from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_community.llms import HuggingFaceHub
 from langchain_core.runnables import RunnableParallel, RunnablePassthrough
 from langchain_core.documents import Document
+
 huggingface_token = os.environ.get("HUGGINGFACE_TOKEN")
 
 # Memory database to store question-answer pairs
 memory_database = {}
+conversation_history = []
 
 def load_and_split_document_basic(file):
     """Loads and splits the document into pages."""
@@ -57,8 +59,13 @@ def clear_cache():
         return "No cache to clear."
 
 prompt = """
-Answer the question based only on the following context:
+Answer the question based on the following context and conversation history:
+Conversation History:
+{history}
+
+Context from documents:
 {context}
+
 Question: {question}
 
 Provide a concise and direct answer to the question:
@@ -81,21 +88,46 @@ def generate_chunked_response(model, prompt, max_tokens=1000, max_chunks=5):
     for i in range(max_chunks):
         chunk = model(prompt + full_response, max_new_tokens=max_tokens)
         chunk = chunk.strip()
-        # Check for final sentence endings
         if chunk.endswith((".", "!", "?")):
             full_response += chunk
             break
         full_response += chunk
     return full_response.strip()
 
-def response(database, model, question):
-    prompt_val = ChatPromptTemplate.from_template(prompt)
-    retriever = database.as_retriever()
-    context = retriever.get_relevant_documents(question)
-    context_str = "\n".join([doc.page_content for doc in context])
-    formatted_prompt = prompt_val.format(context=context_str, question=question)
-    ans = generate_chunked_response(model, formatted_prompt)
-    return ans.split("Question:")[-1].strip()  # Return only the answer part
+def manage_conversation_history(question, answer, history, max_history=5):
+    history.append({"question": question, "answer": answer})
+    if len(history) > max_history:
+        history.pop(0)
+    return history
+
+def ask_question(question, temperature, top_p, repetition_penalty):
+    global conversation_history
+    
+    if not question:
+        return "Please enter a question."
+    
+    if question in memory_database:
+        answer = memory_database[question]
+    else:
+        embed = get_embeddings()
+        database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
+        model = get_model(temperature, top_p, repetition_penalty)
+        
+        history_str = "\n".join([f"Q: {item['question']}\nA: {item['answer']}" for item in conversation_history])
+        prompt_val = ChatPromptTemplate.from_template(prompt)
+        retriever = database.as_retriever()
+        relevant_docs = retriever.get_relevant_documents(question)
+        context_str = "\n".join([doc.page_content for doc in relevant_docs])
+        formatted_prompt = prompt_val.format(history=history_str, context=context_str, question=question)
+        
+        answer = generate_chunked_response(model, formatted_prompt)
+        answer = answer.split("Question:")[-1].strip()
+        
+        memory_database[question] = answer
+    
+    conversation_history = manage_conversation_history(question, answer, conversation_history)
+    
+    return answer
 
 def update_vectors(files, use_recursive_splitter):
     if not files:
@@ -114,26 +146,6 @@ def update_vectors(files, use_recursive_splitter):
     
     return f"Vector store updated successfully. Processed {total_chunks} chunks from {len(files)} files."
 
-def ask_question(question, temperature, top_p, repetition_penalty):
-    if not question:
-        return "Please enter a question."
-    
-    # Check if the question exists in the memory database
-    if question in memory_database:
-        return memory_database[question]
-    
-    embed = get_embeddings()
-    database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
-    model = get_model(temperature, top_p, repetition_penalty)
-    
-    # Generate response from document database
-    answer = response(database, model, question)
-    
-    # Store the question and answer in the memory database
-    memory_database[question] = answer
-    
-    return answer
-
 def extract_db_to_excel():
     embed = get_embeddings()
     database = FAISS.load_local("faiss_database", embed, allow_dangerous_deserialization=True)
@@ -150,11 +162,16 @@ def extract_db_to_excel():
 
 def export_memory_db_to_excel():
     data = [{"question": question, "answer": answer} for question, answer in memory_database.items()]
-    df = pd.DataFrame(data)
+    df_memory = pd.DataFrame(data)
+    
+    data_history = [{"question": item["question"], "answer": item["answer"]} for item in conversation_history]
+    df_history = pd.DataFrame(data_history)
     
     with NamedTemporaryFile(delete=False, suffix='.xlsx') as tmp:
         excel_path = tmp.name
-        df.to_excel(excel_path, index=False)
+        with pd.ExcelWriter(excel_path, engine='openpyxl') as writer:
+            df_memory.to_excel(writer, sheet_name='Memory Database', index=False)
+            df_history.to_excel(writer, sheet_name='Conversation History', index=False)
     
     return excel_path
 
@@ -171,14 +188,21 @@ with gr.Blocks() as demo:
     update_button.click(update_vectors, inputs=[file_input, use_recursive_splitter], outputs=update_output)
     
     with gr.Row():
-        question_input = gr.Textbox(label="Ask a question about your documents")
-        temperature_slider = gr.Slider(label="Temperature", minimum=0.0, maximum=1.0, value=0.5, step=0.1)
-        top_p_slider = gr.Slider(label="Top P", minimum=0.0, maximum=1.0, value=0.9, step=0.1)
-        repetition_penalty_slider = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.0, step=0.1)
-        submit_button = gr.Button("Submit")
-    
-    answer_output = gr.Textbox(label="Answer")
-    submit_button.click(ask_question, inputs=[question_input, temperature_slider, top_p_slider, repetition_penalty_slider], outputs=answer_output)
+        with gr.Column(scale=2):
+            chatbot = gr.Chatbot(label="Conversation")
+            question_input = gr.Textbox(label="Ask a question about your documents")
+            submit_button = gr.Button("Submit")
+        with gr.Column(scale=1):
+            temperature_slider = gr.Slider(label="Temperature", minimum=0.0, maximum=1.0, value=0.5, step=0.1)
+            top_p_slider = gr.Slider(label="Top P", minimum=0.0, maximum=1.0, value=0.9, step=0.1)
+            repetition_penalty_slider = gr.Slider(label="Repetition Penalty", minimum=1.0, maximum=2.0, value=1.0, step=0.1)
+    
+    def chat(question, history):
+        answer = ask_question(question, temperature_slider.value, top_p_slider.value, repetition_penalty_slider.value)
+        history.append((question, answer))
+        return "", history
+    
+    submit_button.click(chat, inputs=[question_input, chatbot], outputs=[question_input, chatbot])
     
     extract_button = gr.Button("Extract Database to Excel")
     excel_output = gr.File(label="Download Excel File")