knowledge_graph.py

import json
import networkx as nx
from pyvis.network import Network

def create_knowledge_graph(events_file="data/parsed_event_data.json", 
                          characters_file="data/parsed_character_metadata.json",
                          output_file="data/knowledge_graph.html"):
    """
    Create a knowledge graph from parsed events and character data
    
    Args:
        events_file (str): Path to parsed event data JSON
        characters_file (str): Path to character metadata JSON  
        output_file (str): Path to save the HTML knowledge graph
    """
    
    with open(events_file, "r") as f:
        parsed_data = json.load(f)

    with open(characters_file, "r") as f:
        character_metadata = json.load(f)

    story_title = list(parsed_data.keys())[0]
    events = parsed_data[story_title]

    character_metadata.pop("Story_name", None)

    G = nx.MultiDiGraph() 

    for char, desc in character_metadata.items():
        G.add_node(char, type="character", description=desc)

    for event in events:
        event_name = event.get("event_name", "")
        event_desc = event.get("description", "")
        G.add_node(event_name, type="event", description=event_desc)

        for obj in event.get("objects", []):
            G.add_node(obj, type="object")
            G.add_edge(event_name, obj, relation="has_object")

        for env in event.get("environment", []):
            G.add_node(env, type="environment")
            G.add_edge(event_name, env, relation="is_in")

        for actor in event.get("actors", []):
            matched = False
            for char in character_metadata:
                if char.lower() in actor.lower():
                    G.add_edge(char, event_name, relation="participates_in")
                    matched = True
                    break
            if not matched:
                G.add_node(actor, type="actor")
                G.add_edge(actor, event_name, relation="participates_in")

    for char, desc in character_metadata.items():
        G.nodes[char]["description"] = desc

    net = Network(height="800px", width="100%", notebook=True, directed=True, cdn_resources="in_line")
    net.from_nx(G)

    color_map = {
        "character": "#ffb347",
        "event": "#87ceeb",
        "object": "#90ee90",
        "environment": "#f08080",
        "actor": "#dda0dd"
    }
    
    for node in net.nodes:
        ntype = G.nodes[node["id"]].get("type", "actor")
        node["color"] = color_map.get(ntype, "#cccccc")
        node["title"] = G.nodes[node["id"]].get("description", "")

    net.write_html(output_file)
    print(f"Knowledge graph saved as {output_file}")

if __name__ == "__main__":
    create_knowledge_graph()