Make Graph.iter into an _async_ contextmanager (#958)

Author: dmontagu

docs/agents.md

@@ -114,7 +114,7 @@ agent = Agent('openai:gpt-4o')
 async def main():
     nodes = []
     # Begin an AgentRun, which is an async-iterable over the nodes of the agent's graph
-    with agent.iter('What is the capital of France?') as agent_run:
+    async with agent.iter('What is the capital of France?') as agent_run:
         async for node in agent_run:
             # Each node represents a step in the agent's execution
             nodes.append(node)
@@ -163,7 +163,7 @@ agent = Agent('openai:gpt-4o')
 
 
 async def main():
-    with agent.iter('What is the capital of France?') as agent_run:
+    async with agent.iter('What is the capital of France?') as agent_run:
         node = agent_run.next_node  # (1)!
 
         all_nodes = [node]
@@ -282,7 +282,7 @@ async def main():
     user_prompt = 'What will the weather be like in Paris on Tuesday?'
 
     # Begin a node-by-node, streaming iteration
-    with weather_agent.iter(user_prompt, deps=WeatherService()) as run:
+    async with weather_agent.iter(user_prompt, deps=WeatherService()) as run:
         async for node in run:
             if Agent.is_user_prompt_node(node):
                 # A user prompt node => The user has provided input

docs/graph.md

@@ -679,7 +679,7 @@ count_down_graph = Graph(nodes=[CountDown])
 
 async def main():
     state = CountDownState(counter=3)
-    with count_down_graph.iter(CountDown(), state=state) as run:  # (1)!
+    async with count_down_graph.iter(CountDown(), state=state) as run:  # (1)!
         async for node in run:  # (2)!
             print('Node:', node)
             #> Node: CountDown()
@@ -712,7 +712,7 @@ from count_down import CountDown, CountDownState, count_down_graph
 
 async def main():
     state = CountDownState(counter=5)
-    with count_down_graph.iter(CountDown(), state=state) as run:
+    async with count_down_graph.iter(CountDown(), state=state) as run:
         node = run.next_node  # (1)!
         while not isinstance(node, End):  # (2)!
             print('Node:', node)

pydantic_ai_slim/pydantic_ai/agent.py

@@ -294,7 +294,7 @@ async def main():
         """
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
-        with self.iter(
+        async with self.iter(
             user_prompt=user_prompt,
             result_type=result_type,
             message_history=message_history,
@@ -310,8 +310,8 @@ async def main():
         assert (final_result := agent_run.result) is not None, 'The graph run did not finish properly'
         return final_result
 
-    @contextmanager
-    def iter(
+    @asynccontextmanager
+    async def iter(
         self,
         user_prompt: str | Sequence[_messages.UserContent],
         *,
@@ -323,7 +323,7 @@ def iter(
         usage_limits: _usage.UsageLimits | None = None,
         usage: _usage.Usage | None = None,
         infer_name: bool = True,
-    ) -> Iterator[AgentRun[AgentDepsT, Any]]:
+    ) -> AsyncIterator[AgentRun[AgentDepsT, Any]]:
         """A contextmanager which can be used to iterate over the agent graph's nodes as they are executed.
 
         This method builds an internal agent graph (using system prompts, tools and result schemas) and then returns an
@@ -344,7 +344,7 @@ def iter(
 
         async def main():
             nodes = []
-            with agent.iter('What is the capital of France?') as agent_run:
+            async with agent.iter('What is the capital of France?') as agent_run:
                 async for node in agent_run:
                     nodes.append(node)
             print(nodes)
@@ -454,7 +454,7 @@ async def main():
             system_prompt_dynamic_functions=self._system_prompt_dynamic_functions,
         )
 
-        with graph.iter(
+        async with graph.iter(
             start_node,
             state=state,
             deps=graph_deps,
@@ -633,7 +633,7 @@ async def main():
                 self._infer_name(frame.f_back)
 
         yielded = False
-        with self.iter(
+        async with self.iter(
             user_prompt,
             result_type=result_type,
             message_history=message_history,
@@ -1217,7 +1217,7 @@ def is_end_node(
 class AgentRun(Generic[AgentDepsT, ResultDataT]):
     """A stateful, async-iterable run of an [`Agent`][pydantic_ai.agent.Agent].
 
-    You generally obtain an `AgentRun` instance by calling `with my_agent.iter(...) as agent_run:`.
+    You generally obtain an `AgentRun` instance by calling `async with my_agent.iter(...) as agent_run:`.
 
     Once you have an instance, you can use it to iterate through the run's nodes as they execute. When an
     [`End`][pydantic_graph.nodes.End] is reached, the run finishes and [`result`][pydantic_ai.agent.AgentRun.result]
@@ -1232,7 +1232,7 @@ class AgentRun(Generic[AgentDepsT, ResultDataT]):
     async def main():
         nodes = []
         # Iterate through the run, recording each node along the way:
-        with agent.iter('What is the capital of France?') as agent_run:
+        async with agent.iter('What is the capital of France?') as agent_run:
             async for node in agent_run:
                 nodes.append(node)
         print(nodes)
@@ -1346,7 +1346,7 @@ async def next(
         agent = Agent('openai:gpt-4o')
 
         async def main():
-            with agent.iter('What is the capital of France?') as agent_run:
+            async with agent.iter('What is the capital of France?') as agent_run:
                 next_node = agent_run.next_node  # start with the first node
                 nodes = [next_node]
                 while not isinstance(next_node, End):

pydantic_graph/pydantic_graph/graph.py

@@ -2,8 +2,8 @@
 
 import inspect
 import types
-from collections.abc import AsyncIterator, Iterator, Sequence
-from contextlib import ExitStack, contextmanager
+from collections.abc import AsyncIterator, Sequence
+from contextlib import ExitStack, asynccontextmanager
 from dataclasses import dataclass, field
 from functools import cached_property
 from time import perf_counter
@@ -172,24 +172,24 @@ async def main():
         if infer_name and self.name is None:
             self._infer_name(inspect.currentframe())
 
-        with self.iter(start_node, state=state, deps=deps, infer_name=infer_name, span=span) as graph_run:
+        async with self.iter(start_node, state=state, deps=deps, infer_name=infer_name, span=span) as graph_run:
             async for _node in graph_run:
                 pass
 
         final_result = graph_run.result
         assert final_result is not None, 'GraphRun should have a final result'
         return final_result
 
-    @contextmanager
-    def iter(
+    @asynccontextmanager
+    async def iter(
         self: Graph[StateT, DepsT, T],
         start_node: BaseNode[StateT, DepsT, T],
         *,
         state: StateT = None,
         deps: DepsT = None,
         infer_name: bool = True,
         span: LogfireSpan | None = None,
-    ) -> Iterator[GraphRun[StateT, DepsT, T]]:
+    ) -> AsyncIterator[GraphRun[StateT, DepsT, T]]:
         """A contextmanager which can be used to iterate over the graph's nodes as they are executed.
 
         This method returns a `GraphRun` object which can be used to async-iterate over the nodes of this `Graph` as
@@ -569,7 +569,7 @@ class GraphRun(Generic[StateT, DepsT, RunEndT]):
     """A stateful, async-iterable run of a [`Graph`][pydantic_graph.graph.Graph].
 
     You typically get a `GraphRun` instance from calling
-    `with [my_graph.iter(...)][pydantic_graph.graph.Graph.iter] as graph_run:`. That gives you the ability to iterate
+    `async with [my_graph.iter(...)][pydantic_graph.graph.Graph.iter] as graph_run:`. That gives you the ability to iterate
     through nodes as they run, either by `async for` iteration or by repeatedly calling `.next(...)`.
 
     Here's an example of iterating over the graph from [above][pydantic_graph.graph.Graph]:
@@ -579,7 +579,7 @@ class GraphRun(Generic[StateT, DepsT, RunEndT]):
 
     async def main():
         state = MyState(1)
-        with never_42_graph.iter(Increment(), state=state) as graph_run:
+        async with never_42_graph.iter(Increment(), state=state) as graph_run:
             node_states = [(graph_run.next_node, deepcopy(graph_run.state))]
             async for node in graph_run:
                 node_states.append((node, deepcopy(graph_run.state)))
@@ -593,7 +593,7 @@ async def main():
             '''
 
         state = MyState(41)
-        with never_42_graph.iter(Increment(), state=state) as graph_run:
+        async with never_42_graph.iter(Increment(), state=state) as graph_run:
             node_states = [(graph_run.next_node, deepcopy(graph_run.state))]
             async for node in graph_run:
                 node_states.append((node, deepcopy(graph_run.state)))
@@ -684,7 +684,7 @@ async def next(
 
         async def main():
             state = MyState(48)
-            with never_42_graph.iter(Increment(), state=state) as graph_run:
+            async with never_42_graph.iter(Increment(), state=state) as graph_run:
                 next_node = graph_run.next_node  # start with the first node
                 node_states = [(next_node, deepcopy(graph_run.state))]
 

tests/test_streaming.py

@@ -761,7 +761,7 @@ def result_validator_simple(data: str) -> str:
     messages: list[str] = []
 
     stream_usage: Usage | None = None
-    with agent.iter('Hello') as run:
+    async with agent.iter('Hello') as run:
         async for node in run:
             if agent.is_model_request_node(node):
                 async with node.stream(run.ctx) as stream:
@@ -800,7 +800,7 @@ def result_validator_simple(data: str) -> str:
     run: AgentRun
     stream: AgentStream
     messages: list[ModelResponse] = []
-    with agent.iter('Hello') as run:
+    async with agent.iter('Hello') as run:
         async for node in run:
             if agent.is_model_request_node(node):
                 async with node.stream(run.ctx) as stream:
@@ -843,7 +843,7 @@ def result_validator(data: ResultType | NotResultType) -> ResultType | NotResult
         return ResultType(value=data.value + ' (validated)')
 
     outputs: list[ResultType] = []
-    with agent.iter('test') as run:
+    async with agent.iter('test') as run:
         async for node in run:
             if agent.is_model_request_node(node):
                 async with node.stream(run.ctx) as stream: