more docs

pydantic · samuelcolvin · Mar 14, 2025 · Feb 19, 2025 · Feb 23, 2025 · Feb 24, 2025
commit 294cbd2ee2c910e9569605b38554038df8d31883
diff --git a/docs/api/pydantic_graph/nodes.md b/docs/api/pydantic_graph/nodes.md
@@ -3,6 +3,7 @@
 ::: pydantic_graph.nodes
     options:
         members:
+            - StateT
             - GraphRunContext
             - BaseNode
             - End

diff --git a/docs/api/pydantic_graph/persistence.md b/docs/api/pydantic_graph/persistence.md
@@ -3,3 +3,5 @@
 ::: pydantic_graph.persistence
 
 ::: pydantic_graph.persistence.in_mem
+
+::: pydantic_graph.persistence.file
diff --git a/docs/graph.md b/docs/graph.md
diff --git a/examples/pydantic_ai_examples/question_graph.py b/examples/pydantic_ai_examples/question_graph.py
@@ -38,24 +38,24 @@ class QuestionState:
 
 
 @dataclass
-class GenerateQuestion(BaseNode[QuestionState]):
-    async def run(self, ctx: GraphRunContext[QuestionState]) -> Ask:
+class Ask(BaseNode[QuestionState]):
+    async def run(self, ctx: GraphRunContext[QuestionState]) -> Answer:
         result = await ask_agent.run(
             'Ask a simple question with a single correct answer.',
             message_history=ctx.state.ask_agent_messages,
         )
         ctx.state.ask_agent_messages += result.all_messages()
         ctx.state.question = result.data
-        return Ask(result.data)
+        return Answer(result.data)
 
 
 @dataclass
-class Ask(BaseNode[QuestionState]):
+class Answer(BaseNode[QuestionState]):
     question: str
 
-    async def run(self, ctx: GraphRunContext[QuestionState]) -> Answer:
+    async def run(self, ctx: GraphRunContext[QuestionState]) -> Evaluate:
         answer = input(f'{self.question}: ')
-        return Answer(answer)
+        return Evaluate(answer)
 
 
 class EvaluationResult(BaseModel, use_attribute_docstrings=True):
@@ -73,7 +73,7 @@ class EvaluationResult(BaseModel, use_attribute_docstrings=True):
 
 
 @dataclass
-class Answer(BaseNode[QuestionState, None, str]):
+class Evaluate(BaseNode[QuestionState, None, str]):
     answer: str
 
     async def run(
@@ -96,21 +96,20 @@ async def run(
 class Reprimand(BaseNode[QuestionState]):
     comment: str
 
-    async def run(self, ctx: GraphRunContext[QuestionState]) -> GenerateQuestion:
+    async def run(self, ctx: GraphRunContext[QuestionState]) -> Ask:
         print(f'Comment: {self.comment}')
-        # > Comment: Vichy is no longer the capital of France.
         ctx.state.question = None
-        return GenerateQuestion()
+        return Ask()
 
 
 question_graph = Graph(
-    nodes=(GenerateQuestion, Ask, Answer, Reprimand), state_type=QuestionState
+    nodes=(Ask, Answer, Evaluate, Reprimand), state_type=QuestionState
 )
 
 
 async def run_as_continuous():
     state = QuestionState()
-    node = GenerateQuestion()
+    node = Ask()
     end = await question_graph.run(node, state=state)
     print('END:', end.output)
 
@@ -124,10 +123,10 @@ async def run_as_cli(answer: str | None):
         assert answer is not None, (
             'answer required, usage "uv run -m pydantic_ai_examples.question_graph cli <answer>"'
         )
-        node = Answer(answer)
+        node = Evaluate(answer)
     else:
         state = QuestionState()
-        node = GenerateQuestion()
+        node = Ask()
     # debug(state, node)
 
     with logfire.span('run questions graph'):
@@ -139,7 +138,7 @@ async def run_as_cli(answer: str | None):
                 print('history:', '\n'.join(str(e.node) for e in history), sep='\n')
                 print('Finished!')
                 break
-            elif isinstance(node, Ask):
+            elif isinstance(node, Answer):
                 print(node.question)
                 break
             # otherwise just continue
@@ -165,12 +164,7 @@ async def run_as_cli(answer: str | None):
         sys.exit(1)
 
     if sub_command == 'mermaid':
-        print(question_graph.mermaid_code(start_node=GenerateQuestion))
-        print(
-            question_graph.mermaid_save(
-                'question_graph.jpg', start_node=GenerateQuestion
-            )
-        )
+        print(question_graph.mermaid_code(start_node=Ask))
     elif sub_command == 'continuous':
         asyncio.run(run_as_continuous())
     else:

diff --git a/pydantic_graph/pydantic_graph/graph.py b/pydantic_graph/pydantic_graph/graph.py
@@ -14,8 +14,8 @@
 from typing_inspection import typing_objects
 
 from . import _utils, exceptions, mermaid
-from .nodes import BaseNode, DepsT, End, GraphRunContext, NodeDef, RunEndT
-from .persistence import StatePersistence, StateT, set_nodes_type_context
+from .nodes import BaseNode, DepsT, End, GraphRunContext, NodeDef, RunEndT, StateT
+from .persistence import BaseStatePersistence, set_nodes_type_context
 from .persistence.in_mem import SimpleStatePersistence
 
 # while waiting for https://github.com/pydantic/logfire/issues/745
@@ -125,7 +125,7 @@ async def run(
         *,
         state: StateT = None,
         deps: DepsT = None,
-        persistence: StatePersistence[StateT, T] | None = None,
+        persistence: BaseStatePersistence[StateT, T] | None = None,
         infer_name: bool = True,
         span: LogfireSpan | None = None,
     ) -> GraphRunResult[StateT, T]:
@@ -181,7 +181,7 @@ def run_sync(
         *,
         state: StateT = None,
         deps: DepsT = None,
-        persistence: StatePersistence[StateT, T] | None = None,
+        persistence: BaseStatePersistence[StateT, T] | None = None,
         infer_name: bool = True,
     ) -> GraphRunResult[StateT, T]:
         """Synchronously run the graph.
@@ -215,7 +215,7 @@ async def iter(
         *,
         state: StateT = None,
         deps: DepsT = None,
-        persistence: StatePersistence[StateT, T] | None = None,
+        persistence: BaseStatePersistence[StateT, T] | None = None,
         infer_name: bool = True,
         span: AbstractContextManager[Any] | None = None,
     ) -> AsyncIterator[GraphRun[StateT, DepsT, T]]:
@@ -268,7 +268,7 @@ async def iter(
     async def next(
         self: Graph[StateT, DepsT, T],
         node: BaseNode[StateT, DepsT, T],
-        persistence: StatePersistence[StateT, T],
+        persistence: BaseStatePersistence[StateT, T],
         *,
         state: StateT = None,
         deps: DepsT = None,
@@ -302,17 +302,20 @@ async def next(
 
     async def next_from_persistence(
         self: Graph[StateT, DepsT, T],
-        persistence: StatePersistence[StateT, T],
+        persistence: BaseStatePersistence[StateT, T],
         *,
         deps: DepsT = None,
         infer_name: bool = True,
     ) -> BaseNode[StateT, DepsT, Any] | End[T]:
+        """Run the next node in the graph from a snapshot stored in persistence."""
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
+        self.set_persistence_types(persistence)
 
         snapshot = await persistence.retrieve_next()
         if snapshot is None:
             raise exceptions.GraphRuntimeError('Unable to restore snapshot from state persistence.')
+        assert snapshot.id is not None, 'Snapshot ID should be set'
         snapshot.node.set_snapshot_id(snapshot.id)
 
         run = GraphRun[StateT, DepsT, T](
@@ -325,7 +328,7 @@ async def next_from_persistence(
         )
         return await run.next()
 
-    def set_persistence_types(self, persistence: StatePersistence[StateT, RunEndT]) -> None:
+    def set_persistence_types(self, persistence: BaseStatePersistence[StateT, RunEndT]) -> None:
         with set_nodes_type_context([node_def.node for node_def in self.node_defs.values()]):
             persistence.set_types(lambda: self._inferred_types)
 
@@ -590,7 +593,7 @@ def __init__(
         *,
         graph: Graph[StateT, DepsT, RunEndT],
         start_node: BaseNode[StateT, DepsT, RunEndT],
-        persistence: StatePersistence[StateT, RunEndT],
+        persistence: BaseStatePersistence[StateT, RunEndT],
         state: StateT,
         deps: DepsT,
         snapshot_id: str | None = None,
@@ -738,4 +741,4 @@ class GraphRunResult(Generic[StateT, RunEndT]):
 
     output: RunEndT
     state: StateT
-    persistence: StatePersistence[StateT, RunEndT] = field(repr=False)
+    persistence: BaseStatePersistence[StateT, RunEndT] = field(repr=False)
diff --git a/pydantic_graph/pydantic_graph/nodes.py b/pydantic_graph/pydantic_graph/nodes.py
@@ -4,20 +4,18 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, is_dataclass
 from functools import cache
-from typing import TYPE_CHECKING, Any, ClassVar, Generic, get_type_hints
+from typing import Any, ClassVar, Generic, get_type_hints
 from uuid import uuid4
 
 from typing_extensions import Never, Self, TypeVar, get_origin
 
 from . import _utils, exceptions
 
-if TYPE_CHECKING:
-    from .persistence import StateT
-else:
-    StateT = TypeVar('StateT', default=None)
+__all__ = 'GraphRunContext', 'BaseNode', 'End', 'Edge', 'NodeDef', 'DepsT', 'StateT', 'RunEndT'
 
-__all__ = 'GraphRunContext', 'BaseNode', 'End', 'Edge', 'NodeDef', 'DepsT', 'RunEndT'
 
+StateT = TypeVar('StateT', default=None)
+"""Type variable for the state in a graph."""
 RunEndT = TypeVar('RunEndT', covariant=True, default=None)
 """Covariant type variable for the return type of a graph [`run`][pydantic_graph.graph.Graph.run]."""
 NodeRunEndT = TypeVar('NodeRunEndT', covariant=True, default=Never)

diff --git a/pydantic_graph/pydantic_graph/persistence/__init__.py b/pydantic_graph/pydantic_graph/persistence/__init__.py
@@ -10,29 +10,29 @@
 import pydantic
 from typing_extensions import TypeVar
 
-from ..nodes import BaseNode, End, RunEndT
+from ..nodes import BaseNode, End
 from . import _utils
 
 __all__ = (
     'StateT',
     'NodeSnapshot',
     'EndSnapshot',
     'Snapshot',
-    'StatePersistence',
+    'BaseStatePersistence',
     'set_nodes_type_context',
     'SnapshotStatus',
 )
 
-StateT = TypeVar('StateT', default=None)
-"""Type variable for the state in a graph."""
+StateT = TypeVar('StateT', default=Any)
+RunEndT = TypeVar('RunEndT', covariant=True, default=Any)
 
 UNSET_SNAPSHOT_ID = '__unset__'
 SnapshotStatus = Literal['created', 'pending', 'running', 'success', 'error']
 """The status of a snapshot.
 
 - `'created'`: The snapshot has been created but not yet run.
 - `'pending'`: The snapshot has been retrieved with
-  [`retrieve_next`][pydantic_graph.persistence.StatePersistence.retrieve_next] but not yet run.
+  [`retrieve_next`][pydantic_graph.persistence.BaseStatePersistence.retrieve_next] but not yet run.
 - `'running'`: The snapshot is currently running.
 - `'success'`: The snapshot has been run successfully.
 - `'error'`: The snapshot has been run but an error occurred.
@@ -101,7 +101,7 @@ def node(self) -> End[RunEndT]:
 """
 
 
-class StatePersistence(ABC, Generic[StateT, RunEndT]):
+class BaseStatePersistence(ABC, Generic[StateT, RunEndT]):
     """Abstract base class for storing the state of a graph."""
 
     @abstractmethod
@@ -164,6 +164,9 @@ def record_run(self, snapshot_id: str) -> AbstractAsyncContextManager[None]:
     async def retrieve_next(self) -> NodeSnapshot[StateT, RunEndT] | None:
         """Retrieve a node snapshot with status `'created`' and set its status to `'pending'`.
 
+        This is used by [`Graph.next_from_persistence`][pydantic_graph.graph.Graph.next_from_persistence]
+        to get the next node to run.
+
         Returns: The snapshot, or `None` if no snapshot with status `'created`' exists.
         """
         raise NotImplementedError

diff --git a/pydantic_graph/pydantic_graph/persistence/_utils.py b/pydantic_graph/pydantic_graph/persistence/_utils.py
@@ -23,7 +23,8 @@ def __get_pydantic_core_schema__(
             nodes = nodes_type_context.get()
         except LookupError as e:
             raise RuntimeError(
-                'Unable to build a Pydantic schema for `BaseNode` without setting `nodes_type_context`.'
+                'Unable to build a Pydantic schema for `BaseNode` without setting `nodes_type_context`. '
+                'You should build Pydantic schemas for snapshots using `StatePersistence.set_types()`.'
             ) from e
         if len(nodes) == 1:
             nodes_type = nodes[0]

diff --git a/pydantic_graph/pydantic_graph/persistence/file.py b/pydantic_graph/pydantic_graph/persistence/file.py
@@ -10,29 +10,28 @@
 from typing import Any, Callable
 
 import pydantic
-from typing_extensions import TypeVar
 
 from .. import _utils as _graph_utils, exceptions
 from ..nodes import BaseNode, End
 from . import (
+    BaseStatePersistence,
     EndSnapshot,
     NodeSnapshot,
+    RunEndT,
     Snapshot,
     SnapshotStatus,
-    StatePersistence,
+    StateT,
     _utils,
     build_snapshot_list_type_adapter,
 )
 
-StateT = TypeVar('StateT', default=Any)
-RunEndT = TypeVar('RunEndT', default=Any)
-
 
 @dataclass
-class FileStatePersistence(StatePersistence[StateT, RunEndT]):
-    """State persistence that just hold the latest snapshot."""
+class FileStatePersistence(BaseStatePersistence[StateT, RunEndT]):
+    """File based state persistence that hold a list of snapshots in a JSON file."""
 
     json_file: Path
+    """Path to the JSON file where the snapshots are stored."""
     _snapshots_type_adapter: pydantic.TypeAdapter[list[Snapshot[StateT, RunEndT]]] | None = field(
         default=None, init=False, repr=False
     )

diff --git a/pydantic_graph/pydantic_graph/persistence/in_mem.py b/pydantic_graph/pydantic_graph/persistence/in_mem.py
@@ -13,25 +13,23 @@
 from typing import Any, Callable
 
 import pydantic
-from typing_extensions import TypeVar
 
 from .. import exceptions
 from ..nodes import BaseNode, End
 from . import (
+    BaseStatePersistence,
     EndSnapshot,
     NodeSnapshot,
+    RunEndT,
     Snapshot,
-    StatePersistence,
+    StateT,
     _utils,
     build_snapshot_list_type_adapter,
 )
 
-StateT = TypeVar('StateT', default=Any)
-RunEndT = TypeVar('RunEndT', default=Any)
-
 
 @dataclass
-class SimpleStatePersistence(StatePersistence[StateT, RunEndT]):
+class SimpleStatePersistence(BaseStatePersistence[StateT, RunEndT]):
     """Simple in memory state persistence that just hold the latest snapshot.
 
     If no state persistence implementation is provided when running a graph, this is used by default.
@@ -85,8 +83,8 @@ async def load(self) -> list[Snapshot[StateT, RunEndT]]:
 
 
 @dataclass
-class FullStatePersistence(StatePersistence[StateT, RunEndT]):
-    """In memory state persistence that hold a history of nodes."""
+class FullStatePersistence(BaseStatePersistence[StateT, RunEndT]):
+    """In memory state persistence that hold a list of snapshots."""
 
     deep_copy: bool = True
     """Whether to deep copy the state and nodes when storing them.
Original file line number	Diff line number	Diff line change
Expand Up		@@ -3,3 +3,5 @@
		::: pydantic_graph.persistence

		::: pydantic_graph.persistence.in_mem

		::: pydantic_graph.persistence.file