Blame for code/graph/graph.py

"""In-memory graph truth and tensor conversion used by GRM inference."""

from __future__ import annotations

import csv
import re
from collections.abc import Iterable, Iterator, Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import torch

from graph.text_features import TextFeaturiser
from graph.types import GraphExample, GraphPath, GraphTensors, PathStep
from graph.utils import stable_hash

#: Modulus for stable global type identities.  Large enough that distinct
#: relation or type names rarely collide, while the encoder reduces further into
#: its configured vocabulary size.
GLOBAL_TYPE_MODULUS = 1_048_576


@dataclass(frozen=True)
class EdgeRecord:
    source_id: str
    relation_id: str
    target_id: str
    source_type: str
    target_type: str
    relation_name: str
    attributes: Mapping[str, str]


@dataclass(frozen=True)
class NodeRecord:
    node_id: str
    node_type: str
    attributes: Mapping[str, str]


class GraphStore:
    """Authoritative graph store with deterministic traversal primitives."""

    def __init__(
        self,
        nodes: Mapping[str, str],
        edges: Iterable[EdgeRecord],
        *,
        graph_version: str = "graph@0",
        node_attributes: Mapping[str, Mapping[str, str]] | None = None,
        relation_attributes: Mapping[str, Mapping[str, str]] | None = None,
    ) -> None:
        self.nodes = dict(nodes)
        if not self.nodes:
            raise ValueError("graph must contain at least one node")
        self.edges = tuple(edges)
        self.graph_version = graph_version
        self.node_attributes = {
            key: dict(value) for key, value in (node_attributes or {}).items()
        }
        self.relation_attributes = {
            key: dict(value) for key, value in (relation_attributes or {}).items()
        }
        self._out: dict[str, tuple[EdgeRecord, ...]] = {}
        for edge in self.edges:
            if edge.source_id not in self.nodes or edge.target_id not in self.nodes:
                raise ValueError(f"edge references unknown node: {edge}")
            self._out.setdefault(edge.source_id, ())
            self._out[edge.source_id] = (*self._out[edge.source_id], edge)

    def outgoing(self, source_id: str) -> tuple[EdgeRecord, ...]:
        return self._out.get(source_id, ())

    def resolve_entities(
        self, mentions: Iterable[str], text: str | None = None
    ) -> tuple[str, ...]:
        """Resolve exact/case-insensitive IDs; never invents an entity."""
        lower = {key.lower(): key for key in self.nodes}
        resolved: list[str] = []
        for mention in mentions:
            if mention in self.nodes:
                resolved.append(mention)
            elif mention.lower() in lower:
                resolved.append(lower[mention.lower()])
        if not resolved and text:
            # Whole-token match only: a substring match resolves "US" inside "bus".
            for key in self.nodes:
                if re.search(rf"(?<!\w){re.escape(key)}(?!\w)", text, re.IGNORECASE):
                    resolved.append(key)
        return tuple(dict.fromkeys(resolved))

    def has_edge(self, source_id: str, relation_id: str, target_id: str) -> bool:
        return any(
            edge.relation_id == relation_id and edge.target_id == target_id
            for edge in self.outgoing(source_id)
        )

    def iter_paths(
        self, start_ids: Iterable[str], max_hops: int
    ) -> Iterator[GraphPath]:
        """Enumerate simple directed paths up to ``max_hops`` in stable order."""
        for start_id in start_ids:
            if start_id not in self.nodes:
                continue
            yield from self._walk(start_id, (), max_hops)

    def _walk(
        self, current_id: str, steps: tuple[PathStep, ...], remaining: int
    ) -> Iterator[GraphPath]:
        if steps:
            yield GraphPath(steps=steps)
        if remaining == 0:
            return
        for edge in self.outgoing(current_id):
            visited = (
                {steps[0].source_id, *(step.target_id for step in steps)}
                if steps
                else {current_id}
            )
            if edge.target_id in visited:
                continue
            yield from self._walk(
                edge.target_id,
                (
                    *steps,
                    PathStep(
                        source_id=edge.source_id,
                        relation_id=edge.relation_id,
                        target_id=edge.target_id,
                        source_type=edge.source_type,
                        target_type=edge.target_type,
                        edge_attributes=edge.attributes,
                    ),
                ),
                remaining - 1,
            )

    def to_example(
        self,
        *,
        feature_dim: int = 8,
        edge_feature_dim: int | None = None,
        featuriser: TextFeaturiser | None = None,
    ) -> GraphExample:
        """Convert to typed tensors.

        ``edge_feature_dim`` defaults to ``feature_dim`` but must be passed
        separately whenever the encoder's ``edge_input_dim`` differs from its
        ``graph_input_dim``.  Emitting both at one width silently worked only
        because every earlier config happened to set the two equal; the shipped
        3B config does not.
        """
        edge_dim = feature_dim if edge_feature_dim is None else edge_feature_dim
        node_ids = tuple(self.nodes)
        node_index = {node_id: i for i, node_id in enumerate(node_ids)}
        node_types = tuple(self.nodes[node_id] for node_id in node_ids)
        node_type_index = {name: i for i, name in enumerate(sorted(set(node_types)))}
        relation_ids = tuple(sorted({edge.relation_id for edge in self.edges}))
        relation_index = {name: i for i, name in enumerate(relation_ids)}
        edge_index = torch.tensor(
            [
                [node_index[e.source_id] for e in self.edges],
                [node_index[e.target_id] for e in self.edges],
            ],
            dtype=torch.long,
        )
        if not self.edges:
            edge_index = torch.empty((2, 0), dtype=torch.long)
        if featuriser is not None:
            # Dense semantic features: a name's vector reflects the words it is made
            # of, so an unseen relation inherits the geometry of its parts instead of
            # landing on an arbitrary unused slot.
            node_features = featuriser.encode_identifiers(
                [f"{node_id} {self.nodes[node_id]}" for node_id in node_ids]
            )
            edge_features = (
                featuriser.encode_identifiers(
                    [
                        f"{edge.relation_id} {edge.relation_name or ''}"
                        for edge in self.edges
                    ]
                )
                if self.edges
                else torch.zeros((0, featuriser.dimension), dtype=torch.float32)
            )
        else:
            node_features = torch.zeros(
                (len(node_ids), feature_dim), dtype=torch.float32
            )
            for i, node_id in enumerate(node_ids):
                attributes = self.node_attributes.get(node_id, {})
                node_features[
                    i,
                    stable_hash(
                        _feature_key(node_id, self.nodes[node_id], attributes),
                        feature_dim,
                    ),
                ] = 1.0
            edge_features = torch.zeros(
                (len(self.edges), edge_dim), dtype=torch.float32
            )
            for i, edge in enumerate(self.edges):
                attributes = {
                    **self.relation_attributes.get(edge.relation_id, {}),
                    **dict(edge.attributes),
                }
                edge_features[
                    i,
                    stable_hash(
                        _feature_key(edge.relation_id, edge.relation_name, attributes),
                        edge_dim,
                    ),
                ] = 1.0
        tensors = GraphTensors(
            node_ids=torch.arange(len(node_ids), dtype=torch.long),
            node_type_ids=torch.tensor(
                [node_type_index[t] for t in node_types], dtype=torch.long
            ),
            node_features=node_features,
            edge_index=edge_index,
            edge_type_ids=torch.tensor(
                [relation_index[e.relation_id] for e in self.edges], dtype=torch.long
            ),
            # Stable identities: a large modulus keeps collisions rare, and the
            # encoder reduces into its own table size.  These, not the local
            # indices, are what the learned embeddings index.
            node_type_hashes=torch.tensor(
                [
                    stable_hash(node_type, GLOBAL_TYPE_MODULUS)
                    for node_type in node_types
                ],
                dtype=torch.long,
            ),
            edge_type_hashes=torch.tensor(
                [stable_hash(e.relation_id, GLOBAL_TYPE_MODULUS) for e in self.edges],
                dtype=torch.long,
            ),
            edge_features=edge_features,
            graph_id=torch.zeros(len(node_ids), dtype=torch.long),
        )
        return GraphExample(
            graph=tensors,
            node_external_ids=node_ids,
            node_types=node_types,
            relation_external_ids=relation_ids,
            relation_names=relation_ids,
            edge_attributes=tuple(e.attributes for e in self.edges),
            node_attributes=tuple(
                self.node_attributes.get(node_id, {}) for node_id in node_ids
            ),
            relation_attributes=tuple(
                self.relation_attributes.get(relation_id, {})
                for relation_id in relation_ids
            ),
            graph_version=self.graph_version,
        )

    @classmethod
    def from_example(cls, example: GraphExample) -> GraphStore:
        # One ``tolist`` per tensor: indexing a tensor element by element costs
        # a dispatch per edge and dominated this conversion on large graphs.
        sources, targets = example.graph.edge_index.tolist()
        relations = example.graph.edge_type_ids.tolist()
        nodes, types = example.node_external_ids, example.node_types
        edges = tuple(
            EdgeRecord(
                source_id=nodes[source],
                relation_id=example.relation_external_ids[relation],
                target_id=nodes[target],
                source_type=types[source],
                target_type=types[target],
                relation_name=example.relation_names[relation],
                attributes=example.edge_attributes[i]
                if example.edge_attributes
                else {},
            )
            for i, (source, target, relation) in enumerate(
                zip(sources, targets, relations, strict=True)
            )
        )
        return cls(
            dict(zip(example.node_external_ids, example.node_types, strict=True)),
            edges,
            graph_version=example.graph_version,
            node_attributes=dict(
                zip(example.node_external_ids, example.node_attributes, strict=True)
            )
            if example.node_attributes
            else None,
            relation_attributes=dict(
                zip(
                    example.relation_external_ids,
                    example.relation_attributes,
                    strict=True,
                )
            )
            if example.relation_attributes
            else None,
        )

    @classmethod
    def from_dict(
        cls, payload: Mapping[str, Any], *, graph_version: str | None = None
    ) -> GraphStore:
        """Adapt the existing graph-service JSON shape without weakening core types."""
        raw_nodes = payload.get("nodes", ())
        nodes: dict[str, str] = {}
        node_attributes: dict[str, Mapping[str, str]] = {}
        if isinstance(raw_nodes, Mapping):
            for node_id, value in raw_nodes.items():
                if isinstance(value, Mapping):
                    nodes[str(node_id)] = str(value.get("type", "entity"))
                    node_attributes[str(node_id)] = {
                        str(key): str(item)
                        for key, item in value.get("attributes", {}).items()
                    }
                else:
                    nodes[str(node_id)] = str(value)
        else:
            for value in raw_nodes:
                node_id = value.get("id") or value.get("node_id")
                if not node_id:
                    raise ValueError("graph JSON node requires id or node_id")
                nodes[str(node_id)] = str(value.get("type", "entity"))
                node_attributes[str(node_id)] = {
                    str(key): str(item)
                    for key, item in value.get("attributes", {}).items()
                }
        relation_names: dict[str, str] = {}
        relation_attributes: dict[str, dict[str, str]] = {}
        raw_relations = payload.get("relations", ())
        if isinstance(raw_relations, Mapping):
            relation_items = (
                {"id": key, **value} if isinstance(value, Mapping) else {"id": key}
                for key, value in raw_relations.items()
            )
        else:
            relation_items = iter(raw_relations)
        for value in relation_items:
            if not isinstance(value, Mapping):
                continue
            relation_id = (
                value.get("id") or value.get("relation_id") or value.get("name")
            )
            if not relation_id:
                continue
            relation_id = str(relation_id)
            relation_names[relation_id] = str(value.get("name", relation_id))
            relation_attributes[relation_id] = {
                str(key): str(item) for key, item in value.get("attributes", {}).items()
            }

        edges: list[EdgeRecord] = []
        for value in payload.get("edges", ()):
            source = value.get("source") or value.get("source_id")
            relation = value.get("relation") or value.get("relation_id")
            target = value.get("target") or value.get("target_id")
            if not source or not relation or not target:
                raise ValueError(
                    "graph JSON edge requires source, relation, and target"
                )
            edges.append(
                EdgeRecord(
                    source_id=str(source),
                    relation_id=str(relation),
                    target_id=str(target),
                    source_type=nodes.get(str(source), "entity"),
                    target_type=nodes.get(str(target), "entity"),
                    relation_name=relation_names.get(
                        str(relation), str(value.get("relation_name", relation))
                    ),
                    attributes={
                        str(key): str(item)
                        for key, item in value.get("attributes", {}).items()
                    },
                )
            )
        return cls(
            nodes,
            edges,
            graph_version=graph_version or str(payload.get("graph_version", "json@0")),
            node_attributes=node_attributes,
            relation_attributes=relation_attributes,
        )

    @classmethod
    def from_csv(
        cls, directory: str | Path, *, graph_version: str = "csv@0"
    ) -> GraphStore:
        root = Path(directory)
        nodes: dict[str, str] = {}
        node_attributes: dict[str, dict[str, str]] = {}
        with (root / "nodes.csv").open(newline="", encoding="utf-8") as handle:
            for row in csv.DictReader(handle):
                node_id = row.get("id") or row.get("node_id") or row.get("name")
                if not node_id:
                    raise ValueError("nodes.csv requires id, node_id, or name")
                nodes[node_id] = row.get("type", "entity")
                node_attributes[node_id] = {
                    key: value
                    for key, value in row.items()
                    if key not in {"id", "node_id", "name", "type"}
                    and value is not None
                }
        relation_names: dict[str, str] = {}
        relation_attributes: dict[str, dict[str, str]] = {}
        relation_path = root / "relations.csv"
        if relation_path.exists():
            with relation_path.open(newline="", encoding="utf-8") as handle:
                for row in csv.DictReader(handle):
                    rid = row.get("id") or row.get("relation_id") or row.get("name")
                    if rid:
                        relation_names[rid] = row.get("name", rid)
                        relation_attributes[rid] = {
                            key: value
                            for key, value in row.items()
                            if key not in {"id", "relation_id", "name"}
                            and value is not None
                        }
        edges: list[EdgeRecord] = []
        with (root / "edges.csv").open(newline="", encoding="utf-8") as handle:
            for row in csv.DictReader(handle):
                source = row.get("source") or row.get("source_id")
                relation = row.get("relation") or row.get("relation_id")
                target = row.get("target") or row.get("target_id")
                if not source or not relation or not target:
                    raise ValueError("edges.csv requires source, relation, and target")
                edges.append(
                    EdgeRecord(
                        source_id=source,
                        relation_id=relation,
                        target_id=target,
                        source_type=nodes.get(source, "entity"),
                        target_type=nodes.get(target, "entity"),
                        relation_name=relation_names.get(relation, relation),
                        attributes={
                            k: v
                            for k, v in row.items()
                            if k
                            not in {
                                "source",
                                "source_id",
                                "relation",
                                "relation_id",
                                "target",
                                "target_id",
                            }
                            and v is not None
                        },
                    )
                )
        return cls(
            nodes,
            edges,
            graph_version=graph_version,
            node_attributes=node_attributes,
            relation_attributes=relation_attributes,
        )


def _feature_key(identifier: str, kind: str, attributes: Mapping[str, str]) -> str:
    serialized = ";".join(f"{key}={attributes[key]}" for key in sorted(attributes))
    return f"{identifier}|{kind}|{serialized}"