Update app.py

app.py

@@ -27,7 +27,6 @@ client = Groq(api_key=groq_api_key)
 
 hf_token = hf_api_key
 
-memory = ConversationBufferMemory()
 embeddings = HuggingFaceEmbeddings(model_name="heydariAI/persian-embeddings")
 vector_store = InMemoryVectorStore(embeddings)
 
@@ -35,7 +34,6 @@ DATASET_NAME = "chat_history"
 try:
     dataset = load_dataset(DATASET_NAME, use_auth_token=hf_token)
 except Exception:
-
     dataset = Dataset.from_dict({"Timestamp": [], "User": [], "ParvizGPT": []})
 
 def save_chat_to_dataset(user_message, bot_message):
@@ -66,7 +64,7 @@ def process_pdf_with_langchain(pdf_path):
         logger.error(f"Error processing PDF: {e}")
         raise
 
-def generate_response(query, retriever=None, use_pdf_context=False):
+def generate_response(query, memory, retriever=None, use_pdf_context=False):
     try:
         knowledge = ""
 
@@ -105,33 +103,33 @@ def generate_response(query, retriever=None, use_pdf_context=False):
                 logger.error(f"Attempt {attempt + 1} failed: {e}")
                 time.sleep(2)
 
-        return response
+        return response, memory
     except Exception as e:
         logger.error(f"Error generating response: {e}")
-        return f"Error: {e}"
+        return f"Error: {e}", memory
 
-def gradio_interface(user_message, chat_box, pdf_file=None, use_pdf_context=False):
+def gradio_interface(user_message, chat_box, memory, pdf_file=None, use_pdf_context=False):
     global retriever
     if pdf_file is not None and use_pdf_context:  
         try:
             retriever = process_pdf_with_langchain(pdf_file.name)
         except Exception as e:
-            return chat_box + [("Error", f"Error processing PDF: {e}")]
+            return chat_box + [("Error", f"Error processing PDF: {e}")], memory
 
     chat_box.append(("ParvizGPT", "در حال پردازش..."))
     
-    response = generate_response(user_message, retriever=retriever, use_pdf_context=use_pdf_context)
+    response, memory = generate_response(user_message, memory, retriever=retriever, use_pdf_context=use_pdf_context)
     
     chat_box[-1] = ("You", user_message) 
     chat_box.append(("ParvizGPT", response))
     
     save_chat_to_dataset(user_message, response)
     
-    return chat_box
+    return chat_box, memory
 
-def clear_memory():
+def clear_memory(memory):
     memory.clear()
-    return []
+    return [], memory
 
 retriever = None
 
@@ -143,8 +141,11 @@ with gr.Blocks() as interface:
     clear_memory_btn = gr.Button("Clear Memory", interactive=True)
     pdf_file = gr.File(label="Upload PDF for Context (Optional)", type="filepath", interactive=True, scale=1)
     submit_btn = gr.Button("Submit")
-    submit_btn.click(gradio_interface, inputs=[user_message, chat_box, pdf_file, use_pdf_context], outputs=chat_box)
-    user_message.submit(gradio_interface, inputs=[user_message, chat_box, pdf_file, use_pdf_context], outputs=chat_box)
-    clear_memory_btn.click(clear_memory, inputs=[], outputs=chat_box)
+    
+    memory_state = gr.State(ConversationBufferMemory())
+    
+    submit_btn.click(gradio_interface, inputs=[user_message, chat_box, memory_state, pdf_file, use_pdf_context], outputs=[chat_box, memory_state])
+    user_message.submit(gradio_interface, inputs=[user_message, chat_box, memory_state, pdf_file, use_pdf_context], outputs=[chat_box, memory_state])
+    clear_memory_btn.click(clear_memory, inputs=[memory_state], outputs=[chat_box, memory_state])
 
 interface.launch()