lore-engine-poc/plugins/world.py

"""
world plugin — pure Neo4j queries.

Tools:
- entity_context(name, world_id?): one-hop summary of a Person / Faction / Location / Item.
- was_true_at(relation, subject, object, at_time, world_id?): time-bounded edge lookup.
- state_at(entity, at_time, world_id?): comprehensive snapshot of an entity at a time.
- list_worlds(): distinct world_id values present in the graph.

Per docs/01-ontology.md, every node carries a `world_id` namespace and
all read tools filter by it. The default is "default", preserving v1
behaviour for callers that don't pass the parameter.
"""
from server import get_neo4j, REGISTRY

DEFAULT_WORLD = "default"


def _world(args):
    """Return the world_id from args, defaulting to DEFAULT_WORLD."""
    return args.get("world_id") or DEFAULT_WORLD


def _q(query, params=None):
    """Run a single read query against Neo4j, return list of dicts."""
    driver = get_neo4j()
    with driver.session() as s:
        result = s.run(query, params or {})
        return [dict(r) for r in result]


@REGISTRY.tool(
    name="entity_context",
    description="One-hop summary of a named entity (Person, Faction, Location, Item, Event) in a given world. Returns labels, properties, and immediate relations.",
    input_schema={
        "type": "object",
        "properties": {
            "name": {"type": "string", "description": "Entity name to look up"},
            "world_id": {"type": "string", "default": DEFAULT_WORLD,
                         "description": "World namespace; defaults to 'default'"},
        },
        "required": ["name"],
    },
)
def entity_context(args):
    name = args["name"]
    world_id = _world(args)
    rows = _q("""
        MATCH (e)
        WHERE (e:Person OR e:Faction OR e:Location OR e:Item OR e:Event)
          AND (e.name = $name OR e.id = $name)
          AND coalesce(e.world_id, $fallback_world) = $world_id
        OPTIONAL MATCH (e)-[r]->(other)
        WHERE type(r) IN ['MEMBER_OF','RULED','LOCATED_IN','PART_OF','PARENT_OF','SPOUSE_OF','POSSESSES','PARTICIPATED_IN']
        RETURN e, labels(e) AS labels,
               collect(DISTINCT {rel: type(r), to: other.name, to_id: other.id}) AS relations
        LIMIT 1
    """, {"name": name, "world_id": world_id, "fallback_world": DEFAULT_WORLD})
    if not rows:
        return {"found": False, "name": name, "world_id": world_id}
    r = rows[0]
    e = r["e"]
    return {
        "found": True,
        "name": e.get("name"),
        "id": e.get("id"),
        "world_id": e.get("world_id") or DEFAULT_WORLD,
        "labels": r["labels"],
        "properties": {k: v for k, v in dict(e).items() if not k.startswith("_")},
        "relations": [rel for rel in r["relations"] if rel.get("to")],
    }


@REGISTRY.tool(
    name="was_true_at",
    description="Check whether a typed relation was true between subject and object at a given in-fiction time, in a given world. Times use the canonical {era}.{year} format, e.g. '2nd_age.year_340'.",
    input_schema={
        "type": "object",
        "properties": {
            "relation": {"type": "string", "description": "Edge type, e.g. RULED, ALLIED_WITH, MEMBER_OF"},
            "subject":  {"type": "string"},
            "object":   {"type": "string"},
            "at_time":  {"type": "string", "description": "Canonical time string, e.g. '2nd_age.year_340'"},
            "world_id": {"type": "string", "default": DEFAULT_WORLD},
        },
        "required": ["relation", "subject", "object", "at_time"],
    },
)
def was_true_at(args):
    world_id = _world(args)
    rows = _q("""
        MATCH (s {name: $subject, world_id: $world_id})-[r:`%s`]->(o {name: $object, world_id: $world_id})
        WHERE r.valid_from IS NULL OR $at_time >= r.valid_from
          AND r.valid_until IS NULL OR $at_time <= r.valid_until
        RETURN r, s, o
    """ % args["relation"], {
        "subject": args["subject"], "object": args["object"],
        "at_time": args["at_time"], "world_id": world_id,
    })
    if not rows:
        return {"was_true": False, "relation": args["relation"],
                "subject": args["subject"], "object": args["object"],
                "at_time": args["at_time"], "world_id": world_id}
    r = rows[0]["r"]
    return {
        "was_true": True,
        "relation": args["relation"],
        "valid_from": r.get("valid_from"),
        "valid_until": r.get("valid_until"),
        "world_id": world_id,
    }


@REGISTRY.tool(
    name="state_at",
    description="Snapshot of an entity at a given in-fiction time: who/what they were allied with, where they were located, what they held. Scoped to a world.",
    input_schema={
        "type": "object",
        "properties": {
            "entity":  {"type": "string"},
            "at_time": {"type": "string", "description": "Canonical time string, e.g. '2nd_age.year_340'"},
            "world_id": {"type": "string", "default": DEFAULT_WORLD},
        },
        "required": ["entity", "at_time"],
    },
)
def state_at(args):
    world_id = _world(args)
    rows = _q("""
        MATCH (e {name: $entity, world_id: $world_id})
        WHERE e:Person OR e:Faction OR e:Location OR e:Item
        OPTIONAL MATCH (e)-[r]->(other)
        WHERE type(r) IN ['MEMBER_OF','RULED','LOCATED_IN','PART_OF','POSSESSES','ALLIED_WITH','ENEMY_OF']
          AND (r.valid_from IS NULL OR $at_time >= r.valid_from)
          AND (r.valid_until IS NULL OR $at_time <= r.valid_until)
        RETURN e, labels(e) AS labels,
               collect(DISTINCT {rel: type(r), to: other.name}) AS active_relations
        LIMIT 1
    """, {"entity": args["entity"], "at_time": args["at_time"], "world_id": world_id})
    if not rows:
        return {"found": False, "entity": args["entity"], "world_id": world_id}
    r = rows[0]
    return {
        "entity": r["e"].get("name"),
        "at_time": args["at_time"],
        "world_id": world_id,
        "labels": r["labels"],
        "active_relations": [x for x in r["active_relations"] if x.get("to")],
    }


@REGISTRY.tool(
    name="list_worlds",
    description="Admin: list the distinct world_id values present in the graph. Useful to discover what parallel worlds exist.",
    input_schema={"type": "object", "properties": {}},
)
def list_worlds(args):
    rows = _q("""
        MATCH (n)
        WHERE n.world_id IS NOT NULL
        RETURN DISTINCT n.world_id AS world_id
        ORDER BY world_id
    """)
    return rows


def register(registry):
    """Plugin entry point — server.py calls this."""
    # Decorators already registered via the @REGISTRY.tool wrappers above.
    pass