Vedanshipanda/layer10-api

1import json2import networkx as nx3from networkx.readwrite import json_graph4import os5 6def canonicalize(name):7    return str(name).lower().strip() if name else ""8 9def build_memory_graph():10    print("Building Memory Graph...")11    path = os.path.join("data", "extracted_graph.json")12    if not os.path.exists(path):13        print("❌ extracted_graph.json not found!")14        return15    16    with open(path, "r", encoding="utf-8") as f: 17        raw_data = json.load(f)18        19    G = nx.MultiDiGraph() # MultiDiGraph allows parallel edges20 21    for item in raw_data:22        source_url = item.get("source_id", "")23        graph_data = item.get("graph_data", {})24 25        # 1. Add Explicit Entities26        for entity in graph_data.get("entities", []):27            node_id = canonicalize(entity.get("name"))28            if node_id:29                if not G.has_node(node_id):30                    # Added display_name for the UI31                    G.add_node(node_id, 32                               display_name=entity.get("name"), 33                               type=entity.get("type", "Unknown"), 34                               description=entity.get("description", ""), 35                               evidence=[])36                37                G.nodes[node_id]["evidence"].append({38                    "source": source_url, 39                    "quote": entity.get("text_excerpt", "")40                })41 42        # 2. Add Relationships (LOOSE MODE)43        for rel in graph_data.get("relationships", []):44            src, tgt = canonicalize(rel.get("source")), canonicalize(rel.get("target"))45            46            if src and tgt:47                # If the LLM found a relationship but missed the entity, create a placeholder node48                if not G.has_node(src):49                    G.add_node(src, display_name=rel.get("source"), type="Unknown", evidence=[])50                if not G.has_node(tgt):51                    G.add_node(tgt, display_name=rel.get("target"), type="Unknown", evidence=[])52                53                # Changed 'type' to 'relationship' to match the Streamlit frontend54                G.add_edge(src, tgt, 55                           relationship=rel.get("relation_type", "related"), 56                           source_id=source_url)57 58    # Save for Frontend (Ensure Streamlit is reading from knowledge_graph.json)59    output_path = os.path.join("data", "knowledge_graph.json")60    with open(output_path, "w", encoding="utf-8") as f:61        json.dump(json_graph.node_link_data(G), f, indent=4)62        63    print(f"✅ Graph Ready: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges.")64 65if __name__ == "__main__":66    build_memory_graph()