Update context_graph.py

context_graph.py

@@ -1,33 +1,37 @@
-import json, os
-from collections import defaultdict
+import numpy as np
+from sentence_transformers import SentenceTransformer
 
 class ContextGraph:
-    def __init__(self, file_path="knowledge_graph.json"):
-        self.file_path = file_path
-        self.graph = self.load_graph()
+    def __init__(self):
+        self.model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
-    def load_graph(self):
-        if os.path.exists(self.file_path):
-            with open(self.file_path, "r") as f:
-                return json.load(f)
-        return defaultdict(list)
+    def encode(self, text):
+        return self.model.encode(text)
 
-    def save_graph(self):
-        with open(self.file_path, "w") as f:
-            json.dump(self.graph, f, indent=2)
+    def connect(self, memory_texts, threshold=0.35):
+        """
+        Create relationships between memory nodes based on vector similarity.
+        """
+        vectors = [self.encode(t) for t in memory_texts]
+        relationships = []
 
-    def link(self, subject, relation, obj):
-        """Create a semantic relationship like 'Tirrek' -> 'builds' -> 'Motion Empire'."""
-        self.graph[subject].append({"relation": relation, "object": obj})
-        self.save_graph()
+        for i, vec1 in enumerate(vectors):
+            for j, vec2 in enumerate(vectors):
+                if i != j:
+                    similarity = np.dot(vec1, vec2) / (np.linalg.norm(vec1) * np.linalg.norm(vec2))
+                    if similarity >= threshold:
+                        relationships.append((memory_texts[i], memory_texts[j], float(similarity)))
 
-    def query(self, subject):
-        if subject not in self.graph:
-            return f"No links found for {subject}."
-        return "\n".join([f"{subject} —{rel['relation']}→ {rel['object']}" for rel in self.graph[subject]])
-    def retrieve_relevant(self, query_vector, top_k=5):
-    scored = []
-    for mem in self.memories:
-        sim = float(np.dot(mem['vector'], query_vector) / (np.linalg.norm(mem['vector']) * np.linalg.norm(query_vector)))
-        scored.append((mem['text'], sim * mem['weight']))
-    return sorted(scored, key=lambda x: x[1], reverse=True)[:top_k]
+        return relationships
+
+    def score_context(self, query, memories):
+        """
+        Score stored memories against the new query.
+        """
+        query_vec = self.encode(query)
+        scored = []
+        for mem, vec in memories:
+            score = np.dot(query_vec, vec) / (np.linalg.norm(query_vec) * np.linalg.norm(vec))
+            scored.append((mem, float(score)))
+
+        return sorted(scored, key=lambda x: x[1], reverse=True)