remove interrupt, replace with "next()"

pydantic · samuelcolvin · Jan 15, 2025 · Dec 22, 2024 · Dec 22, 2024 · Dec 22, 2024
commit b63ca74f5bd8729c8b9c0dabe2558b98233e160d
diff --git a/examples/pydantic_ai_examples/email_extract_graph.py b/examples/pydantic_ai_examples/email_extract_graph.py
@@ -147,7 +147,8 @@ async def run(
 
 async def main():
     state = State(email_content=email)
-    result, history = await graph.run(state, ExtractEvent())
+    history = []
+    result = await graph.run(state, ExtractEvent())
     debug(result, history)
 
 

diff --git a/pydantic_ai_graph/pydantic_ai_graph/__init__.py b/pydantic_ai_graph/pydantic_ai_graph/__init__.py
@@ -1,13 +1,11 @@
-from .graph import Graph, GraphRun
-from .nodes import BaseNode, End, GraphContext, Interrupt
+from .graph import Graph
+from .nodes import BaseNode, End, GraphContext
 from .state import AbstractState, EndEvent, HistoryStep, NextNodeEvent
 
 __all__ = (
     'Graph',
-    'GraphRun',
     'BaseNode',
     'End',
-    'Interrupt',
     'GraphContext',
     'AbstractState',
     'EndEvent',

diff --git a/pydantic_ai_graph/pydantic_ai_graph/graph.py b/pydantic_ai_graph/pydantic_ai_graph/graph.py
@@ -2,20 +2,20 @@
 
 import inspect
 from collections.abc import Sequence
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from pathlib import Path
 from time import perf_counter
-from typing import TYPE_CHECKING, Generic
+from typing import TYPE_CHECKING, Any, Generic
 
 import logfire_api
 from typing_extensions import Never, ParamSpec, TypeVar, Unpack, assert_never
 
 from . import _utils, mermaid
 from ._utils import get_parent_namespace
-from .nodes import BaseNode, End, GraphContext, Interrupt, NodeDef, RunInterrupt
-from .state import EndEvent, HistoryStep, InterruptEvent, NextNodeEvent, StateT
+from .nodes import BaseNode, End, GraphContext, NodeDef
+from .state import EndEvent, HistoryStep, NextNodeEvent, StateT
 
-__all__ = 'Graph', 'GraphRun'
+__all__ = ('Graph',)
 
 _logfire = logfire_api.Logfire(otel_scope='pydantic-ai-graph')
 
@@ -74,16 +74,48 @@ def _validate_edges(self):
                 b = '\n'.join(f' {be}' for be in bad_edges_list)
                 raise ValueError(f'Nodes are referenced in the graph but not included in the graph:\n{b}')
 
+    async def next(
+        self, state: StateT, node: BaseNode[StateT, RunEndT], history: list[HistoryStep[StateT, RunEndT]]
+    ) -> BaseNode[StateT, Any] | End[RunEndT]:
+        node_id = node.get_id()
+        if node_id not in self.node_defs:
+            raise TypeError(f'Node "{node}" is not in the graph.')
+
+        history_step: NextNodeEvent[StateT, RunEndT] | None = NextNodeEvent(state, node)
+        history.append(history_step)
+
+        ctx = GraphContext(state)
+        with _logfire.span('run node {node_id}', node_id=node_id, node=node):
+            start = perf_counter()
+            next_node = await node.run(ctx)
+            history_step.duration = perf_counter() - start
+        return next_node
+
     async def run(
-        self, state: StateT, node: BaseNode[StateT, RunEndT]
-    ) -> tuple[End[RunEndT] | RunInterrupt[StateT], list[HistoryStep[StateT, RunEndT]]]:
-        if not isinstance(node, self.nodes):
-            raise ValueError(f'Node "{node}" is not in the graph.')
-        run = GraphRun[StateT, RunEndT](state=state)
-        # TODO: Infer the graph name properly
-        result = await run.run(self.name or 'graph', node)
-        history = run.history
-        return result, history
+        self,
+        state: StateT,
+        node: BaseNode[StateT, RunEndT],
+    ) -> tuple[End[RunEndT], list[HistoryStep[StateT, RunEndT]]]:
+        history: list[HistoryStep[StateT, RunEndT]] = []
+
+        with _logfire.span(
+            '{graph_name} run {start=}',
+            graph_name=self.name or 'graph',
+            start=node,
+        ) as run_span:
+            while True:
+                next_node = await self.next(state, node, history=history)
+                if isinstance(next_node, End):
+                    history.append(EndEvent(state, next_node))
+                    run_span.set_attribute('history', history)
+                    return next_node, history
+                elif isinstance(next_node, BaseNode):
+                    node = next_node
+                else:
+                    if TYPE_CHECKING:
+                        assert_never(next_node)
+                    else:
+                        raise TypeError(f'Invalid node type: {type(next_node)}. Expected `BaseNode` or `End`.')
 
     def mermaid_code(
         self,
@@ -101,51 +133,3 @@ def mermaid_image(self, **kwargs: Unpack[mermaid.MermaidConfig]) -> bytes:
 
     def mermaid_save(self, path: Path | str, /, **kwargs: Unpack[mermaid.MermaidConfig]) -> None:
         mermaid.save_image(path, self, **kwargs)
-
-
-@dataclass
-class GraphRun(Generic[StateT, RunEndT]):
-    """Stateful run of a graph."""
-
-    state: StateT
-    history: list[HistoryStep[StateT, RunEndT]] = field(default_factory=list)
-
-    async def run(
-        self, graph_name: str, start: BaseNode[StateT, RunEndT], infer_name: bool = True
-    ) -> End[RunEndT] | RunInterrupt[StateT]:
-        current_node = start
-
-        with _logfire.span(
-            '{graph_name} run {start=}',
-            graph_name=graph_name,
-            start=start,
-        ) as run_span:
-            while True:
-                next_node = await self.step(current_node)
-                if isinstance(next_node, (Interrupt, End)):
-                    if isinstance(next_node, End):
-                        self.history.append(EndEvent(self.state, next_node))
-                    else:
-                        self.history.append(InterruptEvent(self.state, next_node))
-                    run_span.set_attribute('history', self.history)
-                    return next_node
-                elif isinstance(next_node, BaseNode):
-                    current_node = next_node
-                else:
-                    if TYPE_CHECKING:
-                        assert_never(next_node)
-                    else:
-                        raise TypeError(f'Invalid node type: {type(next_node)}. Expected `BaseNode` or `End`.')
-
-    async def step(
-        self, node: BaseNode[StateT, RunEndT]
-    ) -> BaseNode[StateT, RunEndT] | End[RunEndT] | RunInterrupt[StateT]:
-        history_step = NextNodeEvent(self.state, node)
-        self.history.append(history_step)
-
-        ctx = GraphContext(self.state)
-        with _logfire.span('run node {node_id}', node_id=node.get_id()):
-            start = perf_counter()
-            next_node = await node.run(ctx)
-            history_step.duration = perf_counter() - start
-        return next_node
diff --git a/pydantic_ai_graph/pydantic_ai_graph/mermaid.py b/pydantic_ai_graph/pydantic_ai_graph/mermaid.py
@@ -2,7 +2,6 @@
 
 import base64
 from collections.abc import Iterable, Sequence
-from itertools import chain
 from pathlib import Path
 from typing import TYPE_CHECKING, Annotated, Any, Literal
 
@@ -43,19 +42,14 @@ def generate_code(
             raise LookupError(f'Start node "{node_id}" is not in the graph.')
 
     node_order = {node_id: index for index, node_id in enumerate(graph.node_defs)}
-    after_interrupt_nodes = set(chain(*(node_def.after_interrupt_nodes for node_def in graph.node_defs.values())))
 
     lines = ['graph TD']
     for node in graph.nodes:
         node_id = node.get_id()
         node_def = graph.node_defs[node_id]
 
-        # we use square brackets (square box) for nodes that can interrupt the flow,
-        # and round brackets (rounded box) for nodes that cannot interrupt the flow
-        if node_id in after_interrupt_nodes:
-            mermaid_name = f'[{node_id}]'
-        else:
-            mermaid_name = f'({node_id})'
+        # we use round brackets (rounded box) for nodes other than the start and end
+        mermaid_name = f'({node_id})'
         if node_id in start_node_ids:
             lines.append(f'  START --> {node_id}{mermaid_name}')
         if node_def.returns_base_node:

diff --git a/pydantic_ai_graph/pydantic_ai_graph/nodes.py b/pydantic_ai_graph/pydantic_ai_graph/nodes.py
@@ -3,14 +3,14 @@
 from abc import abstractmethod
 from dataclasses import dataclass
 from functools import cache
-from typing import Any, Generic, get_args, get_origin, get_type_hints
+from typing import Any, Generic, get_origin, get_type_hints
 
 from typing_extensions import Never, TypeVar
 
 from . import _utils
 from .state import StateT
 
-__all__ = 'GraphContext', 'BaseNode', 'End', 'Interrupt', 'RunInterrupt', 'NodeDef'
+__all__ = 'GraphContext', 'BaseNode', 'End', 'NodeDef'
 
 RunEndT = TypeVar('RunEndT', default=None)
 NodeRunEndT = TypeVar('NodeRunEndT', covariant=True, default=Never)
@@ -27,9 +27,7 @@ class BaseNode(Generic[StateT, NodeRunEndT]):
     """Base class for a node."""
 
     @abstractmethod
-    async def run(
-        self, ctx: GraphContext[StateT]
-    ) -> BaseNode[StateT, Any] | End[NodeRunEndT] | RunInterrupt[StateT]: ...
+    async def run(self, ctx: GraphContext[StateT]) -> BaseNode[StateT, Any] | End[NodeRunEndT]: ...
 
     @classmethod
     @cache
@@ -47,17 +45,10 @@ def get_node_def(cls, local_ns: dict[str, Any] | None) -> NodeDef[StateT, NodeRu
         next_node_ids: set[str] = set()
         returns_end: bool = False
         returns_base_node: bool = False
-        after_interrupt_nodes: list[str] = []
         for return_type in _utils.get_union_args(return_hint):
             return_type_origin = get_origin(return_type) or return_type
             if return_type_origin is End:
                 returns_end = True
-            elif issubclass(return_type_origin, Interrupt):
-                interrupt_args = get_args(return_type)
-                assert len(interrupt_args) == 1, f'Invalid Interrupt return type: {return_type}'
-                next_node_id = interrupt_args[0].get_id()
-                next_node_ids.add(next_node_id)
-                after_interrupt_nodes.append(next_node_id)
             elif return_type_origin is BaseNode:
                 # TODO: Should we disallow this?
                 returns_base_node = True
@@ -71,7 +62,6 @@ def get_node_def(cls, local_ns: dict[str, Any] | None) -> NodeDef[StateT, NodeRu
             cls.get_id(),
             next_node_ids,
             returns_end,
-            after_interrupt_nodes,
             returns_base_node,
         )
 
@@ -83,19 +73,6 @@ class End(Generic[RunEndT]):
     data: RunEndT
 
 
-InterruptNextNodeT = TypeVar('InterruptNextNodeT', covariant=True, bound=BaseNode[Any, Any])
-
-
-@dataclass
-class Interrupt(Generic[InterruptNextNodeT]):
-    """Type to return from a node to signal that the run should be interrupted."""
-
-    node: InterruptNextNodeT
-
-
-RunInterrupt = Interrupt[BaseNode[StateT, Any]]
-
-
 @dataclass
 class NodeDef(Generic[StateT, NodeRunEndT]):
     """Definition of a node.
@@ -112,7 +89,5 @@ class NodeDef(Generic[StateT, NodeRunEndT]):
     """IDs of the nodes that can be called next."""
     returns_end: bool
     """The node definition returns an `End`, hence the node and end the run."""
-    after_interrupt_nodes: list[str]
-    """Nodes that can be returned within an `Interrupt`."""
     returns_base_node: bool
     """The node definition returns a `BaseNode`, hence any node in the next can be called next."""
diff --git a/pydantic_ai_graph/pydantic_ai_graph/state.py b/pydantic_ai_graph/pydantic_ai_graph/state.py
@@ -10,11 +10,11 @@
 
 from . import _utils
 
-__all__ = 'AbstractState', 'StateT', 'NextNodeEvent', 'EndEvent', 'InterruptEvent', 'HistoryStep'
+__all__ = 'AbstractState', 'StateT', 'NextNodeEvent', 'EndEvent', 'HistoryStep'
 
 if TYPE_CHECKING:
     from pydantic_ai_graph import BaseNode
-    from pydantic_ai_graph.nodes import End, RunInterrupt
+    from pydantic_ai_graph.nodes import End
 
 
 class AbstractState(ABC):
@@ -50,28 +50,10 @@ def __post_init__(self):
         # Copy the state to prevent it from being modified by other code
         self.state = _deep_copy_state(self.state)
 
-    def node_summary(self) -> str:
+    def summary(self) -> str:
         return str(self.node)
 
 
-@dataclass
-class InterruptEvent(Generic[StateT]):
-    """History step describing the interruption of a graph run."""
-
-    state: StateT
-    result: RunInterrupt[StateT]
-    ts: datetime = field(default_factory=_utils.now_utc)
-
-    kind: Literal['interrupt'] = 'interrupt'
-
-    def __post_init__(self):
-        # Copy the state to prevent it from being modified by other code
-        self.state = _deep_copy_state(self.state)
-
-    def node_summary(self) -> str:
-        return str(self.result)
-
-
 @dataclass
 class EndEvent(Generic[StateT, RunEndT]):
     """History step describing the end of a graph run."""
@@ -86,7 +68,7 @@ def __post_init__(self):
         # Copy the state to prevent it from being modified by other code
         self.state = _deep_copy_state(self.state)
 
-    def node_summary(self) -> str:
+    def summary(self) -> str:
         return str(self.result)
 
 
@@ -97,4 +79,4 @@ def _deep_copy_state(state: StateT) -> StateT:
         return state.deep_copy()
 
 
-HistoryStep = Union[NextNodeEvent[StateT, RunEndT], InterruptEvent[StateT], EndEvent[StateT, RunEndT]]
+HistoryStep = Union[NextNodeEvent[StateT, RunEndT], EndEvent[StateT, RunEndT]]