Blame for code/graph/types.py

"""Typed contracts shared by the GRM model, search, verifier, and API."""

from __future__ import annotations

from collections.abc import Mapping
from dataclasses import dataclass, field
from enum import StrEnum
from typing import Any

import torch
from torch import Tensor


class ReasoningStatus(StrEnum):
    SUPPORTED = "supported"
    UNSUPPORTED = "unsupported"
    CONTRADICTED = "contradicted"
    INSUFFICIENT_EVIDENCE = "insufficient_evidence"
    ERROR = "error"


@dataclass(frozen=True)
class GraphTensors:
    """A validated, batch-agnostic graph tensor contract.

    ``edge_index`` is shaped ``[2, E]`` and stores source/target node indices.
    IDs are indices into the graph tensors, never user-facing entity IDs.

    Two kinds of type id are carried, and the distinction is load-bearing:

    ``node_type_ids`` / ``edge_type_ids``
        *Local* indices -- positions within this graph's sorted type and relation
        lists.  They exist for bookkeeping: pooling per-edge states into a
        per-graph relation bank, and mapping a bank row back to a relation name.
    ``node_type_hashes`` / ``edge_type_hashes``
        *Global, graph-independent* identities, used for the learned embeddings.

    Using the local index for the embedding lookup was a significant bug.  Local
    indices come from the alphabetical rank of a name *within one graph*, so a
    graph that simply lacked one relation shifted every later index: ``owns``
    could be embedding row 1 in one graph and row 0 in the next.  The relation and
    type embedding tables could therefore not learn a stable representation of any
    relation, and renaming labels reshuffled everything, which is what collapsed
    the unseen-ontology holdout.
    """

    node_ids: Tensor
    node_type_ids: Tensor
    node_features: Tensor
    edge_index: Tensor
    edge_type_ids: Tensor
    edge_features: Tensor
    graph_id: Tensor
    #: Stable global type identities for the learned embeddings.  Default to the
    #: local ids so older callers keep working, with the drift they always had.
    node_type_hashes: Tensor | None = None
    edge_type_hashes: Tensor | None = None

    def __post_init__(self) -> None:
        _require_rank(self.node_ids, 1, "node_ids")
        _require_rank(self.node_type_ids, 1, "node_type_ids")
        _require_rank(self.node_features, 2, "node_features")
        _require_rank(self.edge_index, 2, "edge_index")
        _require_rank(self.edge_type_ids, 1, "edge_type_ids")
        _require_rank(self.edge_features, 2, "edge_features")
        _require_rank(self.graph_id, 1, "graph_id")
        if self.edge_index.shape[0] != 2:
            raise ValueError(
                f"edge_index must have shape [2, E], got {tuple(self.edge_index.shape)}"
            )
        n = self.node_ids.shape[0]
        e = self.edge_index.shape[1]
        if self.node_type_ids.shape[0] != n or self.node_features.shape[0] != n:
            raise ValueError(
                "node_ids, node_type_ids, and node_features must have the same first dimension"
            )
        if self.edge_type_ids.shape[0] != e or self.edge_features.shape[0] != e:
            raise ValueError(
                "edge_index, edge_type_ids, and edge_features must have the same edge dimension"
            )
        if self.graph_id.shape[0] != n:
            raise ValueError("graph_id must contain one graph ID per node")
        if self.edge_index.numel() and (
            self.edge_index.min() < 0 or self.edge_index.max() >= n
        ):
            raise ValueError("edge_index contains a node index outside node_ids")
        if (
            self.node_ids.dtype not in (torch.int32, torch.int64)
            or self.node_type_ids.dtype not in (torch.int32, torch.int64)
            or self.edge_type_ids.dtype not in (torch.int32, torch.int64)
            or self.graph_id.dtype not in (torch.int32, torch.int64)
        ):
            raise TypeError(
                "node, type, edge-type, and graph IDs must be integer tensors"
            )
        if self.edge_index.dtype not in (torch.int32, torch.int64):
            raise TypeError("edge_index must be an integer tensor")
        if self.node_type_hashes is not None:
            _require_rank(self.node_type_hashes, 1, "node_type_hashes")
            if self.node_type_hashes.shape[0] != n:
                raise ValueError("node_type_hashes must have one entry per node")
        if self.edge_type_hashes is not None:
            _require_rank(self.edge_type_hashes, 1, "edge_type_hashes")
            if self.edge_type_hashes.shape[0] != e:
                raise ValueError("edge_type_hashes must have one entry per edge")
        if (
            not self.node_features.is_floating_point()
            or not self.edge_features.is_floating_point()
        ):
            raise TypeError(
                "node_features and edge_features must be floating-point tensors"
            )

    @property
    def num_nodes(self) -> int:
        return int(self.node_ids.shape[0])

    @property
    def num_edges(self) -> int:
        return int(self.edge_index.shape[1])

    @property
    def node_type_identity(self) -> Tensor:
        """Global node-type ids when present, else the local indices."""
        return (
            self.node_type_ids
            if self.node_type_hashes is None
            else self.node_type_hashes
        )

    @property
    def edge_type_identity(self) -> Tensor:
        """Global relation ids when present, else the local indices."""
        return (
            self.edge_type_ids
            if self.edge_type_hashes is None
            else self.edge_type_hashes
        )

    def to(self, device: torch.device | str) -> GraphTensors:
        return GraphTensors(
            node_ids=self.node_ids.to(device),
            node_type_ids=self.node_type_ids.to(device),
            node_features=self.node_features.to(device),
            edge_index=self.edge_index.to(device),
            edge_type_ids=self.edge_type_ids.to(device),
            edge_features=self.edge_features.to(device),
            graph_id=self.graph_id.to(device),
            node_type_hashes=None
            if self.node_type_hashes is None
            else self.node_type_hashes.to(device),
            edge_type_hashes=None
            if self.edge_type_hashes is None
            else self.edge_type_hashes.to(device),
        )


@dataclass(frozen=True)
class GraphQuery:
    text: str | None = None
    start_entities: tuple[str, ...] = ()
    start_entity_ids: tuple[str, ...] = ()
    allowed_relations: tuple[str, ...] = ()
    required_relations: tuple[str, ...] = ()
    required_relation_sequence: tuple[str, ...] = ()
    forbidden_relations: tuple[str, ...] = ()
    required_types: tuple[str, ...] = ()
    forbidden_types: tuple[str, ...] = ()
    required_intermediate_types: tuple[str, ...] = ()
    max_hops: int = 3
    min_hops: int = 1
    exact_hops: int | None = None
    return_paths: bool = True
    verify: bool = True
    operation: str = "find_entities"
    answer_mode: str = "entity"
    target_entity_ids: tuple[str, ...] = ()

    def __post_init__(self) -> None:
        if self.max_hops < 0 or self.min_hops < 0:
            raise ValueError("hop bounds must be non-negative")
        if self.min_hops > self.max_hops:
            raise ValueError("min_hops cannot exceed max_hops")
        if self.exact_hops is not None and not (
            self.min_hops <= self.exact_hops <= self.max_hops
        ):
            raise ValueError("exact_hops must be within min_hops and max_hops")
        if self.answer_mode not in {
            "entity",
            "relation",
            "boolean",
            "count",
            "aggregate",
        }:
            raise ValueError(f"unsupported answer_mode: {self.answer_mode!r}")

    def canonical_text(self) -> str:
        return "|".join(
            (
                self.operation,
                self.text or "",
                ",".join(self.start_entity_ids or self.start_entities),
                ",".join(self.allowed_relations),
                ",".join(self.required_relations),
                ",".join(self.required_relation_sequence),
                ",".join(self.forbidden_relations),
                ",".join(self.required_types),
                ",".join(self.forbidden_types),
                ",".join(self.required_intermediate_types),
                str(self.min_hops),
                str(self.max_hops),
                str(self.exact_hops),
                self.answer_mode,
                ",".join(self.target_entity_ids),
            )
        )


@dataclass(frozen=True)
class PathStep:
    source_id: str
    relation_id: str
    target_id: str
    source_type: str | None = None
    target_type: str | None = None
    edge_attributes: Mapping[str, str] = field(default_factory=dict)


@dataclass(frozen=True)
class GraphPath:
    steps: tuple[PathStep, ...]
    score: float = 0.0
    verified: bool = False
    query_compatibility: float | None = None

    @property
    def hop_count(self) -> int:
        return len(self.steps)

    @property
    def start_id(self) -> str | None:
        return self.steps[0].source_id if self.steps else None

    @property
    def endpoint_id(self) -> str | None:
        return self.steps[-1].target_id if self.steps else None

    @property
    def relation_sequence(self) -> tuple[str, ...]:
        return tuple(step.relation_id for step in self.steps)


@dataclass(frozen=True)
class ConstraintState:
    allowed_relations: frozenset[str] = frozenset()
    required_relations: frozenset[str] = frozenset()
    required_relation_sequence: tuple[str, ...] = ()
    forbidden_relations: frozenset[str] = frozenset()
    required_types: frozenset[str] = frozenset()
    forbidden_types: frozenset[str] = frozenset()
    required_intermediate_types: frozenset[str] = frozenset()
    min_hops: int = 0
    max_hops: int = 8
    exact_hops: int | None = None
    require_unique_nodes: bool = False

    @classmethod
    def from_query(cls, query: GraphQuery) -> ConstraintState:
        return cls(
            allowed_relations=frozenset(query.allowed_relations),
            required_relations=frozenset(query.required_relations),
            required_relation_sequence=query.required_relation_sequence,
            forbidden_relations=frozenset(query.forbidden_relations),
            required_types=frozenset(query.required_types),
            forbidden_types=frozenset(query.forbidden_types),
            required_intermediate_types=frozenset(query.required_intermediate_types),
            min_hops=query.min_hops,
            max_hops=query.max_hops,
            exact_hops=query.exact_hops,
        )


@dataclass(frozen=True)
class ReasoningState:
    query_embedding: Tensor
    entity_states: Tensor
    relation_states: Tensor
    candidate_nodes: Tensor
    candidate_edges: Tensor
    candidate_paths: tuple[GraphPath, ...]
    path_representations: Tensor
    hop: int
    constraint_state: ConstraintState
    evidence: tuple[GraphPath, ...] = ()
    confidence: Tensor | None = None


@dataclass(frozen=True)
class EntityAnswer:
    entity_id: str
    score: float


@dataclass(frozen=True)
class ReasoningResult:
    status: ReasoningStatus
    answers: tuple[EntityAnswer, ...]
    supporting_paths: tuple[GraphPath, ...]
    conflicting_paths: tuple[GraphPath, ...]
    confidence: float
    calibrated: bool
    model_version: str
    graph_version: str
    reasoning_steps: int
    metadata: Mapping[str, Any] = field(default_factory=dict)
    relation_answers: tuple[str, ...] = ()
    boolean_answer: bool | None = None
    aggregate_value: int | float | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "status": self.status.value,
            "answers": [
                {"entity_id": a.entity_id, "score": a.score} for a in self.answers
            ],
            "supporting_paths": [_path_to_dict(p) for p in self.supporting_paths],
            "conflicting_paths": [_path_to_dict(p) for p in self.conflicting_paths],
            "confidence": self.confidence,
            "calibrated": self.calibrated,
            "model_version": self.model_version,
            "graph_version": self.graph_version,
            "reasoning_steps": self.reasoning_steps,
            "metadata": dict(self.metadata),
            "relation_answers": list(self.relation_answers),
            "boolean_answer": self.boolean_answer,
            "aggregate_value": self.aggregate_value,
        }


@dataclass(frozen=True)
class GraphExample:
    graph: GraphTensors
    node_external_ids: tuple[str, ...]
    node_types: tuple[str, ...]
    relation_external_ids: tuple[str, ...]
    relation_names: tuple[str, ...]
    edge_attributes: tuple[Mapping[str, str], ...] = ()
    node_attributes: tuple[Mapping[str, str], ...] = ()
    relation_attributes: tuple[Mapping[str, str], ...] = ()
    graph_version: str = "synthetic@0"

    def __post_init__(self) -> None:
        if len(self.node_external_ids) != self.graph.num_nodes:
            raise ValueError("node_external_ids must align with graph nodes")
        if len(self.node_types) != self.graph.num_nodes:
            raise ValueError("node_types must align with graph nodes")
        if len(self.edge_attributes) not in (0, self.graph.num_edges):
            raise ValueError("edge_attributes must be empty or align with graph edges")
        if len(self.node_attributes) not in (0, self.graph.num_nodes):
            raise ValueError("node_attributes must be empty or align with graph nodes")
        if len(self.relation_external_ids) != len(self.relation_names):
            raise ValueError("relation IDs and names must align")
        if len(self.relation_attributes) not in (0, len(self.relation_external_ids)):
            raise ValueError(
                "relation_attributes must be empty or align with relations"
            )
        if len(set(self.node_external_ids)) != len(self.node_external_ids):
            raise ValueError("node_external_ids must be unique")

    @property
    def node_index(self) -> Mapping[str, int]:
        return {value: i for i, value in enumerate(self.node_external_ids)}


def _require_rank(value: Tensor, rank: int, name: str) -> None:
    if value.ndim != rank:
        raise ValueError(
            f"{name} must have rank {rank}, got shape {tuple(value.shape)}"
        )


def _path_to_dict(path: GraphPath) -> dict[str, Any]:
    return {
        "steps": [
            {
                "source_id": step.source_id,
                "relation_id": step.relation_id,
                "target_id": step.target_id,
                "source_type": step.source_type,
                "target_type": step.target_type,
                "edge_attributes": dict(step.edge_attributes),
            }
            for step in path.steps
        ],
        "score": path.score,
        "verified": path.verified,
        "query_compatibility": path.query_compatibility,
    }