pydantic-graph 1.11.0__tar.gz → 1.11.1__tar.gz

{pydantic_graph-1.11.0 → pydantic_graph-1.11.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-graph
-Version: 1.11.0
+Version: 1.11.1
 Summary: Graph and state machine library
 Project-URL: Homepage, https://ai.pydantic.dev/graph/tree/main/pydantic_graph
 Project-URL: Source, https://github.com/pydantic/pydantic-ai

{pydantic_graph-1.11.0 → pydantic_graph-1.11.1}/pydantic_graph/beta/graph.py

@@ -10,7 +10,7 @@ from __future__ import annotations as _annotations
 import sys
 import uuid
 from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Iterable, Sequence
-from contextlib import AbstractContextManager, ExitStack, asynccontextmanager, contextmanager
+from contextlib import AbstractContextManager, AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Generic, Literal, TypeGuard, cast, get_args, get_origin, overload
 
@@ -381,18 +381,32 @@ class GraphRun(Generic[StateT, DepsT, OutputT]):
         run_id = GraphRunID(str(uuid.uuid4()))
         initial_fork_stack: ForkStack = (ForkStackItem(StartNode.id, NodeRunID(run_id), 0),)
         self._first_task = GraphTask(node_id=StartNode.id, inputs=inputs, fork_stack=initial_fork_stack)
-        self._iterator_instance = _GraphIterator[StateT, DepsT, OutputT](self.graph, self.state, self.deps)
+        self._iterator_task_group = create_task_group()
+        self._iterator_instance = _GraphIterator[StateT, DepsT, OutputT](
+            self.graph, self.state, self.deps, self._iterator_task_group
+        )
         self._iterator = self._iterator_instance.iter_graph(self._first_task)
 
         self.__traceparent = traceparent
+        self._async_exit_stack = AsyncExitStack()
 
     async def __aenter__(self):
+        self._async_exit_stack.enter_context(_unwrap_exception_groups())
+        await self._async_exit_stack.enter_async_context(self._iterator_task_group)
+        await self._async_exit_stack.enter_async_context(self._iterator_context())
         return self
 
     async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any):
-        self._iterator_instance.iter_stream_sender.close()
-        self._iterator_instance.iter_stream_receiver.close()
-        await self._iterator.aclose()
+        await self._async_exit_stack.__aexit__(exc_type, exc_val, exc_tb)
+
+    @asynccontextmanager
+    async def _iterator_context(self):
+        try:
+            yield
+        finally:
+            self._iterator_instance.iter_stream_sender.close()
+            self._iterator_instance.iter_stream_receiver.close()
+            await self._iterator.aclose()
 
     @overload
     def _traceparent(self, *, required: Literal[False]) -> str | None: ...
@@ -495,13 +509,13 @@ class _GraphIterator(Generic[StateT, DepsT, OutputT]):
     graph: Graph[StateT, DepsT, Any, OutputT]
     state: StateT
     deps: DepsT
+    task_group: TaskGroup
 
     cancel_scopes: dict[TaskID, CancelScope] = field(init=False)
     active_tasks: dict[TaskID, GraphTask] = field(init=False)
     active_reducers: dict[tuple[JoinID, NodeRunID], JoinState] = field(init=False)
     iter_stream_sender: MemoryObjectSendStream[_GraphTaskResult] = field(init=False)
     iter_stream_receiver: MemoryObjectReceiveStream[_GraphTaskResult] = field(init=False)
-    _task_group: TaskGroup | None = field(init=False)
 
     def __post_init__(self):
         self.cancel_scopes = {}
@@ -509,178 +523,176 @@ class _GraphIterator(Generic[StateT, DepsT, OutputT]):
         self.active_reducers = {}
         self.iter_stream_sender, self.iter_stream_receiver = create_memory_object_stream[_GraphTaskResult]()
 
-    @property
-    def task_group(self) -> TaskGroup:
-        if self._task_group is None:
-            raise RuntimeError("This graph iterator hasn't been started")  # pragma: no cover
-        return self._task_group
-
     async def iter_graph(  # noqa C901
         self, first_task: GraphTask
     ) -> AsyncGenerator[EndMarker[OutputT] | Sequence[GraphTask], EndMarker[OutputT] | Sequence[GraphTask]]:
-        with _unwrap_exception_groups():
-            async with self.iter_stream_sender, create_task_group() as self._task_group:
-                try:
-                    # Fire off the first task
-                    self.active_tasks[first_task.task_id] = first_task
-                    self._handle_execution_request([first_task])
-
-                    # Handle task results
-                    async with self.iter_stream_receiver:
-                        while self.active_tasks or self.active_reducers:
-                            async for task_result in self.iter_stream_receiver:  # pragma: no branch
-                                if isinstance(task_result.result, JoinItem):
-                                    maybe_overridden_result = task_result.result
-                                else:
-                                    maybe_overridden_result = yield task_result.result
-                                if isinstance(maybe_overridden_result, EndMarker):
+        async with self.iter_stream_sender:
+            try:
+                # Fire off the first task
+                self.active_tasks[first_task.task_id] = first_task
+                self._handle_execution_request([first_task])
+
+                # Handle task results
+                async with self.iter_stream_receiver:
+                    while self.active_tasks or self.active_reducers:
+                        async for task_result in self.iter_stream_receiver:  # pragma: no branch
+                            if isinstance(task_result.result, JoinItem):
+                                maybe_overridden_result = task_result.result
+                            else:
+                                maybe_overridden_result = yield task_result.result
+                            if isinstance(maybe_overridden_result, EndMarker):
+                                # If we got an end marker, this task is definitely done, and we're ready to
+                                # start cleaning everything up
+                                await self._finish_task(task_result.source.task_id)
+                                if self.active_tasks:
+                                    # Cancel the remaining tasks
                                     self.task_group.cancel_scope.cancel()
-                                    return
-                                elif isinstance(maybe_overridden_result, JoinItem):
-                                    result = maybe_overridden_result
-                                    parent_fork_id = self.graph.get_parent_fork(result.join_id).fork_id
-                                    for i, x in enumerate(result.fork_stack[::-1]):
-                                        if x.fork_id == parent_fork_id:
-                                            # For non-final joins (those that are intermediate nodes of other joins),
-                                            # preserve the fork stack so downstream joins can still associate with the same fork run
-                                            if self.graph.is_final_join(result.join_id):
-                                                # Final join: remove the parent fork from the stack
-                                                downstream_fork_stack = result.fork_stack[: len(result.fork_stack) - i]
-                                            else:
-                                                # Non-final join: preserve the fork stack
-                                                downstream_fork_stack = result.fork_stack
-                                            fork_run_id = x.node_run_id
-                                            break
-                                    else:  # pragma: no cover
-                                        raise RuntimeError('Parent fork run not found')
-
-                                    join_node = self.graph.nodes[result.join_id]
-                                    assert isinstance(join_node, Join), f'Expected a `Join` but got {join_node}'
-                                    join_state = self.active_reducers.get((result.join_id, fork_run_id))
-                                    if join_state is None:
-                                        current = join_node.initial_factory()
-                                        join_state = self.active_reducers[(result.join_id, fork_run_id)] = JoinState(
-                                            current, downstream_fork_stack
-                                        )
-                                    context = ReducerContext(state=self.state, deps=self.deps, join_state=join_state)
-                                    join_state.current = join_node.reduce(context, join_state.current, result.inputs)
-                                    if join_state.cancelled_sibling_tasks:
-                                        await self._cancel_sibling_tasks(parent_fork_id, fork_run_id)
-                                else:
-                                    for new_task in maybe_overridden_result:
-                                        self.active_tasks[new_task.task_id] = new_task
-
-                                tasks_by_id_values = list(self.active_tasks.values())
-                                join_tasks: list[GraphTask] = []
-
-                                for join_id, fork_run_id in self._get_completed_fork_runs(
-                                    task_result.source, tasks_by_id_values
-                                ):
-                                    join_state = self.active_reducers.pop((join_id, fork_run_id))
-                                    join_node = self.graph.nodes[join_id]
-                                    assert isinstance(join_node, Join), f'Expected a `Join` but got {join_node}'
-                                    new_tasks = self._handle_non_fork_edges(
-                                        join_node, join_state.current, join_state.downstream_fork_stack
+                                return
+                            elif isinstance(maybe_overridden_result, JoinItem):
+                                result = maybe_overridden_result
+                                parent_fork_id = self.graph.get_parent_fork(result.join_id).fork_id
+                                for i, x in enumerate(result.fork_stack[::-1]):
+                                    if x.fork_id == parent_fork_id:
+                                        # For non-final joins (those that are intermediate nodes of other joins),
+                                        # preserve the fork stack so downstream joins can still associate with the same fork run
+                                        if self.graph.is_final_join(result.join_id):
+                                            # Final join: remove the parent fork from the stack
+                                            downstream_fork_stack = result.fork_stack[: len(result.fork_stack) - i]
+                                        else:
+                                            # Non-final join: preserve the fork stack
+                                            downstream_fork_stack = result.fork_stack
+                                        fork_run_id = x.node_run_id
+                                        break
+                                else:  # pragma: no cover
+                                    raise RuntimeError('Parent fork run not found')
+
+                                join_node = self.graph.nodes[result.join_id]
+                                assert isinstance(join_node, Join), f'Expected a `Join` but got {join_node}'
+                                join_state = self.active_reducers.get((result.join_id, fork_run_id))
+                                if join_state is None:
+                                    current = join_node.initial_factory()
+                                    join_state = self.active_reducers[(result.join_id, fork_run_id)] = JoinState(
+                                        current, downstream_fork_stack
                                     )
-                                    join_tasks.extend(new_tasks)
-                                if join_tasks:
-                                    for new_task in join_tasks:
-                                        self.active_tasks[new_task.task_id] = new_task
-                                    self._handle_execution_request(join_tasks)
-
-                                if isinstance(maybe_overridden_result, Sequence):
-                                    if isinstance(task_result.result, Sequence):
-                                        new_task_ids = {t.task_id for t in maybe_overridden_result}
-                                        for t in task_result.result:
-                                            if t.task_id not in new_task_ids:
-                                                await self._finish_task(t.task_id, t.node_id)
-                                    self._handle_execution_request(maybe_overridden_result)
-
-                                if task_result.source_is_finished:
-                                    await self._finish_task(task_result.source.task_id, task_result.source.node_id)
-
-                                if not self.active_tasks:
-                                    # if there are no active tasks, we'll be waiting forever for the next result..
-                                    break
-
-                            if self.active_reducers:  # pragma: no branch
-                                # In this case, there are no pending tasks. We can therefore finalize all active reducers
-                                # that don't have intermediate joins which are also active reducers. If a join J2 has an
-                                # intermediate join J1 that shares the same parent fork run, we must finalize J1 first
-                                # because it might produce items that feed into J2.
-                                for (join_id, fork_run_id), join_state in list(self.active_reducers.items()):
-                                    # Check if this join has any intermediate joins that are also active reducers
-                                    should_skip = False
-                                    intermediate_joins = self.graph.intermediate_join_nodes.get(join_id, set())
-
-                                    # Get the parent fork for this join to use for comparison
-                                    join_parent_fork = self.graph.get_parent_fork(join_id)
-
-                                    for intermediate_join_id in intermediate_joins:
-                                        # Check if the intermediate join is also an active reducer with matching fork run
-                                        for (other_join_id, _), other_join_state in self.active_reducers.items():
-                                            if other_join_id == intermediate_join_id:
-                                                # Check if they share the same fork run for this join's parent fork
-                                                # by finding the parent fork's node_run_id in both fork stacks
-                                                join_parent_fork_run_id = None
-                                                other_parent_fork_run_id = None
-
-                                                for fsi in join_state.downstream_fork_stack:  # pragma: no branch
-                                                    if fsi.fork_id == join_parent_fork.fork_id:
-                                                        join_parent_fork_run_id = fsi.node_run_id
-                                                        break
-
-                                                for fsi in other_join_state.downstream_fork_stack:  # pragma: no branch
-                                                    if fsi.fork_id == join_parent_fork.fork_id:
-                                                        other_parent_fork_run_id = fsi.node_run_id
-                                                        break
-
-                                                if (
-                                                    join_parent_fork_run_id
-                                                    and other_parent_fork_run_id
-                                                    and join_parent_fork_run_id == other_parent_fork_run_id
-                                                ):  # pragma: no branch
-                                                    should_skip = True
+                                context = ReducerContext(state=self.state, deps=self.deps, join_state=join_state)
+                                join_state.current = join_node.reduce(context, join_state.current, result.inputs)
+                                if join_state.cancelled_sibling_tasks:
+                                    await self._cancel_sibling_tasks(parent_fork_id, fork_run_id)
+                            else:
+                                for new_task in maybe_overridden_result:
+                                    self.active_tasks[new_task.task_id] = new_task
+
+                            tasks_by_id_values = list(self.active_tasks.values())
+                            join_tasks: list[GraphTask] = []
+
+                            for join_id, fork_run_id in self._get_completed_fork_runs(
+                                task_result.source, tasks_by_id_values
+                            ):
+                                join_state = self.active_reducers.pop((join_id, fork_run_id))
+                                join_node = self.graph.nodes[join_id]
+                                assert isinstance(join_node, Join), f'Expected a `Join` but got {join_node}'
+                                new_tasks = self._handle_non_fork_edges(
+                                    join_node, join_state.current, join_state.downstream_fork_stack
+                                )
+                                join_tasks.extend(new_tasks)
+                            if join_tasks:
+                                for new_task in join_tasks:
+                                    self.active_tasks[new_task.task_id] = new_task
+                                self._handle_execution_request(join_tasks)
+
+                            if isinstance(maybe_overridden_result, Sequence):
+                                if isinstance(task_result.result, Sequence):
+                                    new_task_ids = {t.task_id for t in maybe_overridden_result}
+                                    for t in task_result.result:
+                                        if t.task_id not in new_task_ids:
+                                            await self._finish_task(t.task_id)
+                                self._handle_execution_request(maybe_overridden_result)
+
+                            if task_result.source_is_finished:
+                                await self._finish_task(task_result.source.task_id)
+
+                            if not self.active_tasks:
+                                # if there are no active tasks, we'll be waiting forever for the next result..
+                                break
+
+                        if self.active_reducers:  # pragma: no branch
+                            # In this case, there are no pending tasks. We can therefore finalize all active reducers
+                            # that don't have intermediate joins which are also active reducers. If a join J2 has an
+                            # intermediate join J1 that shares the same parent fork run, we must finalize J1 first
+                            # because it might produce items that feed into J2.
+                            for (join_id, fork_run_id), join_state in list(self.active_reducers.items()):
+                                # Check if this join has any intermediate joins that are also active reducers
+                                should_skip = False
+                                intermediate_joins = self.graph.intermediate_join_nodes.get(join_id, set())
+
+                                # Get the parent fork for this join to use for comparison
+                                join_parent_fork = self.graph.get_parent_fork(join_id)
+
+                                for intermediate_join_id in intermediate_joins:
+                                    # Check if the intermediate join is also an active reducer with matching fork run
+                                    for (other_join_id, _), other_join_state in self.active_reducers.items():
+                                        if other_join_id == intermediate_join_id:
+                                            # Check if they share the same fork run for this join's parent fork
+                                            # by finding the parent fork's node_run_id in both fork stacks
+                                            join_parent_fork_run_id = None
+                                            other_parent_fork_run_id = None
+
+                                            for fsi in join_state.downstream_fork_stack:  # pragma: no branch
+                                                if fsi.fork_id == join_parent_fork.fork_id:
+                                                    join_parent_fork_run_id = fsi.node_run_id
+                                                    break
+
+                                            for fsi in other_join_state.downstream_fork_stack:  # pragma: no branch
+                                                if fsi.fork_id == join_parent_fork.fork_id:
+                                                    other_parent_fork_run_id = fsi.node_run_id
                                                     break
-                                        if should_skip:
-                                            break
 
+                                            if (
+                                                join_parent_fork_run_id
+                                                and other_parent_fork_run_id
+                                                and join_parent_fork_run_id == other_parent_fork_run_id
+                                            ):  # pragma: no branch
+                                                should_skip = True
+                                                break
                                     if should_skip:
-                                        continue
-
-                                    self.active_reducers.pop(
-                                        (join_id, fork_run_id)
-                                    )  # we're handling it now, so we can pop it
-                                    join_node = self.graph.nodes[join_id]
-                                    assert isinstance(join_node, Join), f'Expected a `Join` but got {join_node}'
-                                    new_tasks = self._handle_non_fork_edges(
-                                        join_node, join_state.current, join_state.downstream_fork_stack
-                                    )
-                                    maybe_overridden_result = yield new_tasks
-                                    if isinstance(maybe_overridden_result, EndMarker):  # pragma: no cover
-                                        # This is theoretically reachable but it would be awkward.
-                                        # Probably a better way to get coverage here would be to unify the code pat
-                                        # with the other `if isinstance(maybe_overridden_result, EndMarker):`
-                                        self.task_group.cancel_scope.cancel()
-                                        return
-                                    for new_task in maybe_overridden_result:
-                                        self.active_tasks[new_task.task_id] = new_task
-                                    new_task_ids = {t.task_id for t in maybe_overridden_result}
-                                    for t in new_tasks:
-                                        # Same note as above about how this is theoretically reachable but we should
-                                        # just get coverage by unifying the code paths
-                                        if t.task_id not in new_task_ids:  # pragma: no cover
-                                            await self._finish_task(t.task_id, t.node_id)
-                                    self._handle_execution_request(maybe_overridden_result)
-                except GeneratorExit:
-                    self._task_group.cancel_scope.cancel()
-                    return
+                                        break
+
+                                if should_skip:
+                                    continue
+
+                                self.active_reducers.pop(
+                                    (join_id, fork_run_id)
+                                )  # we're handling it now, so we can pop it
+                                join_node = self.graph.nodes[join_id]
+                                assert isinstance(join_node, Join), f'Expected a `Join` but got {join_node}'
+                                new_tasks = self._handle_non_fork_edges(
+                                    join_node, join_state.current, join_state.downstream_fork_stack
+                                )
+                                maybe_overridden_result = yield new_tasks
+                                if isinstance(maybe_overridden_result, EndMarker):  # pragma: no cover
+                                    # This is theoretically reachable but it would be awkward.
+                                    # Probably a better way to get coverage here would be to unify the code pat
+                                    # with the other `if isinstance(maybe_overridden_result, EndMarker):`
+                                    self.task_group.cancel_scope.cancel()
+                                    return
+                                for new_task in maybe_overridden_result:
+                                    self.active_tasks[new_task.task_id] = new_task
+                                new_task_ids = {t.task_id for t in maybe_overridden_result}
+                                for t in new_tasks:
+                                    # Same note as above about how this is theoretically reachable but we should
+                                    # just get coverage by unifying the code paths
+                                    if t.task_id not in new_task_ids:  # pragma: no cover
+                                        await self._finish_task(t.task_id)
+                                self._handle_execution_request(maybe_overridden_result)
+            except GeneratorExit:
+                self.task_group.cancel_scope.cancel()
+                return
 
         raise RuntimeError(  # pragma: no cover
             'Graph run completed, but no result was produced. This is either a bug in the graph or a bug in the graph runner.'
         )
 
-    async def _finish_task(self, task_id: TaskID, node_id: str) -> None:
+    async def _finish_task(self, task_id: TaskID) -> None:
         # node_id is just included for debugging right now
         scope = self.cancel_scopes.pop(task_id, None)
         if scope is not None:
@@ -898,7 +910,7 @@ class _GraphIterator(Generic[StateT, DepsT, OutputT]):
                 else:
                     pass
         for task_id in task_ids_to_cancel:
-            await self._finish_task(task_id, 'sibling')
+            await self._finish_task(task_id)
 
 
 def _is_any_iterable(x: Any) -> TypeGuard[Iterable[Any]]: