"""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))