Update context_graph.py

context_graph.py

@@ -1,5 +1,4 @@
-import json
-import os
+import json, os
 
 class ContextGraph:
     def __init__(self, file_path="context_graph.json"):
@@ -7,42 +6,37 @@ class ContextGraph:
         self._init_graph()
 
     def _init_graph(self):
-        # Create the graph file if it doesn't exist
         if not os.path.exists(self.file_path):
             with open(self.file_path, "w") as f:
                 json.dump({"context_links": {}}, f)
 
-    def link_context(self, agent_id, key, value):
-        # Load the current context graph
+    def link_context(self, agent_id, category, key, value):
         with open(self.file_path, "r") as f:
             graph = json.load(f)
 
-        # Add the new link under the correct agent ID
         if agent_id not in graph["context_links"]:
             graph["context_links"][agent_id] = {}
-        graph["context_links"][agent_id][key] = value
+        if category not in graph["context_links"][agent_id]:
+            graph["context_links"][agent_id][category] = {}
+
+        graph["context_links"][agent_id][category][key] = value
 
-        # Save the updated graph
         with open(self.file_path, "w") as f:
             json.dump(graph, f, indent=2)
 
-        os.sync()  # Force Hugging Face to write to disk
-
-    def query_context(self, agent_id, keyword):
-        # Search for any stored memory entries that contain the keyword
+    def query_context(self, agent_id, category=None, keyword=None):
         with open(self.file_path, "r") as f:
             graph = json.load(f)
 
         agent_data = graph["context_links"].get(agent_id, {})
+        if category:
+            data = agent_data.get(category, {})
+        else:
+            # Merge all categories
+            data = {k: v for c in agent_data.values() for k, v in c.items()}
+
         results = [
-            f"{k}: {v}" for k, v in agent_data.items()
-            if keyword.lower() in k.lower() or keyword.lower() in v.lower()
+            f"{k}: {v}" for k, v in data.items()
+            if not keyword or keyword.lower() in k.lower() or keyword.lower() in v.lower()
         ]
-
-        return results if results else ["No matching context found."]
-
-    def get_context(self, agent_id):
-        # Simple alias to maintain backward compatibility with AgentCore
-        with open(self.file_path, "r") as f:
-            graph = json.load(f)
-        return graph["context_links"].get(agent_id, {})
+        return results or ["No context found."]