src/graph/graph_builder.py

"""Graph builder — populate rdflib or WoGraph from a scenario dict.

A *scenario* describes a family tree, kinship relations, a will clause,
and assets.  This module converts that dict into RDF triples (rdflib) or
WoGraph entities, ready for the reasoning agent to explore.
"""

from __future__ import annotations

from typing import Any

import httpx
from rdflib import Graph, Literal, Namespace, RDF, URIRef, XSD

# Canonical namespace for the kinship ontology
EX = Namespace("http://example.org/kinship#")


# ---------------------------------------------------------------------------
# Canonical test scenario
# ---------------------------------------------------------------------------

def get_arthur_scenario() -> dict:
    """Return the canonical Arthur / Cormac / Diane / Edward scenario.

    Returns:
        A scenario dict consumable by :func:`build_rdflib_graph` and
        :func:`build_wograph`.
    """
    return {
        "persons": [
            {
                "id": "Arthur",
                "name": "Arthur Bellini",
                "age": 82,
                "alive": False,
                "role": "patriarch",
            },
            {
                "id": "Cormac",
                "name": "Cormac Bellini",
                "age": 55,
                "alive": True,
            },
            {
                "id": "Diane",
                "name": "Diane Moretti",
                "age": 52,
                "alive": True,
            },
            {
                "id": "Edward",
                "name": "Edward Bellini",
                "age": 28,
                "alive": True,
            },
        ],
        "relations": [
            {"type": "isBiologicalParentOf", "from": "Arthur", "to": "Cormac"},
            {"type": "isMarriedTo", "from": "Cormac", "to": "Diane"},
            {"type": "isBiologicalParentOf", "from": "Cormac", "to": "Edward"},
            {"type": "isBiologicalParentOf", "from": "Diane", "to": "Edward"},
        ],
        "will": {
            "id": "ArthurWill",
            "clause_text": (
                "The estate passes entirely to my eldest living biological "
                "grandchild. Spouses and in-laws are strictly excluded from "
                "eligibility."
            ),
            "benefactor": "Arthur",
            "excludes": ["isMarriedTo"],
            "requires": ["isBiologicalParentOf"],
        },
        "assets": [
            {
                "id": "Estate",
                "name": "Bellini Estate",
                "value": 5_000_000,
                "owner": "Arthur",
            },
        ],
        "gold_answer": "Edward",
        "optimal_hops": 3,
    }


# ---------------------------------------------------------------------------
# rdflib builder
# ---------------------------------------------------------------------------

def build_rdflib_graph(scenario: dict) -> Graph:
    """Create an in-memory ``rdflib.Graph`` from *scenario*.

    Converts persons, relations, will clauses, and assets into RDF triples
    using the ``EX`` (``http://example.org/kinship#``) namespace.

    Args:
        scenario: A scenario dict (see :func:`get_arthur_scenario` for the
            canonical shape).

    Returns:
        A populated ``rdflib.Graph``.
    """
    g = Graph()
    g.bind("ex", EX)

    # -- persons -------------------------------------------------------------
    for person in scenario.get("persons", []):
        uri = EX[person["id"]]
        g.add((uri, RDF.type, EX.Person))
        g.add((uri, EX.fullName, Literal(person["name"])))
        if "age" in person:
            g.add((uri, EX.age, Literal(person["age"], datatype=XSD.integer)))
        if "alive" in person:
            g.add((uri, EX.isAlive, Literal(person["alive"], datatype=XSD.boolean)))
        if "role" in person:
            g.add((uri, EX.role, Literal(person["role"])))

    # -- relations -----------------------------------------------------------
    for rel in scenario.get("relations", []):
        subj = EX[rel["from"]]
        pred = EX[rel["type"]]
        obj = EX[rel["to"]]
        g.add((subj, pred, obj))

        # Symmetric relations: marriage
        if rel["type"] == "isMarriedTo":
            g.add((obj, pred, subj))

    # -- will / testament ----------------------------------------------------
    will_data = scenario.get("will")
    if will_data:
        will_uri = EX[will_data["id"]]
        g.add((will_uri, RDF.type, EX.Will))
        g.add((will_uri, EX.clauseText, Literal(will_data["clause_text"])))
        g.add((will_uri, EX.benefactor, EX[will_data["benefactor"]]))
        for excl in will_data.get("excludes", []):
            g.add((will_uri, EX.excludesRelation, EX[excl]))
        for req in will_data.get("requires", []):
            g.add((will_uri, EX.requiresRelation, EX[req]))

    # -- assets --------------------------------------------------------------
    for asset in scenario.get("assets", []):
        asset_uri = EX[asset["id"]]
        g.add((asset_uri, RDF.type, EX.Asset))
        g.add((asset_uri, EX.name, Literal(asset["name"])))
        if "value" in asset:
            g.add(
                (asset_uri, EX.value, Literal(asset["value"], datatype=XSD.integer))
            )
        if "owner" in asset:
            g.add((asset_uri, EX.ownedBy, EX[asset["owner"]]))

    return g


# ---------------------------------------------------------------------------
# WoGraph builder
# ---------------------------------------------------------------------------

def _person_to_jsonld(person: dict) -> dict[str, Any]:
    """Convert a person dict to a JSON-LD payload for WoGraph."""
    iri = str(EX[person["id"]])
    doc: dict[str, Any] = {
        "@context": {"ex": str(EX)},
        "@id": iri,
        "@type": "ex:Person",
        "ex:fullName": person["name"],
    }
    if "age" in person:
        doc["ex:age"] = person["age"]
    if "alive" in person:
        doc["ex:isAlive"] = person["alive"]
    if "role" in person:
        doc["ex:role"] = person["role"]
    return doc


def _relation_to_jsonld(rel: dict) -> dict[str, Any]:
    """Convert a relation dict to a JSON-LD patch payload."""
    subj_iri = str(EX[rel["from"]])
    obj_iri = str(EX[rel["to"]])
    pred = f"ex:{rel['type']}"
    return {
        "@context": {"ex": str(EX)},
        "@id": subj_iri,
        pred: {"@id": obj_iri},
    }


def _will_to_jsonld(will_data: dict) -> dict[str, Any]:
    """Convert a will dict to a JSON-LD payload."""
    iri = str(EX[will_data["id"]])
    doc: dict[str, Any] = {
        "@context": {"ex": str(EX)},
        "@id": iri,
        "@type": "ex:Will",
        "ex:clauseText": will_data["clause_text"],
        "ex:benefactor": {"@id": str(EX[will_data["benefactor"]])},
    }
    if will_data.get("excludes"):
        doc["ex:excludesRelation"] = [
            {"@id": str(EX[e])} for e in will_data["excludes"]
        ]
    if will_data.get("requires"):
        doc["ex:requiresRelation"] = [
            {"@id": str(EX[r])} for r in will_data["requires"]
        ]
    return doc


def _asset_to_jsonld(asset: dict) -> dict[str, Any]:
    """Convert an asset dict to a JSON-LD payload."""
    iri = str(EX[asset["id"]])
    doc: dict[str, Any] = {
        "@context": {"ex": str(EX)},
        "@id": iri,
        "@type": "ex:Asset",
        "ex:name": asset["name"],
    }
    if "value" in asset:
        doc["ex:value"] = asset["value"]
    if "owner" in asset:
        doc["ex:ownedBy"] = {"@id": str(EX[asset["owner"]])}
    return doc


def build_wograph(
    scenario: dict,
    endpoint: str,
    api_key: str,
) -> None:
    """Create entities in WoGraph via ``PUT /entities``.

    Uploads persons, relations, will clauses, and assets as JSON-LD
    documents to the WoGraph API.

    Args:
        scenario: A scenario dict.
        endpoint: Base URL of the WoGraph instance.
        api_key: Bearer token for authentication.
    """
    base = endpoint.rstrip("/")
    headers = {
        "Authorization": f"Key {api_key}",
        "Content-Type": "application/json",
    }

    with httpx.Client(base_url=base, headers=headers, timeout=30.0) as client:
        # 1. Create persons
        for person in scenario.get("persons", []):
            doc = _person_to_jsonld(person)
            resp = client.put("/entities", json=doc, params={"id": doc["@id"]})
            resp.raise_for_status()

        # 2. Add relations (PATCH each subject entity)
        for rel in scenario.get("relations", []):
            doc = _relation_to_jsonld(rel)
            resp = client.put("/entities", json=doc, params={"id": doc["@id"]})
            resp.raise_for_status()

            # Symmetric: marriage
            if rel["type"] == "isMarriedTo":
                inverse = {
                    "@context": {"ex": str(EX)},
                    "@id": str(EX[rel["to"]]),
                    f"ex:{rel['type']}": {"@id": str(EX[rel["from"]])},
                }
                resp = client.put(
                    "/entities", json=inverse, params={"id": inverse["@id"]}
                )
                resp.raise_for_status()

        # 3. Will
        will_data = scenario.get("will")
        if will_data:
            doc = _will_to_jsonld(will_data)
            resp = client.put("/entities", json=doc, params={"id": doc["@id"]})
            resp.raise_for_status()

        # 4. Assets
        for asset in scenario.get("assets", []):
            doc = _asset_to_jsonld(asset)
            resp = client.put("/entities", json=doc, params={"id": doc["@id"]})
            resp.raise_for_status()