Update app.py

app.py

@@ -41,6 +41,7 @@ def get_groq_api_key():
 def rrf_score(rank: int, k: int = RRF_K) -> float:
     return 1.0 / (k + rank)
 
+# ========================= MEMORY LAYER =========================
 # ========================= MEMORY LAYER =========================
 class HybridMemory:
     def __init__(self):
@@ -77,32 +78,47 @@ class HybridMemory:
             "timestamp": datetime.now().isoformat(),
             "type": memory_type,
         })
+
+        # Add to Vector Memory
         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])
+
+        # Add to Graph Memory
         self._add_to_graph(text)
         self.save_memory()
 
     def _add_to_graph(self, text: str):
+        """Safely add entities to graph"""
         try:
+            # Get nlp from session_state or load it directly
             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)
+            if not nlp:
+                nlp = load_spacy()  # fallback
+                st.session_state.nlp = nlp
+
+            doc = nlp(text[:30000])
+            entities = [ent.text.strip().lower() for ent in doc.ents if len(ent.text.strip()) > 1]
+            
+            if not entities:
+                return  # No entities found
+
+            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.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:
+                        self.graph_memory.add_edge(a, b, weight=1)
+        except Exception as e:
+            # st.warning(f"Graph memory update skipped: {e}")  # uncomment for debug
             pass
 
     def retrieve_memory(self, query: str, top_k: int = 6):