Update app.py

Add memory context

app.py

@@ -108,7 +108,7 @@ def retrieve_context_per_file(question, state, k=10):
         # Get top-k indices for this file
         top_k_idx = distances.argsort()[: min(k, len(chunks))]
 
-        # Collect chunks (optionally tag with filename)
+        # Tag chunks with filename for transparency
         selected = [f"[{filename}]\n{chunks[i]}" for i in top_k_idx]
         all_selected.extend(selected)
 
@@ -158,7 +158,6 @@ def load_pdf(files, state):
     """
     files: current list from the UI (after add/remove).
     We rebuild state from scratch to stay in sync.
-    State structure:
 
     state = {
         "files": {
@@ -169,7 +168,9 @@ def load_pdf(files, state):
             ...
         },
         "pdf_count": int,
-        "total_chunks": int
+        "total_chunks": int,
+        "history": [...],
+        "facts": [...]
     }
     """
     # If user cleared all files
@@ -270,14 +271,58 @@ def answer_question(question, state):
     else:
         print("HF_TOKEN is None or empty")
 
-    if not question.strip():
+    question = question.strip()
+    if not question:
         return "Please enter a question."
 
     if not state["files"]:
         return "Please upload at least one PDF first."
 
-    context = retrieve_context_per_file(question, state, k=10)
-    answer = generate_answer(question, context)
+    # Simple persistent memory hook: "Remember that ..."
+    lower_q = question.lower()
+    if lower_q.startswith("remember that"):
+        fact = question[len("remember that"):].strip(" .")
+        if fact:
+            state["facts"].append(fact)
+            return f"Okay, I will remember that {fact}."
+        else:
+            return "What should I remember?"
+
+    # RAG context (per-file balanced retrieval)
+    rag_context = retrieve_context_per_file(question, state, k=10)
+
+    # Conversation memory (last 5 turns)
+    memory_context = ""
+    if state["history"]:
+        memory_context = "\n\n".join(
+            [
+                f"User: {h['user']}\nAssistant: {h['assistant']}"
+                for h in state["history"][-5:]
+            ]
+        )
+
+    # Persistent facts
+    facts_context = ""
+    if state["facts"]:
+        facts_context = "User Facts:\n" + "\n".join(state["facts"]) + "\n\n"
+
+    # Merge all contexts
+    full_context_parts = []
+    if facts_context:
+        full_context_parts.append(facts_context)
+    if memory_context:
+        full_context_parts.append("Conversation Memory:\n" + memory_context)
+    if rag_context:
+        full_context_parts.append("Document Context:\n" + rag_context)
+
+    full_context = "\n\n".join(full_context_parts)
+
+    # Generate answer
+    answer = generate_answer(question, full_context)
+
+    # Save conversational memory
+    state["history"].append({"user": question, "assistant": answer})
+
     return answer
 
 
@@ -293,6 +338,8 @@ with gr.Blocks(title="PDF Q&A Chatbot") as demo:
             "files": {},          # filename -> {chunks, embeddings}
             "pdf_count": 0,
             "total_chunks": 0,
+            "history": [],        # conversational memory
+            "facts": [],          # persistent user facts
         }
     )