Upload app.py

app.py

@@ -0,0 +1,333 @@
+# app.py - Hybrid Memory-Augmented Graph RAG with Groq
+import hashlib
+import os
+from collections import defaultdict
+from datetime import datetime
+
+import networkx as nx
+import numpy as np
+import spacy
+import streamlit as st
+from langchain_community.vectorstores import FAISS
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_groq import ChatGroq
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from pypdf import PdfReader
+from rank_bm25 import BM25Okapi
+from sentence_transformers import CrossEncoder, SentenceTransformer
+from sklearn.metrics.pairwise import cosine_similarity
+
+# ========================= CONSTANTS =========================
+EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+RERANKER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"
+LLM_MODEL_NAME = "llama-3.3-70b-versatile"
+DEFAULT_TOP_K = 5
+RRF_K = 60
+
+# ========================= HELPERS =========================
+def get_groq_api_key():
+    try:
+        return st.secrets.get("GROQ_API_KEY") or os.getenv("GROQ_API_KEY")
+    except:
+        return os.getenv("GROQ_API_KEY")
+
+def rrf_score(rank: int, k: int = RRF_K) -> float:
+    return 1.0 / (k + rank)
+
+# ========================= MEMORY LAYER =========================
+class HybridMemory:
+    def __init__(self):
+        self.embeddings = HuggingFaceEmbeddings(model_name=EMBED_MODEL_NAME)
+        self.vector_memory = None
+        self.graph_memory = nx.Graph()
+        self.memory_dir = "user_memory"
+        os.makedirs(self.memory_dir, exist_ok=True)
+        self.load_memory()
+
+    def load_memory(self):
+        vector_path = os.path.join(self.memory_dir, "vector_memory")
+        graph_path = os.path.join(self.memory_dir, "memory_graph.gpickle")
+        
+        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):
+            self.graph_memory = nx.read_gpickle(graph_path)
+
+    def save_memory(self):
+        if self.vector_memory:
+            self.vector_memory.save_local(os.path.join(self.memory_dir, "vector_memory"))
+        nx.write_gpickle(self.graph_memory, os.path.join(self.memory_dir, "memory_graph.gpickle"))
+
+    def add_memory(self, text: str, memory_type: str = "fact", metadata: dict = None):
+        if metadata is None:
+            metadata = {}
+        metadata.update({
+            "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()
+
+    def _add_to_graph(self, text: str):
+        try:
+            nlp = st.session_state.get("nlp")
+            if nlp:
+                doc = nlp(text[:30000])
+                entities = [ent.text.strip().lower() for ent in doc.ents if len(ent.text.strip()) > 1]
+                for ent in entities:
+                    if not self.graph_memory.has_node(ent):
+                        self.graph_memory.add_node(ent, count=1)
+                    else:
+                        self.graph_memory.nodes[ent]["count"] = self.graph_memory.nodes[ent].get("count", 0) + 1
+                for i in range(len(entities)):
+                    for j in range(i + 1, len(entities)):
+                        a, b = entities[i], entities[j]
+                        if self.graph_memory.has_edge(a, b):
+                            self.graph_memory[a][b]["weight"] += 1
+                        else:
+                            self.graph_memory.add_edge(a, b, weight=1)
+        except:
+            pass
+
+    def retrieve_memory(self, query: str, top_k: int = 6):
+        if not self.vector_memory:
+            return []
+        docs = self.vector_memory.similarity_search(query, k=top_k)
+        return [{"text": doc.page_content, "metadata": doc.metadata} for doc in docs]
+
+    def get_stats(self):
+        return {
+            "vector_memories": len(self.vector_memory.index_to_docstore_id) if self.vector_memory else 0,
+            "graph_nodes": self.graph_memory.number_of_nodes(),
+            "graph_edges": self.graph_memory.number_of_edges()
+        }
+
+# ========================= ORIGINAL PIPELINE FUNCTIONS =========================
+@st.cache_resource
+def load_embed_model(): return SentenceTransformer(EMBED_MODEL_NAME)
+@st.cache_resource
+def load_reranker(): return CrossEncoder(RERANKER_MODEL)
+@st.cache_resource
+def load_lc_embeddings(): return HuggingFaceEmbeddings(model_name=EMBED_MODEL_NAME)
+@st.cache_resource
+def load_spacy():
+    try:
+        return spacy.load("en_core_web_sm")
+    except OSError:
+        from spacy.cli import download
+        download("en_core_web_sm")
+        return spacy.load("en_core_web_sm")
+
+def load_pdf_documents(uploaded_files):
+    docs = []
+    for f in uploaded_files:
+        reader = PdfReader(f)
+        text = "\n".join([page.extract_text() or "" for page in reader.pages])
+        if text.strip():
+            docs.append({"text": text, "source": f.name})
+    return docs
+
+def hybrid_chunking(text, embed_model, base=300, final=1000, thresh=0.6):
+    splitter = RecursiveCharacterTextSplitter(chunk_size=base, chunk_overlap=50)
+    base_chunks = splitter.split_text(text)
+    if not base_chunks: return []
+    embeddings = embed_model.encode(base_chunks, show_progress_bar=False)
+    chunks, current = [], base_chunks[0]
+    anchor = embeddings[0].reshape(1, -1)
+    for i in range(1, len(base_chunks)):
+        next_emb = embeddings[i].reshape(1, -1)
+        sim = cosine_similarity(anchor, next_emb)[0][0]
+        if sim > thresh and len(current) + len(base_chunks[i]) < final:
+            current += " " + base_chunks[i]
+        else:
+            chunks.append(current)
+            current = base_chunks[i]
+            anchor = next_emb
+    chunks.append(current)
+    return chunks
+
+def chunk_documents(docs, embed_model, method, **kwargs):
+    result = []
+    for doc in docs:
+        if method == "Hybrid (Recursive + Semantic)":
+            texts = hybrid_chunking(doc["text"], embed_model, **kwargs)
+        else:
+            splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+            texts = splitter.split_text(doc["text"])
+        for t in texts:
+            result.append({"text": t, "source": doc["source"]})
+    return result
+
+def build_knowledge_graph(chunks, nlp):
+    G = nx.Graph()
+    entity_to_chunks = defaultdict(list)
+    for idx, chunk in enumerate(chunks):
+        doc = nlp(chunk["text"][:50000])
+        entities = [ent.text.strip().lower() for ent in doc.ents if len(ent.text.strip()) > 1]
+        for e in set(entities):
+            if not G.has_node(e):
+                G.add_node(e)
+            entity_to_chunks[e].append(idx)
+        for i in range(len(entities)):
+            for j in range(i+1, len(entities)):
+                a, b = entities[i], entities[j]
+                if G.has_edge(a, b):
+                    G[a][b]["weight"] += 1
+                else:
+                    G.add_edge(a, b, weight=1)
+    return G, dict(entity_to_chunks)
+
+def build_indexes(chunks, lc_embeddings):
+    texts = [c["text"] for c in chunks]
+    bm25 = BM25Okapi([t.lower().split() for t in texts])
+    faiss_index = FAISS.from_texts(texts, 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)
+    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]
+    faiss_docs = faiss_index.similarity_search(query, k=top_k*3)
+
+    candidates = []
+    seen = set()
+    for i in bm25_idx[:top_k*2]:
+        c = chunks[i]
+        if c["text"] not in seen:
+            candidates.append(c)
+            seen.add(c["text"])
+    for doc in faiss_docs:
+        c = chunk_lookup.get(doc.page_content)
+        if c and c["text"] not in seen:
+            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)
+
+    scored = sorted(zip(scores, all_texts), reverse=True)
+    final_texts = [t for _, t in scored[:top_k]]
+
+    final_chunks = []
+    for t in final_texts:
+        for c in candidates:
+            if c["text"] == t:
+                final_chunks.append(c)
+                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):
+    context = "\n\n---\n\n".join([f"[Source: {c['source']}]\n{c['text']}" for c in retrieved])
+    prompt = f"""Answer using ONLY the provided context. Cite sources.
+Context:
+{context}
+
+Question: {query}
+Answer:"""
+    return llm.invoke(prompt).content
+
+# ========================= STREAMLIT APP =========================
+st.set_page_config(page_title="Hybrid Memory Graph RAG", page_icon="🧠", layout="wide")
+st.title("🧠 Hybrid Memory-Augmented Graph RAG")
+
+# Session State
+if "memory" not in st.session_state:
+    st.session_state.memory = HybridMemory()
+if "chunks" not in st.session_state:
+    st.session_state.chunks = None
+
+# Sidebar
+with st.sidebar:
+    st.header("Configuration")
+    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!")
+
+# Load models
+embed_model = load_embed_model()
+reranker = load_reranker()
+nlp = load_spacy()
+lc_embeddings = load_lc_embeddings()
+
+groq_key = get_groq_api_key()
+if not groq_key:
+    st.error("Groq API key not found!")
+    st.stop()
+
+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)
+
+if uploaded_files and st.session_state.chunks is None:
+    with st.spinner("Processing documents..."):
+        docs = load_pdf_documents(uploaded_files)
+        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)
+        
+        st.session_state.chunks = chunks
+        st.session_state.bm25 = bm25
+        st.session_state.faiss_index = faiss_index
+        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")
+
+# Query
+query = st.text_input("Ask a question about your documents:")
+
+if query and st.session_state.chunks:
+    with st.spinner("Retrieving from Documents + Memory..."):
+        retrieved, debug = hybrid_memory_retrieve(
+            query, st.session_state.chunks, st.session_state.bm25,
+            st.session_state.faiss_index, st.session_state.chunk_lookup,
+            st.session_state.graph, st.session_state.entity_to_chunks,
+            reranker, nlp, st.session_state.memory, top_k=top_k, hops=graph_hops
+        )
+    
+    with st.spinner("Generating answer..."):
+        answer = generate_answer(llm, query, retrieved)
+    
+    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")
+
+    with st.expander("📊 Memory Stats"):
+        st.json(st.session_state.memory.get_stats())
+
+    with st.expander("🔍 Retrieved Sources"):
+        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.caption("Hybrid Graph RAG + Persistent Memory • Powered by Groq + NetworkX + FAISS")