context_graph.py

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):
        if not os.path.exists(self.file_path):
            with open(self.file_path, "w") as f:
                json.dump({"links": {}}, f)

    def link_context(self, agent_id, key, value):
        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
        with open(self.file_path, "w") as f:
            json.dump(graph, f, indent=2)
        print(f"[ContextGraph] Linked ({key} → {value}) for {agent_id}")

    def get_context(self, agent_id):
        with open(self.file_path, "r") as f:
            graph = json.load(f)
        return graph["links"].get(agent_id, {})