Update context_graph.py

context_graph.py

@@ -2,31 +2,40 @@ import json
 import os
 
 class ContextGraph:
-    """
-    Stores relationships between prompts and responses.
-    Think of it as a 'map' of what the agent has learned.
-    """
-
     def __init__(self, file_path="context_graph.json"):
         self.file_path = file_path
         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({"links": {}}, f)
+                json.dump({"context_links": {}}, f)
 
     def link_context(self, agent_id, key, value):
+        # Load the current context graph
         with open(self.file_path, "r") as f:
             graph = json.load(f)
-        if agent_id not in graph["links"]:
-            graph["links"][agent_id] = {}
-        graph["links"][agent_id][key] = value
+
+        # 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
+
+        # Save the updated graph
         with open(self.file_path, "w") as f:
             json.dump(graph, f, indent=2)
-        print(f"[ContextGraph] Linked ({key} → {value}) for {agent_id}")
+        os.sync()  # Force Hugging Face to write to disk
 
-    def get_context(self, agent_id):
+    def query_context(self, agent_id, keyword):
+        # Search for any stored memory entries that contain the keyword
         with open(self.file_path, "r") as f:
             graph = json.load(f)
-        return graph["links"].get(agent_id, {})
+
+        agent_data = graph["context_links"].get(agent_id, {})
+        results = [
+            f"{k}: {v}" for k, v in agent_data.items()
+            if keyword.lower() in k.lower() or keyword.lower() in v.lower()
+        ]
+
+        return results if results else ["No matching context found."]