Update context_graph.py

context_graph.py

@@ -1,28 +1,26 @@
 import json
 import os
 
-class ContextGraph:
-    def __init__(self, agent_id):
-        self.file_path = f"context_{agent_id}.json"
+class ContextMemoryGraph:
+    def __init__(self, file_path="context_graph.json"):
+        self.file_path = file_path
         self._init_graph()
 
     def _init_graph(self):
         if not os.path.exists(self.file_path):
             with open(self.file_path, "w") as f:
-                json.dump({"nodes": [], "edges": []}, f)
+                json.dump({"context_links": {}}, f)
 
-    def add_node(self, node_type, content):
+    def link_context(self, agent_id, key, value):
         with open(self.file_path, "r") as f:
             graph = json.load(f)
-        node_id = len(graph["nodes"]) + 1
-        graph["nodes"].append({"id": node_id, "type": node_type, "content": content})
+        if agent_id not in graph["context_links"]:
+            graph["context_links"][agent_id] = {}
+        graph["context_links"][agent_id][key] = value
         with open(self.file_path, "w") as f:
             json.dump(graph, f, indent=2)
-        return node_id
 
-    def connect(self, node1, node2):
+    def get_context(self, agent_id):
         with open(self.file_path, "r") as f:
             graph = json.load(f)
-        graph["edges"].append({"from": node1, "to": node2})
-        with open(self.file_path, "w") as f:
-            json.dump(graph, f, indent=2)
+        return graph["context_links"].get(agent_id, {})