Update hf_demo.py

hf_demo.py

@@ -364,12 +364,40 @@ class PolicyEngine:
             return True
         return False
 
+# ==============================================================================
+# UPGRADED RAG MEMORY WITH SENTENCE-TRANSFORMERS
+# ==============================================================================
 class RAGMemory:
-    """Persistent RAG memory with SQLite and simple embeddings."""
+    """Persistent RAG memory with SQLite and sentence‑transformer embeddings."""
     def __init__(self):
         self.db_path = f"{settings.data_dir}/memory.db"
         self._init_db()
         self.embedding_cache = {}
+        self._sentence_model = None  # lazy loaded
+
+    def _get_sentence_model(self):
+        """Lazy load the sentence‑transformer model."""
+        if self._sentence_model is None:
+            from sentence_transformers import SentenceTransformer
+            # Using all-MiniLM-L6-v2 – fast and good for semantic similarity
+            self._sentence_model = SentenceTransformer('all-MiniLM-L6-v2')
+        return self._sentence_model
+
+    def _build_incident_text(self, action: str) -> str:
+        """Create a descriptive text from the action."""
+        # You can enrich this with more context (risk level, component, etc.)
+        return f"Action: {action}"
+
+    def _simple_embedding(self, text: str) -> List[float]:
+        """Generate embedding using sentence‑transformer."""
+        if text in self.embedding_cache:
+            return self.embedding_cache[text]
+
+        model = self._get_sentence_model()
+        # encode returns a numpy array; convert to list for JSON storage
+        embedding = model.encode(text, convert_to_numpy=True).tolist()
+        self.embedding_cache[text] = embedding
+        return embedding
 
     def _init_db(self):
         try:
@@ -420,25 +448,12 @@ class RAGMemory:
             if conn:
                 conn.close()
 
-    def _simple_embedding(self, text: str) -> List[float]:
-        if text in self.embedding_cache:
-            return self.embedding_cache[text]
-        words = text.lower().split()
-        trigrams = set()
-        for word in words:
-            for i in range(len(word) - 2):
-                trigrams.add(word[i:i+3])
-        vector = [hash(t) % 1000 / 1000.0 for t in sorted(trigrams)[:100]]
-        while len(vector) < 100:
-            vector.append(0.0)
-        vector = vector[:100]
-        self.embedding_cache[text] = vector
-        return vector
-
     def store_incident(self, action: str, risk_score: float, risk_level: RiskLevel,
                        confidence: float, allowed: bool, gates: List[Dict]):
         action_hash = hashlib.sha256(action.encode()).hexdigest()[:50]
-        embedding = json.dumps(self._simple_embedding(action))
+        # Build a descriptive text and generate embedding
+        incident_text = self._build_incident_text(action)
+        embedding = json.dumps(self._simple_embedding(incident_text))
         try:
             with self._get_db() as conn:
                 conn.execute('''
@@ -462,7 +477,9 @@ class RAGMemory:
             logger.error(f"Failed to store incident: {e}")
 
     def find_similar(self, action: str, limit: int = 5) -> List[Dict]:
-        query_embedding = self._simple_embedding(action)
+        # Build query embedding from the action text
+        query_text = self._build_incident_text(action)
+        query_embedding = self._simple_embedding(query_text)
         try:
             with self._get_db() as conn:
                 cursor = conn.execute('SELECT * FROM incidents ORDER BY timestamp DESC LIMIT 100')