Blame for code/graph/reasoner.py

"""Query-conditioned iterative expansion with neural transition scoring.

Two changes matter here relative to the earlier implementation.

First, candidate expansion no longer filters transitions by
``required_relation_sequence``.  That field used to carry the gold composition,
so the search was walking a path it had been handed; with the field gone the
search must consider every type-legal transition and rely on the model to rank
them.

Second, ranking is done by the reasoning core and the next-hop head rather than
by an MLP over static embeddings, so the beam is steered by a state that has
attended over the graph.  Hard constraints (PLAN section 17) still prune, because
they are genuine query semantics rather than leaked supervision.
"""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from dataclasses import dataclass

import torch
from torch import Tensor

from graph.composition import compose_relation_sequence
from graph.constraints import path_satisfies_constraints
from graph.graph import EdgeRecord, GraphStore
from graph.heads import NextHopHead
from graph.pat import PathAwareTransformer, PATOutput
from graph.reasoning_core import ReasoningCore
from graph.types import (
    ConstraintState,
    GraphPath,
    GraphQuery,
    PathStep,
    ReasoningState,
)


@dataclass(frozen=True)
class SearchOutput:
    """Completed candidate paths plus the state that produced them."""

    paths: tuple[GraphPath, ...]
    path_states: Tensor
    reasoning_states: tuple[ReasoningState, ...] = ()
    pat_output: PATOutput | None = None
    #: Reasoning-core state aligned with ``paths``; consumed by path selection.
    candidate_states: Tensor | None = None
    pooled_state: Tensor | None = None


class BeamSearchReasoner:
    """Expand graph-local transitions under neural next-hop scoring."""

    def __init__(
        self,
        next_hop_head: NextHopHead,
        beam_width: int,
        candidate_width: int,
        max_hops: int,
        relation_top_k: int = 16,
        max_completed: int = 256,
    ) -> None:
        self.next_hop_head = next_hop_head
        self.beam_width = beam_width
        self.candidate_width = candidate_width
        self.max_hops = max_hops
        self.relation_top_k = relation_top_k
        self.max_completed = max_completed

    def search(
        self,
        graph: GraphStore,
        query: GraphQuery,
        query_state: Tensor,
        node_states: Mapping[str, Tensor],
        relation_states: Mapping[str, Tensor],
        pat: PathAwareTransformer,
        reasoning_core: ReasoningCore,
        memory: Tensor,
        memory_mask: Tensor | None,
        *,
        use_pat: bool = True,
        beam_width: int | None = None,
        query_conditioned: bool = True,
    ) -> SearchOutput:
        constraints = ConstraintState.from_query(query)
        width = beam_width or self.beam_width
        starts = query.start_entity_ids or query.start_entities
        conditioning = (
            query_state if query_conditioned else torch.zeros_like(query_state)
        )

        beam: list[GraphPath] = []
        for start in starts:
            for edge in self._candidate_edges(graph, start, constraints, ()):
                candidate = self._append((), edge)
                if self._partial_allowed(candidate, constraints):
                    beam.append(candidate)

        completed: list[GraphPath] = []
        reasoning_states: list[ReasoningState] = []
        hop = 0
        horizon = min(query.max_hops, self.max_hops)

        while beam and hop < horizon:
            hop += 1
            path_states, pat_output = self._encode_paths(
                beam,
                node_states,
                relation_states,
                conditioning,
                constraints,
                pat,
                use_pat,
            )
            core = reasoning_core(
                conditioning, memory, memory_mask, path_states.unsqueeze(0), hop=hop
            )
            candidate_states = core.candidate_states[0]
            reasoning_states.append(
                self._reasoning_state(
                    beam,
                    conditioning,
                    node_states,
                    relation_states,
                    path_states,
                    hop,
                    constraints,
                )
            )

            for path in beam:
                scored = GraphPath(
                    path.steps, path.score, path.verified, path.query_compatibility
                )
                if (
                    constraints.min_hops <= scored.hop_count
                    and path_satisfies_constraints(scored, constraints, graph)
                ):
                    completed.append(scored)

            if hop >= horizon or len(completed) >= self.max_completed:
                break

            beam = self._expand(
                graph,
                beam,
                candidate_states,
                node_states,
                relation_states,
                constraints,
                width,
            )

        if not completed:
            empty = torch.empty((0, conditioning.shape[-1]), device=conditioning.device)
            return SearchOutput((), empty, tuple(reasoning_states), None, None, None)

        ranked = self._dedupe(completed)[: self.max_completed]
        path_states, pat_output = self._encode_paths(
            ranked,
            node_states,
            relation_states,
            conditioning,
            constraints,
            pat,
            use_pat,
        )
        core = reasoning_core(
            conditioning, memory, memory_mask, path_states.unsqueeze(0), hop=horizon
        )
        return SearchOutput(
            paths=ranked,
            path_states=path_states,
            reasoning_states=tuple(reasoning_states),
            pat_output=pat_output,
            candidate_states=core.candidate_states[0],
            pooled_state=core.state[0],
        )

    # -- expansion ----------------------------------------------------------

    def _expand(
        self,
        graph: GraphStore,
        beam: Sequence[GraphPath],
        candidate_states: Tensor,
        node_states: Mapping[str, Tensor],
        relation_states: Mapping[str, Tensor],
        constraints: ConstraintState,
        width: int,
    ) -> list[GraphPath]:
        """Score every legal extension and keep the best ``width`` of them."""
        extensions: list[GraphPath] = []
        rows: list[Tensor] = []
        sources: list[Tensor] = []
        relations: list[Tensor] = []
        targets: list[Tensor] = []
        zero = torch.zeros_like(candidate_states[0])
        for index, path in enumerate(beam):
            endpoint = path.endpoint_id
            if endpoint is None:
                continue
            for edge in self._candidate_edges(
                graph, endpoint, constraints, path.relation_sequence
            ):
                extended = self._append(path.steps, edge)
                if not self._partial_allowed(extended, constraints):
                    continue
                extensions.append(extended)
                rows.append(candidate_states[index])
                sources.append(node_states.get(edge.source_id, zero))
                relations.append(relation_states.get(edge.relation_id, zero))
                targets.append(node_states.get(edge.target_id, zero))
        if not extensions:
            return []
        scores = self.next_hop_head(
            torch.stack(rows),
            torch.stack(sources),
            torch.stack(relations),
            torch.stack(targets),
        )
        order = torch.argsort(scores, descending=True)[:width]
        return [
            GraphPath(extensions[int(i)].steps, float(scores[int(i)].item()))
            for i in order.tolist()
        ]

    def _encode_paths(
        self,
        paths: Sequence[GraphPath],
        node_states: Mapping[str, Tensor],
        relation_states: Mapping[str, Tensor],
        conditioning: Tensor,
        constraints: ConstraintState,
        pat: PathAwareTransformer,
        use_pat: bool,
    ) -> tuple[Tensor, PATOutput | None]:
        """Encode candidate paths with PAT, falling back to relation composition.

        The no-PAT ablation must still produce a path representation, otherwise
        it would be an ablation of the whole path channel rather than of PAT.
        Relation composition supplies that representation (PLAN section 34).
        """
        composed = torch.stack(
            [
                compose_relation_sequence(
                    path.relation_sequence, relation_states, conditioning
                )
                for path in paths
            ]
        )
        if not use_pat:
            return composed, None
        pat_output = pat.encode_paths(
            paths, node_states, relation_states, conditioning, constraints
        )
        return pat_output.path_states + composed, pat_output

    def _reasoning_state(
        self,
        beam: Sequence[GraphPath],
        conditioning: Tensor,
        node_states: Mapping[str, Tensor],
        relation_states: Mapping[str, Tensor],
        path_states: Tensor,
        hop: int,
        constraints: ConstraintState,
    ) -> ReasoningState:
        base = conditioning[0] if conditioning.ndim == 2 else conditioning
        zero = torch.zeros_like(base)
        return ReasoningState(
            query_embedding=base,
            entity_states=torch.stack(
                [node_states.get(path.endpoint_id or "", zero) for path in beam]
            ),
            relation_states=torch.stack(
                [
                    relation_states.get(path.steps[-1].relation_id, zero)
                    if path.steps
                    else zero
                    for path in beam
                ]
            ),
            candidate_nodes=torch.arange(len(beam), device=base.device),
            candidate_edges=torch.arange(len(beam), device=base.device),
            candidate_paths=tuple(beam),
            path_representations=path_states,
            hop=hop,
            constraint_state=constraints,
            evidence=tuple(beam),
            confidence=path_states.norm(dim=-1),
        )

    @staticmethod
    def _dedupe(paths: Sequence[GraphPath]) -> tuple[GraphPath, ...]:
        unique: dict[tuple[tuple[str, str, str], ...], GraphPath] = {}
        for path in sorted(paths, key=lambda item: item.score, reverse=True):
            key = tuple((s.source_id, s.relation_id, s.target_id) for s in path.steps)
            unique.setdefault(key, path)
        return tuple(unique.values())

    @staticmethod
    def _append(steps: Sequence[PathStep], edge: EdgeRecord) -> GraphPath:
        return GraphPath(
            steps=(
                *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,
                ),
            )
        )

    def _candidate_edges(
        self,
        graph: GraphStore,
        node_id: str,
        constraints: ConstraintState,
        prefix: tuple[str, ...],
    ) -> tuple[EdgeRecord, ...]:
        """Legal outgoing transitions, bounded by ``candidate_width``.

        Only hard query constraints prune here.  ``required_relation_sequence``
        is honoured when a caller genuinely supplies one (the constraints stage
        and the API do), but task generation no longer sets it, so on the main
        task this enumerates every type-legal transition.
        """
        selected: list[EdgeRecord] = []
        relations: set[str] = set()
        required = constraints.required_relation_sequence
        for edge in graph.outgoing(node_id):
            if (
                required
                and len(prefix) < len(required)
                and edge.relation_id != required[len(prefix)]
            ):
                continue
            if edge.relation_id in constraints.forbidden_relations:
                continue
            if (
                constraints.allowed_relations
                and edge.relation_id not in constraints.allowed_relations
            ):
                continue
            if (
                edge.relation_id not in relations
                and len(relations) >= self.relation_top_k
            ):
                continue
            relations.add(edge.relation_id)
            selected.append(edge)
            if len(selected) >= self.candidate_width:
                break
        return tuple(selected)

    @staticmethod
    def _partial_allowed(path: GraphPath, constraints: ConstraintState) -> bool:
        if not path.steps:
            return True
        if constraints.forbidden_relations.intersection(path.relation_sequence):
            return False
        if constraints.allowed_relations and not set(path.relation_sequence).issubset(
            constraints.allowed_relations
        ):
            return False
        required = constraints.required_relation_sequence
        if (
            required
            and tuple(path.relation_sequence) != required[: len(path.relation_sequence)]
        ):
            return False
        if constraints.forbidden_types.intersection(
            {step.target_type for step in path.steps if step.target_type}
        ):
            return False
        ids = [path.steps[0].source_id, *(step.target_id for step in path.steps)]
        # Simple paths only: a repeated node is a cycle, never new evidence.
        return len(ids) == len(set(ids))