State persistence (#955)

Co-authored-by: David Montague <[email protected]>

Author: samuelcolvin

Makefile

@@ -49,11 +49,11 @@ test: ## Run tests and collect coverage data
 
 .PHONY: test-all-python
 test-all-python: ## Run tests on Python 3.9 to 3.13
-	UV_PROJECT_ENVIRONMENT=.venv39 uv run --python 3.9 --all-extras coverage run -p -m pytest
-	UV_PROJECT_ENVIRONMENT=.venv310 uv run --python 3.10 --all-extras coverage run -p -m pytest
-	UV_PROJECT_ENVIRONMENT=.venv311 uv run --python 3.11 --all-extras coverage run -p -m pytest
-	UV_PROJECT_ENVIRONMENT=.venv312 uv run --python 3.12 --all-extras coverage run -p -m pytest
-	UV_PROJECT_ENVIRONMENT=.venv313 uv run --python 3.13 --all-extras coverage run -p -m pytest
+	UV_PROJECT_ENVIRONMENT=.venv39 uv run --python 3.9 --all-extras --all-packages coverage run -p -m pytest
+	UV_PROJECT_ENVIRONMENT=.venv310 uv run --python 3.10 --all-extras --all-packages coverage run -p -m pytest
+	UV_PROJECT_ENVIRONMENT=.venv311 uv run --python 3.11 --all-extras --all-packages coverage run -p -m pytest
+	UV_PROJECT_ENVIRONMENT=.venv312 uv run --python 3.12 --all-extras --all-packages coverage run -p -m pytest
+	UV_PROJECT_ENVIRONMENT=.venv313 uv run --python 3.13 --all-extras --all-packages coverage run -p -m pytest
 	@uv run coverage combine
 	@uv run coverage report
 

docs/api/pydantic_graph/nodes.md

@@ -3,6 +3,7 @@
 ::: pydantic_graph.nodes
     options:
         members:
+            - StateT
             - GraphRunContext
             - BaseNode
             - End

docs/api/pydantic_graph/persistence.md

@@ -0,0 +1,7 @@
+# `pydantic_graph.persistence`
+
+::: pydantic_graph.persistence
+
+::: pydantic_graph.persistence.in_mem
+
+::: pydantic_graph.persistence.file

docs/api/pydantic_graph/state.md

@@ -9,11 +9,16 @@
 
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Annotated
 
 import logfire
-from devtools import debug
-from pydantic_graph import BaseNode, Edge, End, Graph, GraphRunContext, HistoryStep
+from groq import BaseModel
+from pydantic_graph import (
+    BaseNode,
+    End,
+    Graph,
+    GraphRunContext,
+)
+from pydantic_graph.persistence.file import FileStatePersistence
 
 from pydantic_ai import Agent
 from pydantic_ai.format_as_xml import format_as_xml
@@ -41,22 +46,23 @@ async def run(self, ctx: GraphRunContext[QuestionState]) -> Answer:
         )
         ctx.state.ask_agent_messages += result.all_messages()
         ctx.state.question = result.data
-        return Answer()
+        return Answer(result.data)
 
 
 @dataclass
 class Answer(BaseNode[QuestionState]):
-    answer: str | None = None
+    question: str
 
     async def run(self, ctx: GraphRunContext[QuestionState]) -> Evaluate:
-        assert self.answer is not None
-        return Evaluate(self.answer)
+        answer = input(f'{self.question}: ')
+        return Evaluate(answer)
 
 
-@dataclass
-class EvaluationResult:
+class EvaluationResult(BaseModel, use_attribute_docstrings=True):
     correct: bool
+    """Whether the answer is correct."""
     comment: str
+    """Comment on the answer, reprimand the user if the answer is wrong."""
 
 
 evaluate_agent = Agent(
@@ -67,101 +73,76 @@ class EvaluationResult:
 
 
 @dataclass
-class Evaluate(BaseNode[QuestionState]):
+class Evaluate(BaseNode[QuestionState, None, str]):
     answer: str
 
     async def run(
         self,
         ctx: GraphRunContext[QuestionState],
-    ) -> Congratulate | Reprimand:
+    ) -> End[str] | Reprimand:
         assert ctx.state.question is not None
         result = await evaluate_agent.run(
             format_as_xml({'question': ctx.state.question, 'answer': self.answer}),
             message_history=ctx.state.evaluate_agent_messages,
         )
         ctx.state.evaluate_agent_messages += result.all_messages()
         if result.data.correct:
-            return Congratulate(result.data.comment)
+            return End(result.data.comment)
         else:
             return Reprimand(result.data.comment)
 
 
-@dataclass
-class Congratulate(BaseNode[QuestionState, None, None]):
-    comment: str
-
-    async def run(
-        self, ctx: GraphRunContext[QuestionState]
-    ) -> Annotated[End, Edge(label='success')]:
-        print(f'Correct answer! {self.comment}')
-        return End(None)
-
-
 @dataclass
 class Reprimand(BaseNode[QuestionState]):
     comment: str
 
     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 Ask()
 
 
 question_graph = Graph(
-    nodes=(Ask, Answer, Evaluate, Congratulate, Reprimand), state_type=QuestionState
+    nodes=(Ask, Answer, Evaluate, Reprimand), state_type=QuestionState
 )
 
 
 async def run_as_continuous():
     state = QuestionState()
     node = Ask()
-    history: list[HistoryStep[QuestionState, None]] = []
-    with logfire.span('run questions graph'):
-        while True:
-            node = await question_graph.next(node, history, state=state)
-            if isinstance(node, End):
-                debug([e.data_snapshot() for e in history])
-                break
-            elif isinstance(node, Answer):
-                assert state.question
-                node.answer = input(f'{state.question} ')
-            # otherwise just continue
+    end = await question_graph.run(node, state=state)
+    print('END:', end.output)
 
 
 async def run_as_cli(answer: str | None):
-    history_file = Path('question_graph_history.json')
-    history = (
-        question_graph.load_history(history_file.read_bytes())
-        if history_file.exists()
-        else []
-    )
-
-    if history:
-        last = history[-1]
-        assert last.kind == 'node', 'expected last step to be a node'
-        state = last.state
-        assert answer is not None, 'answer is required to continue from history'
-        node = Answer(answer)
+    persistence = FileStatePersistence(Path('question_graph.json'))
+    persistence.set_graph_types(question_graph)
+
+    if snapshot := await persistence.load_next():
+        state = snapshot.state
+        assert answer is not None, (
+            'answer required, usage "uv run -m pydantic_ai_examples.question_graph cli <answer>"'
+        )
+        node = Evaluate(answer)
     else:
         state = QuestionState()
         node = Ask()
-    debug(state, node)
+    # debug(state, node)
 
-    with logfire.span('run questions graph'):
+    async with question_graph.iter(node, state=state, persistence=persistence) as run:
         while True:
-            node = await question_graph.next(node, history, state=state)
+            node = await run.next()
             if isinstance(node, End):
-                debug([e.data_snapshot() for e in history])
+                print('END:', node.data)
+                history = await persistence.load_all()
+                print('history:', '\n'.join(str(e.node) for e in history), sep='\n')
                 print('Finished!')
                 break
             elif isinstance(node, Answer):
-                print(state.question)
+                print(node.question)
                 break
             # otherwise just continue
 
-    history_file.write_bytes(question_graph.dump_history(history, indent=2))
-
 
 if __name__ == '__main__':
     import asyncio

docs/graph.md

@@ -66,7 +66,7 @@ nav:
       - api/providers.md
       - api/pydantic_graph/graph.md
       - api/pydantic_graph/nodes.md
-      - api/pydantic_graph/state.md
+      - api/pydantic_graph/persistence.md
       - api/pydantic_graph/mermaid.md
       - api/pydantic_graph/exceptions.md
 

examples/pydantic_ai_examples/question_graph.py

@@ -475,8 +475,8 @@ async def main():
             start_node,
             state=state,
             deps=graph_deps,
-            infer_name=False,
             span=use_span(run_span, end_on_exit=True),
+            infer_name=False,
         ) as graph_run:
             yield AgentRun(graph_run)
 

mkdocs.yml

@@ -53,7 +53,4 @@ fives_graph = Graph(nodes=[DivisibleBy5, Increment])
 result = fives_graph.run_sync(DivisibleBy5(4))
 print(result.output)
 #> 5
-# the full history is quite verbose (see below), so we'll just print the summary
-print([item.data_snapshot() for item in result.history])
-#> [DivisibleBy5(foo=4), Increment(foo=4), DivisibleBy5(foo=5), End(data=5)]
 ```

pydantic_ai_slim/pydantic_ai/agent.py

@@ -1,7 +1,8 @@
 from .exceptions import GraphRuntimeError, GraphSetupError
 from .graph import Graph, GraphRun, GraphRunResult
 from .nodes import BaseNode, Edge, End, GraphRunContext
-from .state import EndStep, HistoryStep, NodeStep
+from .persistence import EndSnapshot, NodeSnapshot, Snapshot
+from .persistence.in_mem import FullStatePersistence, SimpleStatePersistence
 
 __all__ = (
     'Graph',
@@ -11,9 +12,11 @@
     'End',
     'GraphRunContext',
     'Edge',
-    'EndStep',
-    'HistoryStep',
-    'NodeStep',
+    'EndSnapshot',
+    'Snapshot',
+    'NodeSnapshot',
     'GraphSetupError',
     'GraphRuntimeError',
+    'SimpleStatePersistence',
+    'FullStatePersistence',
 )