Update app.py

app.py

@@ -18,6 +18,12 @@ from rank_bm25 import BM25Okapi
 from sentence_transformers import CrossEncoder, SentenceTransformer
 from sklearn.metrics.pairwise import cosine_similarity
 
+# For DOCX support
+try:
+    from docx import Document
+except ImportError:
+    pass  # Will show warning if missing
+
 # ========================= CONSTANTS =========================
 EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
 RERANKER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
@@ -52,7 +58,6 @@ class HybridMemory:
     def load_memory(self):
         vector_path = os.path.join(self.memory_dir, "vector_memory")
         graph_path = os.path.join(self.memory_dir, "memory_graph.pkl")
-
         if os.path.exists(vector_path):
             self.vector_memory = FAISS.load_local(vector_path, self.embeddings, allow_dangerous_deserialization=True)
         if os.path.exists(graph_path):
@@ -72,12 +77,10 @@ class HybridMemory:
             "timestamp": datetime.now().isoformat(),
             "type": memory_type,
         })
-
         if self.vector_memory is None:
             self.vector_memory = FAISS.from_texts([text], self.embeddings, metadatas=[metadata])
         else:
             self.vector_memory.add_texts([text], metadatas=[metadata])
-
         self._add_to_graph(text)
         self.save_memory()
 
@@ -141,13 +144,31 @@ def load_spacy():
         download("en_core_web_sm")
         return spacy.load("en_core_web_sm")
 
-def load_pdf_documents(uploaded_files):
+def load_documents(uploaded_files):
+    """Support PDF, TXT, and DOCX"""
     docs = []
     for f in uploaded_files:
-        reader = PdfReader(f)
-        text = "\n".join([page.extract_text() or "" for page in reader.pages])
+        text = ""
+        source = f.name
+        
+        try:
+            if f.name.lower().endswith('.pdf'):
+                reader = PdfReader(f)
+                text = "\n".join([page.extract_text() or "" for page in reader.pages])
+            elif f.name.lower().endswith('.txt'):
+                text = f.read().decode("utf-8")
+            elif f.name.lower().endswith(('.docx', '.doc')):
+                doc = Document(f)
+                text = "\n".join([paragraph.text for paragraph in doc.paragraphs])
+            else:
+                st.warning(f"Unsupported file type: {f.name}")
+                continue
+        except Exception as e:
+            st.warning(f"Error reading {f.name}: {str(e)}")
+            continue
+
         if text.strip():
-            docs.append({"text": text, "source": f.name})
+            docs.append({"text": text, "source": source})
     return docs
 
 def hybrid_chunking(text, embed_model, base=300, final=1000, thresh=0.6):
@@ -207,10 +228,7 @@ def build_indexes(chunks, lc_embeddings):
     chunk_lookup = {c["text"]: c for c in chunks}
     return bm25, faiss_index, chunk_lookup
 
-def hybrid_memory_retrieve(query, chunks, bm25, faiss_index, chunk_lookup,
-                           graph, entity_to_chunks, reranker, nlp, memory,
-                           top_k=5, hops=2):
-    # Document + Graph Retrieval (simplified)
+def hybrid_memory_retrieve(query, chunks, bm25, faiss_index, chunk_lookup, graph, entity_to_chunks, reranker, nlp, memory, top_k=5, hops=2):
     texts = [c["text"] for c in chunks]
     bm25_scores = bm25.get_scores(query.lower().split())
     bm25_idx = np.argsort(bm25_scores)[::-1][:top_k*3]
@@ -229,10 +247,8 @@ def hybrid_memory_retrieve(query, chunks, bm25, faiss_index, chunk_lookup,
             candidates.append(c)
             seen.add(c["text"])
 
-    # Memory Retrieval
     mem_items = memory.retrieve_memory(query, top_k=top_k*2)
 
-    # Reranking
     all_texts = [c["text"] for c in candidates] + [m["text"] for m in mem_items]
     pairs = [(query, t) for t in all_texts]
     scores = reranker.predict(pairs)
@@ -248,7 +264,6 @@ def hybrid_memory_retrieve(query, chunks, bm25, faiss_index, chunk_lookup,
                 break
         else:
             final_chunks.append({"text": t, "source": "🧠 Memory"})
-
     return final_chunks[:top_k], {"memory_used": len(mem_items)}
 
 def generate_answer(llm, query, retrieved):
@@ -278,7 +293,6 @@ with st.sidebar:
     chunk_method = st.selectbox("Chunking Method", ["Hybrid (Recursive + Semantic)", "Recursive Only"])
     top_k = st.slider("Top-K Results", 3, 15, DEFAULT_TOP_K)
     graph_hops = st.slider("Graph Hops", 1, 3, 2)
-    
     if st.button("Clear Memory"):
         st.session_state.memory = HybridMemory()
         st.success("Memory cleared!")
@@ -296,12 +310,19 @@ if not groq_key:
 
 llm = ChatGroq(temperature=0.2, model_name=LLM_MODEL_NAME, groq_api_key=groq_key)
 
-# File Upload
-uploaded_files = st.file_uploader("Upload PDFs", type=["pdf"], accept_multiple_files=True)
+# File Upload - Now supports PDF, TXT, DOCX
+uploaded_files = st.file_uploader(
+    "Upload Documents (PDF, TXT, DOCX)", 
+    type=["pdf", "txt", "docx"], 
+    accept_multiple_files=True
+)
 
 if uploaded_files and st.session_state.chunks is None:
     with st.spinner("Processing documents..."):
-        docs = load_pdf_documents(uploaded_files)
+        docs = load_documents(uploaded_files)
+        if not docs:
+            st.error("No readable text found in uploaded files.")
+            st.stop()
         chunks = chunk_documents(docs, embed_model, chunk_method)
         bm25, faiss_index, chunk_lookup = build_indexes(chunks, lc_embeddings)
         graph, entity_to_chunks = build_knowledge_graph(chunks, nlp)
@@ -312,7 +333,7 @@ if uploaded_files and st.session_state.chunks is None:
         st.session_state.chunk_lookup = chunk_lookup
         st.session_state.graph = graph
         st.session_state.entity_to_chunks = entity_to_chunks
-        st.success(f"Loaded {len(chunks)} chunks")
+        st.success(f"Loaded {len(chunks)} chunks from {len(docs)} document(s)")
 
 # Query
 query = st.text_input("Ask a question about your documents:")
@@ -332,7 +353,6 @@ if query and st.session_state.chunks:
     st.subheader("Answer")
     st.write(answer)
 
-    # Store new memories
     if retrieved:
         summary = f"Q: {query}\nA: {answer[:300]}..."
         st.session_state.memory.add_memory(summary, memory_type="conversation")
@@ -344,8 +364,7 @@ if query and st.session_state.chunks:
         for i, chunk in enumerate(retrieved, 1):
             st.markdown(f"**{i}.** {chunk['source']}")
             st.write(chunk["text"][:500] + "..." if len(chunk["text"]) > 500 else chunk["text"])
-
 else:
-    st.info("Upload PDFs and ask a question to begin.")
+    st.info("Upload documents (PDF, TXT, DOCX) and ask a question to begin.")
 
 st.caption("Hybrid Graph RAG + Persistent Memory • Powered by Groq + NetworkX + FAISS")