adding graph tests

pydantic · samuelcolvin · Jan 15, 2025 · Dec 22, 2024 · Dec 22, 2024 · Dec 22, 2024
commit 4f69960044de1aaa9444c3331a0bc1093e57d110
diff --git a/pydantic_graph/pydantic_graph/__init__.py b/pydantic_graph/pydantic_graph/__init__.py
@@ -1,3 +1,4 @@
+from .exceptions import GraphRuntimeError, GraphSetupError
 from .graph import Graph
 from .nodes import BaseNode, Edge, End, GraphContext
 from .state import AbstractState, EndEvent, HistoryStep, NodeEvent
@@ -12,4 +13,6 @@
     'EndEvent',
     'HistoryStep',
     'NodeEvent',
+    'GraphSetupError',
+    'GraphRuntimeError',
 )
diff --git a/pydantic_graph/pydantic_graph/exceptions.py b/pydantic_graph/pydantic_graph/exceptions.py
@@ -0,0 +1,20 @@
+class GraphSetupError(TypeError):
+    """Error caused by an incorrectly configured graph."""
+
+    message: str
+    """Description of the mistake."""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(message)
+
+
+class GraphRuntimeError(RuntimeError):
+    """Error caused by an issue during graph execution."""
+
+    message: str
+    """The error message."""
+
+    def __init__(self, message: str):
+        self.message = message
+        super().__init__(message)
diff --git a/pydantic_graph/pydantic_graph/graph.py b/pydantic_graph/pydantic_graph/graph.py
@@ -12,6 +12,7 @@
 
 from . import _utils, mermaid
 from ._utils import get_parent_namespace
+from .exceptions import GraphRuntimeError, GraphSetupError
 from .nodes import BaseNode, End, GraphContext, NodeDef
 from .state import EndEvent, HistoryStep, NodeEvent, StateT
 
@@ -44,42 +45,42 @@ def __init__(
         _nodes_by_id: dict[str, type[BaseNode[StateT, RunEndT]]] = {}
         for node in nodes:
             node_id = node.get_id()
-            if (existing_node := _nodes_by_id.get(node_id)) and existing_node is not node:
-                raise ValueError(f'Node ID "{node_id}" is not unique — found in {existing_node} and {node}')
+            if existing_node := _nodes_by_id.get(node_id):
+                raise GraphSetupError(f'Node ID `{node_id}` is not unique — found in {existing_node} and {node}')
             else:
                 _nodes_by_id[node_id] = node
         self.nodes = tuple(_nodes_by_id.values())
 
         parent_namespace = get_parent_namespace(inspect.currentframe())
-        self.node_defs: dict[str, NodeDef[StateT, RunEndT]] = {}
-        for node in self.nodes:
-            self.node_defs[node.get_id()] = node.get_node_def(parent_namespace)
+        self.node_defs: dict[str, NodeDef[StateT, RunEndT]] = {
+            node.get_id(): node.get_node_def(parent_namespace) for node in self.nodes
+        }
 
         self._validate_edges()
 
     def _validate_edges(self):
-        known_node_ids = set(self.node_defs.keys())
+        known_node_ids = self.node_defs.keys()
         bad_edges: dict[str, list[str]] = {}
 
         for node_id, node_def in self.node_defs.items():
-            node_bad_edges = node_def.next_node_edges.keys() - known_node_ids
-            for bad_edge in node_bad_edges:
-                bad_edges.setdefault(bad_edge, []).append(f'"{node_id}"')
+            for edge in node_def.next_node_edges.keys():
+                if edge not in known_node_ids:
+                    bad_edges.setdefault(edge, []).append(f'`{node_id}`')
 
         if bad_edges:
-            bad_edges_list = [f'"{k}" is referenced by {_utils.comma_and(v)}' for k, v in bad_edges.items()]
+            bad_edges_list = [f'`{k}` is referenced by {_utils.comma_and(v)}' for k, v in bad_edges.items()]
             if len(bad_edges_list) == 1:
-                raise ValueError(f'{bad_edges_list[0]} but not included in the graph.')
+                raise GraphSetupError(f'{bad_edges_list[0]} but not included in the graph.')
             else:
                 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}')
+                raise GraphSetupError(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.')
+            raise GraphRuntimeError(f'Node `{node}` is not in the graph.')
 
         history_step: NodeEvent[StateT, RunEndT] | None = NodeEvent(state, node)
         history.append(history_step)
@@ -95,7 +96,7 @@ async def run(
         self,
         state: StateT,
         node: BaseNode[StateT, RunEndT],
-    ) -> tuple[End[RunEndT], list[HistoryStep[StateT, RunEndT]]]:
+    ) -> tuple[RunEndT, list[HistoryStep[StateT, RunEndT]]]:
         history: list[HistoryStep[StateT, RunEndT]] = []
 
         with _logfire.span(
@@ -108,24 +109,33 @@ async def run(
                 if isinstance(next_node, End):
                     history.append(EndEvent(state, next_node))
                     run_span.set_attribute('history', history)
-                    return next_node, history
+                    return next_node.data, 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`.')
+                        raise GraphRuntimeError(
+                            f'Invalid node return type: `{type(next_node).__name__}`. Expected `BaseNode` or `End`.'
+                        )
 
     def mermaid_code(
         self,
         *,
         start_node: Sequence[mermaid.NodeIdent] | mermaid.NodeIdent | None = None,
         highlighted_nodes: Sequence[mermaid.NodeIdent] | mermaid.NodeIdent | None = None,
         highlight_css: str = mermaid.DEFAULT_HIGHLIGHT_CSS,
+        edge_labels: bool = True,
+        notes: bool = True,
     ) -> str:
         return mermaid.generate_code(
-            self, start_node=start_node, highlighted_nodes=highlighted_nodes, highlight_css=highlight_css
+            self,
+            start_node=start_node,
+            highlighted_nodes=highlighted_nodes,
+            highlight_css=highlight_css,
+            edge_labels=edge_labels,
+            notes=notes,
         )
 
     def mermaid_image(self, **kwargs: Unpack[mermaid.MermaidConfig]) -> bytes:

diff --git a/pydantic_graph/pydantic_graph/mermaid.py b/pydantic_graph/pydantic_graph/mermaid.py
@@ -59,7 +59,7 @@ def generate_code(
         if node_id in start_node_ids:
             lines.append(f'  [*] --> {node_id}')
         if node_def.returns_base_node:
-            for next_node_id in graph.nodes:
+            for next_node_id in graph.node_defs:
                 lines.append(f'  {node_id} --> {next_node_id}')
         else:
             for next_node_id, edge in node_def.next_node_edges.items():

diff --git a/tests/test_graph.py b/tests/test_graph.py
@@ -2,12 +2,14 @@
 
 from dataclasses import dataclass
 from datetime import timezone
+from functools import cache
 from typing import Union
 
 import pytest
+from dirty_equals import IsStr
 from inline_snapshot import snapshot
 
-from pydantic_graph import BaseNode, End, EndEvent, Graph, GraphContext, NodeEvent
+from pydantic_graph import BaseNode, End, EndEvent, Graph, GraphContext, GraphRuntimeError, GraphSetupError, NodeEvent
 
 from .conftest import IsFloat, IsNow
 
@@ -42,7 +44,7 @@ async def run(self, ctx: GraphContext) -> Union[String2Length, End[int]]:  # noq
     g = Graph[None, int](nodes=(Float2String, String2Length, Double))
     result, history = await g.run(None, Float2String(3.14))
     # len('3.14') * 2 == 8
-    assert result == End(8)
+    assert result == 8
     assert history == snapshot(
         [
             NodeEvent(
@@ -72,7 +74,7 @@ async def run(self, ctx: GraphContext) -> Union[String2Length, End[int]]:  # noq
     )
     result, history = await g.run(None, Float2String(3.14159))
     # len('3.14159') == 7, 21 * 2 == 42
-    assert result == End(42)
+    assert result == 42
     assert history == snapshot(
         [
             NodeEvent(
@@ -112,3 +114,108 @@ async def run(self, ctx: GraphContext) -> Union[String2Length, End[int]]:  # noq
             ),
         ]
     )
+
+
+def test_one_bad_node():
+    class Float2String(BaseNode):
+        async def run(self, ctx: GraphContext) -> String2Length:
+            return String2Length()
+
+    class String2Length(BaseNode[None, None]):
+        async def run(self, ctx: GraphContext) -> End[None]:
+            return End(None)
+
+    with pytest.raises(GraphSetupError) as exc_info:
+        Graph(nodes=(Float2String,))
+
+    assert exc_info.value.message == snapshot(
+        '`String2Length` is referenced by `Float2String` but not included in the graph.'
+    )
+
+
+def test_two_bad_nodes():
+    class Float2String(BaseNode):
+        input_data: float
+
+        async def run(self, ctx: GraphContext) -> String2Length | Double:
+            raise NotImplementedError()
+
+    class String2Length(BaseNode[None, None]):
+        input_data: str
+
+        async def run(self, ctx: GraphContext) -> End[None]:
+            return End(None)
+
+    class Double(BaseNode[None, None]):
+        async def run(self, ctx: GraphContext) -> End[None]:
+            return End(None)
+
+    with pytest.raises(GraphSetupError) as exc_info:
+        Graph(nodes=(Float2String,))
+
+    assert exc_info.value.message == snapshot("""\
+Nodes are referenced in the graph but not included in the graph:
+ `String2Length` is referenced by `Float2String`
+ `Double` is referenced by `Float2String`\
+""")
+
+
+def test_duplicate_id():
+    class Foo(BaseNode):
+        async def run(self, ctx: GraphContext) -> Bar:
+            return Bar()
+
+    class Bar(BaseNode[None, None]):
+        async def run(self, ctx: GraphContext) -> End[None]:
+            return End(None)
+
+        @classmethod
+        @cache
+        def get_id(cls) -> str:
+            return 'Foo'
+
+    with pytest.raises(GraphSetupError) as exc_info:
+        Graph(nodes=(Foo, Bar))
+
+    assert exc_info.value.message == snapshot(IsStr(regex='Node ID `Foo` is not unique — found in.+'))
+
+
+async def test_run_node_not_in_graph():
+    @dataclass
+    class Foo(BaseNode):
+        async def run(self, ctx: GraphContext) -> Bar:
+            return Bar()
+
+    @dataclass
+    class Bar(BaseNode[None, None]):
+        async def run(self, ctx: GraphContext) -> End[None]:
+            return Spam()  # type: ignore
+
+    @dataclass
+    class Spam(BaseNode[None, None]):
+        async def run(self, ctx: GraphContext) -> End[None]:
+            return End(None)
+
+    g = Graph(nodes=(Foo, Bar))
+    with pytest.raises(GraphRuntimeError) as exc_info:
+        await g.run(None, Foo())
+
+    assert exc_info.value.message == snapshot('Node `test_run_node_not_in_graph.<locals>.Spam()` is not in the graph.')
+
+
+async def test_run_return_other():
+    @dataclass
+    class Foo(BaseNode):
+        async def run(self, ctx: GraphContext) -> Bar:
+            return Bar()
+
+    @dataclass
+    class Bar(BaseNode[None, None]):
+        async def run(self, ctx: GraphContext) -> End[None]:
+            return 42  # type: ignore
+
+    g = Graph(nodes=(Foo, Bar))
+    with pytest.raises(GraphRuntimeError) as exc_info:
+        await g.run(None, Foo())
+
+    assert exc_info.value.message == snapshot('Invalid node return type: `int`. Expected `BaseNode` or `End`.')