lore-engine-poc-v3/scripts/02_demo.py

"""02_demo — full tool tour against the loaded codex.

Run:
    python3 scripts/02_demo.py

Calls every read tool shipped in slice 0, 2, and 4 against the
in-memory graph, plus the consistency engine and a small
write/read round-trip. The output is the executable contract
for what each tool returns.
"""

from __future__ import annotations

import argparse
import json
import pickle
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from lore_engine_poc.consistency_tools import (
    get_anachronisms,
    get_contradictions,
    get_orphans,
    latest_run,
    run_consistency_check,
)
from lore_engine_poc.read_tools import (
    ancestors_of,
    descendants_of,
    entities_present,
    entity_context,
    event_chain,
    events_during,
    list_offspring,
    location_hierarchy,
    lookup,
    lore_about,
    timeline,
    true_during,
)
from lore_engine_poc.tools import Graph, was_true_at
from lore_engine_poc.write_tools import add_entity, add_relation


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument(
        "--query",
        action="append",
        default=[],
        metavar="RELATION,SUBJECT,OBJECT,TIME",
        help="Run a specific query. Repeatable.",
    )
    return p.parse_args()


def load_graph():
    pkl = ROOT / "lore_engine_poc" / ".graph.pkl"
    if not pkl.exists():
        print("No cached graph. Run `python3 scripts/01_ingest.py` first.", file=sys.stderr)
        sys.exit(1)
    with open(pkl, "rb") as f:
        data = pickle.load(f)
    return data["graph"]


def run_query(graph, spec: str) -> dict:
    parts = [p.strip() for p in spec.split(",")]
    if len(parts) != 4:
        raise SystemExit(f"--query expects RELATION,SUBJECT,OBJECT,TIME (got {spec!r})")
    relation, subject, object_, at_time = parts
    return was_true_at(graph, relation, subject, object_, at_time)


def _section(title: str, payload) -> None:
    print("=" * 72)
    print(title)
    if isinstance(payload, (dict, list)):
        print(json.dumps(payload, indent=2, ensure_ascii=False, default=str))
    else:
        print(payload)


def main() -> int:
    args = parse_args()
    graph = load_graph()

    queries: list[str] = list(args.query)
    if not queries:
        queries = [
            "MEMBER_OF,Roland Raventhorne,House Raventhorne,3rd_age.year_345",
            "MEMBER_OF,Aldric Raventhorne,House Raventhorne,3rd_age.year_345",
            "SIBLING_OF,Roland Raventhorne,Aldric Raventhorne,3rd_age.year_345",
            "LOCATED_IN,Mardsville,Uul' Dar,3rd_age.year_345",
            "PART_OF,Voldramir,Underdark,3rd_age.year_345",
            "MEMBER_OF,Roland Raventhorne,Iron Mountain Trading Company,3rd_age.year_345",
            "ALLIED_WITH,House Raventhorne,House Quche,3rd_age.year_345",
            "PARENT_OF,Theron Ashveil,Aldric Raventhorne,3rd_age.year_325",
            "PARENT_OF,Theron Ashveil,Aldric Raventhorne,3rd_age.year_400",
            "PARENT_OF,Maric Vyr,Elara Raventhorne,3rd_age.year_250",
            "PARENT_OF,Maric Vyr,Elara Raventhorne,3rd_age.year_340",
        ]

    for q in queries:
        _section(f"was_true_at: {q}", run_query(graph, q))

    # Slice 4 — every read tool, exercised once against the codex.
    _section("lookup('Aldric')", lookup(graph, "Aldric", type_="Person"))
    _section(
        "entity_context('Aldric Raventhorne')",
        entity_context(graph, "Aldric Raventhorne"),
    )
    _section(
        "entity_context('Aldric Raventhorne', at=3rd_age.year_345)",
        entity_context(graph, "Aldric Raventhorne", at_time="3rd_age.year_345"),
    )
    _section(
        "true_during(MEMBER_OF, Aldric, 3rd_age)",
        true_during(graph, "MEMBER_OF", "Aldric Raventhorne", "3rd_age"),
    )
    _section(
        "entities_present(Mardsville, 3rd_age.year_345)",
        entities_present(graph, "Mardsville", "3rd_age.year_345"),
    )
    _section(
        "timeline(Aldric Raventhorne)",
        timeline(graph, "Aldric Raventhorne"),
    )
    _section(
        "list_offspring(Theron Ashveil)",
        list_offspring(graph, "Theron Ashveil"),
    )
    _section(
        "ancestors_of(Aldric Raventhorne)",
        ancestors_of(graph, "Aldric Raventhorne", generations=3),
    )
    _section(
        "descendants_of(Theron Ashveil)",
        descendants_of(graph, "Theron Ashveil", generations=3),
    )
    _section(
        "location_hierarchy(Voldramir, up)",
        location_hierarchy(graph, "Voldramir", direction="up"),
    )
    _section(
        "events_during(3rd_age)",
        events_during(graph, "3rd_age"),
    )
    _section(
        "lore_about(Aldric Raventhorne)",
        lore_about(graph, "Aldric Raventhorne", limit=3),
    )

    # Slice 2: run the consistency engine on the same graph.
    print("=" * 72)
    print("Consistency check:")
    run = run_consistency_check(graph)
    print(f"  run id:        {run.id}")
    print(f"  rules_run:     {run.rules_run}")
    print(f"  contradictions: {run.violations_found}")
    print(f"  anachronisms:   {run.anachronisms_found}")
    print(f"  orphans:        {run.orphans_found}")
    cs = get_contradictions(limit=3)
    if cs:
        print("  Top contradictions:")
        for c in cs:
            print(f"    {c.subject} -[{c.predicate}]-> {c.claim_a} | {c.claim_b}")
    os = get_orphans(limit=3)
    if os:
        print("  Top orphans:")
        for o in os:
            print(f"    {o.entity_name} ({o.entity_type}): {o.reason}")
    latest = latest_run()
    assert latest is not None
    assert latest.id == run.id

    # Slice 4.7 — write tools (round-trip: write, then read).
    _section(
        "write: add_entity + add_relation (round-trip)",
        _write_round_trip(graph),
    )

    return 0


def _write_round_trip(graph: Graph) -> dict:
    """Demonstrate the slice-4.7 write tools: add a Person, add
    a MEMBER_OF edge, then prove the read tools see it.

    Returns the envelope sequence so the demo's stdout is
    self-documenting.
    """
    e1 = add_entity(graph, "Person", "DemoCharacter")
    e2 = add_entity(graph, "Faction", "DemoFaction")
    r1 = add_relation(
        graph, "DemoCharacter", "MEMBER_OF", "DemoFaction",
        valid_from="3rd_age.year_300",
        valid_until="3rd_age.year_330",
    )
    # Read back via lookup + true_during.
    found = lookup(graph, "DemoCharacter")
    td = true_during(graph, "MEMBER_OF", "DemoCharacter", "3rd_age.year_315")
    return {
        "add_entity DemoCharacter": e1,
        "add_entity DemoFaction": e2,
        "add_relation": r1,
        "lookup DemoCharacter": found,
        "true_during DemoCharacter @ year_315": td,
    }


if __name__ == "__main__":
    raise SystemExit(main())