prompt.py

import json
import networkx as nx

def load_knowledge_graph(events_file="data/parsed_event_data.json", 
                        characters_file="data/parsed_character_metadata.json"):
    """Recreate the knowledge graph from saved data"""
    
    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")
    
    return G

def generate_scene_prompts(output_file="data/event_prompts.json"):
    """Generate enhanced scene-aware prompts for each event"""
    
    G = load_knowledge_graph()
    event_prompts = {}
    
    for node, data in G.nodes(data=True):
        if data.get("type") == "event":
            event_name = node
            event_desc = data.get("description", "")
            
            characters = []
            objects = []
            environments = []
            
            # Get objects and environments from outgoing edges
            for neighbor in G.neighbors(node):
                edge_data = G.get_edge_data(node, neighbor)
                if edge_data:
                    for k, v in edge_data.items():
                        relation = v.get("relation", "")
                        ntype = G.nodes[neighbor].get("type", "")
                        if relation == "has_object" and ntype == "object":
                            objects.append(neighbor)
                        elif relation == "is_in" and ntype == "environment":
                            environments.append(neighbor)
            
            # Get characters from incoming edges
            for pred in G.predecessors(node):
                edge_data = G.get_edge_data(pred, node)
                if edge_data:
                    for k, v in edge_data.items():
                        relation = v.get("relation", "")
                        ntype = G.nodes[pred].get("type", "")
                        if relation == "participates_in" and ntype == "character":
                            char_desc = G.nodes[pred].get("description", pred)
                            characters.append(f"{pred} ({char_desc})")
            
            # Build enhanced prompt
            prompt_template = f"""Scene: {event_desc}

Visual Elements:
- Characters: {', '.join(characters) if characters else 'None specified'}
- Key Objects: {', '.join(objects) if objects else 'None specified'}  
- Environment/Setting: {', '.join(environments) if environments else 'None specified'}

Generate a vivid, detailed scene description that captures the mood, atmosphere, and visual details of this moment in the story. Focus on creating an immersive experience that brings the scene to life."""

            event_prompts[event_name] = prompt_template
    
    # Save prompts
    with open(output_file, "w") as f:
        json.dump(event_prompts, f, indent=2)
    
    print(f"Enhanced scene prompts saved to {output_file}")
    return event_prompts

def display_prompts():
    """Display all generated prompts"""
    prompts = generate_scene_prompts()
    
    for event, prompt in prompts.items():
        print(f"Event: {event}")
        print(f"Prompt:\n{prompt}")
        print("-" * 80)

if __name__ == "__main__":
    display_prompts()