core/generate/graph_planner.py

"""Graph planner — converts a PropositionGraph into an ArticulationTarget.

The planner walks the graph in topological order and emits an ordered
sequence of articulation steps that the downstream generation pipeline
can execute. Each step carries the proposition node ID, the rhetorical
move, and any constraints inherited from intent classification.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum, unique

from generate.intent import DialogueIntent, IntentTag


@unique
class Relation(Enum):
    ELABORATION = "elaboration"
    CAUSE = "cause"
    CONTRAST = "contrast"
    SEQUENCE = "sequence"
    CORRECTION = "correction"


@unique
class RhetoricalMove(Enum):
    ASSERT = "assert"
    ELABORATE = "elaborate"
    CONTRAST = "contrast"
    SEQUENCE = "sequence"
    CORRECT = "correct"


@dataclass(frozen=True, slots=True)
class GraphEdge:
    source: str
    target: str
    relation: Relation

    def as_dict(self) -> dict[str, str]:
        return {
            "source": self.source,
            "target": self.target,
            "relation": self.relation.value,
        }


@dataclass(frozen=True, slots=True)
class GraphNode:
    node_id: str
    subject: str
    predicate: str
    obj: str
    source_intent: IntentTag

    def as_dict(self) -> dict[str, str]:
        return {
            "node_id": self.node_id,
            "subject": self.subject,
            "predicate": self.predicate,
            "object": self.obj,
            "source_intent": self.source_intent.value,
        }


@dataclass(frozen=True, slots=True)
class PropositionGraph:
    nodes: tuple[GraphNode, ...] = ()
    edges: tuple[GraphEdge, ...] = ()

    def add_node(self, node: GraphNode) -> PropositionGraph:
        return PropositionGraph(nodes=(*self.nodes, node), edges=self.edges)

    def add_edge(self, edge: GraphEdge) -> PropositionGraph:
        return PropositionGraph(nodes=self.nodes, edges=(*self.edges, edge))

    def roots(self) -> tuple[str, ...]:
        targets = frozenset(e.target for e in self.edges)
        return tuple(n.node_id for n in self.nodes if n.node_id not in targets)

    def topo_order(self) -> tuple[str, ...]:
        in_degree: dict[str, int] = {n.node_id: 0 for n in self.nodes}
        for e in self.edges:
            in_degree[e.target] = in_degree.get(e.target, 0) + 1
        queue = sorted(nid for nid, deg in in_degree.items() if deg == 0)
        order: list[str] = []
        while queue:
            nid = queue.pop(0)
            order.append(nid)
            for e in self.edges:
                if e.source == nid:
                    in_degree[e.target] -= 1
                    if in_degree[e.target] == 0:
                        queue.append(e.target)
        return tuple(order)

    def as_dict(self) -> dict[str, object]:
        return {
            "nodes": tuple(n.as_dict() for n in self.nodes),
            "edges": tuple(e.as_dict() for e in self.edges),
        }

    def to_json(self) -> str:
        import json
        return json.dumps(self.as_dict(), sort_keys=True)


@dataclass(frozen=True, slots=True)
class ArticulationStep:
    node_id: str
    move: RhetoricalMove
    predicate: str
    subject: str

    def as_dict(self) -> dict[str, str]:
        return {
            "node_id": self.node_id,
            "move": self.move.value,
            "predicate": self.predicate,
            "subject": self.subject,
        }


@dataclass(frozen=True, slots=True)
class ArticulationTarget:
    steps: tuple[ArticulationStep, ...]
    source_intent: IntentTag

    def as_dict(self) -> dict[str, object]:
        return {
            "steps": tuple(s.as_dict() for s in self.steps),
            "source_intent": self.source_intent.value,
        }


_RELATION_TO_MOVE: dict[Relation, RhetoricalMove] = {
    Relation.ELABORATION: RhetoricalMove.ELABORATE,
    Relation.CAUSE: RhetoricalMove.ELABORATE,
    Relation.CONTRAST: RhetoricalMove.CONTRAST,
    Relation.SEQUENCE: RhetoricalMove.SEQUENCE,
    Relation.CORRECTION: RhetoricalMove.CORRECT,
}


_INTENT_PREDICATES: dict[IntentTag, str] = {
    IntentTag.DEFINITION: "is_defined_as",
    IntentTag.CAUSE: "is_caused_by",
    IntentTag.PROCEDURE: "has_steps",
    IntentTag.COMPARISON: "contrasts_with",
    IntentTag.CORRECTION: "corrects",
    IntentTag.RECALL: "recalls",
    IntentTag.VERIFICATION: "is_verified_as",
}


def graph_from_intent(
    intent: DialogueIntent,
    *,
    prior_node_id: str | None = None,
) -> PropositionGraph:
    """Build a minimal proposition graph from a classified intent."""
    predicate = _INTENT_PREDICATES.get(intent.tag, "addresses")
    graph = PropositionGraph()

    if intent.tag is IntentTag.COMPARISON:
        left = GraphNode(
            node_id="p0",
            subject=intent.subject,
            predicate=predicate,
            obj=intent.secondary_subject or "<pending>",
            source_intent=intent.tag,
        )
        right = GraphNode(
            node_id="p1",
            subject=intent.secondary_subject or "<pending>",
            predicate=predicate,
            obj=intent.subject,
            source_intent=intent.tag,
        )
        edge = GraphEdge(source="p0", target="p1", relation=Relation.CONTRAST)
        return graph.add_node(left).add_node(right).add_edge(edge)

    if intent.tag is IntentTag.CORRECTION:
        root = GraphNode(
            node_id="p0",
            subject=intent.subject,
            predicate=predicate,
            obj=prior_node_id or "<prior>",
            source_intent=intent.tag,
        )
        graph = graph.add_node(root)
        if prior_node_id is not None:
            graph = graph.add_edge(
                GraphEdge(source="p0", target=prior_node_id, relation=Relation.CORRECTION)
            )
        return graph

    root = GraphNode(
        node_id="p0",
        subject=intent.subject,
        predicate=predicate,
        obj="<pending>",
        source_intent=intent.tag,
    )
    return graph.add_node(root)


def ground_graph(
    graph: PropositionGraph,
    recalled_words: tuple[str, ...],
) -> PropositionGraph:
    """Fill <pending> obj slots with recalled words from vault recall.

    Each node whose obj is '<pending>' gets the next available recalled
    word. If there are more nodes than words, remaining slots stay as
    '<pending>'. Comparison nodes get paired words when available.
    """
    words = list(recalled_words)
    new_nodes: list[GraphNode] = []
    for node in graph.nodes:
        if node.obj == "<pending>" and words:
            obj = words.pop(0)
            new_nodes.append(GraphNode(
                node_id=node.node_id,
                subject=node.subject,
                predicate=node.predicate,
                obj=obj,
                source_intent=node.source_intent,
            ))
        else:
            new_nodes.append(node)
    return PropositionGraph(nodes=tuple(new_nodes), edges=graph.edges)


def plan_articulation(graph: PropositionGraph) -> ArticulationTarget:
    """Walk *graph* in topological order and emit an articulation target."""
    node_map = {n.node_id: n for n in graph.nodes}
    incoming: dict[str, Relation | None] = {n.node_id: None for n in graph.nodes}
    for edge in graph.edges:
        if edge.target in incoming:
            incoming[edge.target] = edge.relation

    source_intent = IntentTag.UNKNOWN
    if graph.nodes:
        source_intent = graph.nodes[0].source_intent

    steps: list[ArticulationStep] = []
    for node_id in graph.topo_order():
        node = node_map.get(node_id)
        if node is None:
            continue
        relation = incoming.get(node_id)
        move = _RELATION_TO_MOVE.get(relation, RhetoricalMove.ASSERT) if relation is not None else RhetoricalMove.ASSERT
        steps.append(
            ArticulationStep(
                node_id=node_id,
                move=move,
                predicate=node.predicate,
                subject=node.subject,
            )
        )

    return ArticulationTarget(steps=tuple(steps), source_intent=source_intent)